rules: index extracted bytes by fixed-length prefix buckets for O(1) candidate selection - #2928
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a significant performance enhancement to Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
There was a problem hiding this comment.
Code Review
This pull request introduces a significant performance optimization for matching bytes features by replacing a linear scan with a prefix-based bucketing strategy. The implementation is sound and includes a fallback for short patterns and new tests to ensure correctness. I have one suggestion to make the bucket lookup slightly more efficient and idiomatic.
williballenthin
left a comment
There was a problem hiding this comment.
addressing this TODO is much more about doing the work to profile and benchmark and find the right balance, versus writing the naive code. @devs6186 are you up for it?
otherwise, i don't want to merge this because its more code with unclear real-world performance impact.
…ction Instead of iterating all extracted Bytes features for every bytes-based rule, build a prefix index keyed by fixed bucket sizes (4, 8, 16, 32, 64, 128, 256) once per scope evaluation. Each bytes pattern is looked up in the largest bucket that fits its length, then only candidates sharing that prefix are compared, replacing the previous O(n) linear scan with an O(1) hash lookup. Patterns shorter than the minimum bucket still fall back to the full scan. Adds a test to verify correctness for exact match, startswith match, mismatch, and short-bytes cases. Closes: mandiant#2128
68bb178 to
b14dcf6
Compare
|
Thanks for the feedback @williballenthin , I reworked this to prioritize simplicity and measurable impact.
I also expanded tests in tests/test_match.py for:
Benchmark results (scripts/profile-time.py, 30 repeats, same sample/rules):
Memory profile (scripts/profile-memory.py) was effectively unchanged between multi and single variants in this run. |
…) call The previous implementation rebuilt a `defaultdict` mapping byte prefixes to extracted feature values inside `_match()`, which is called per function/basic-block/instruction. Moving the rule-side index build to `_index_rules_by_feature()` (called once at RuleSet construction) eliminates this per-call allocation and O(R) rule iteration from the hot path. `_match()` now looks up candidate rules via the pre-built `bytes_prefix_index` stored in `_RuleFeatureIndex`, iterating only extracted byte features to compute their prefixes.
|
did some more digging The root cause was that the prefix index was being rebuilt on every Latest commit moves the entire index build into The struct: @dataclass
class _RuleFeatureIndex:
...
bytes_prefix_index: dict[int, list[tuple[str, bytes]]]is populated once and shared across all Happy to benchmark again if that would help — just wanted to get the structural fix in first. |
can you run capa against all the files in capa-testfiles before/after and share the results? or at least mimikatz.exe. the numbers you show above are too small to really feel confident about. i'm afraid there might be too much noise in the measurements. |
When _RuleFeatureIndex gains a new field, pickle.loads() on an older cached ruleset succeeds but the resulting objects silently lack the new field — causing an AttributeError deep in _match() at runtime. Extend load_cached_ruleset() to walk every _RuleFeatureIndex in the loaded ruleset and verify each dataclass field is present on the instance. On mismatch, delete the stale cache and return None so the caller rebuilds from scratch. Production users are unaffected (the version hash in the cache key already invalidates caches across releases); this guard covers developer switching between branches.
- Change _match() guard from bytes_rules to bytes_prefix_index so the guard references the field actually used for candidate selection. - Update stale comment to describe the prefix-bucket strategy. - Clarify bytes_rules dataclass comment (retained for logging only). - Add test_bytes_prefix_index_mixed_short_and_long_patterns covering rules with both short (<4B) and long (>=4B) patterns exercised together.
|
@williballenthin
The evaluate.feature.bytes counter drops dramatically across all three — from 71k to 4.4k on kernel32 which confirms the prefix bucketing is cutting out the vast majority of pointless comparisons. Wall time improvement is ~29% on kernel32, ~34% on microsocks, ~31% on PMA 16-01. I don't have mimikatz.exe_ in my local checkout (it's not in the default testfiles submodule). If you want mimikatz numbers specifically, could you point me to where I can grab it? Happy to rerun. Also pushed a cleanup commit:
|
|
wow, that's a huge improvement! something like 30% faster (wall clock) with this PR? thats well beyond my expectations. do you mind also confirming the i will review this PR carefully and make a plan for merging it. thank you @devs6186! |
williballenthin
left a comment
There was a problem hiding this comment.
the implementation looks good to me, some inline comments to address.
| # Mapping from rule name to list of Bytes features that have to match. | ||
| # All these features will be evaluated whenever a Bytes feature is encountered. | ||
| # Retained for logging and the emptiness check; candidate selection uses bytes_prefix_index. |
There was a problem hiding this comment.
i don't understand this comment. can we remove this structure completely?
lets phrase our comments to explain how the code currently is, not referring to changes/history. "retained" requires the reader to know the history of the file, rather than how it works today.
| # For long patterns, group extracted bytes by their 4-byte prefix and look up | ||
| # only the rules whose pattern prefix matches. | ||
| for feature in bytes_features: | ||
| assert isinstance(feature.value, bytes) |
There was a problem hiding this comment.
can we get rid of this check? mypy should know about the type of value
| # Short-pattern rules (key -1) require a linear scan against all extracted bytes. | ||
| for rule_name, pattern in feature_index.bytes_prefix_index.get(-1, ()): | ||
| for feature in bytes_features: | ||
| assert isinstance(feature.value, bytes) | ||
| if feature.value.startswith(pattern): | ||
| candidate_rule_names.add(rule_name) | ||
| break |
There was a problem hiding this comment.
maybe we can put this behind a check of if -1 in feature_index.bytes_prefix_index to avoid the creation of temporary objects. i don't think there are any rules today that have such a short bytes feature, nor do i expect that we'll add any, so almost certainly this branch won't be taken.
| else: | ||
| return cache.ruleset | ||
|
|
||
| # Verify that every _RuleFeatureIndex in the loaded ruleset has all the fields | ||
| # that the current class definition declares. A cache built by an older dev build | ||
| # (before a new field was added to _RuleFeatureIndex) will be missing those fields | ||
| # on the unpickled objects even though pickle.loads() succeeds. | ||
| try: | ||
| for fi in cache.ruleset._feature_indexes_by_scopes.values(): | ||
| for field in dataclasses.fields(fi): | ||
| getattr(fi, field.name) | ||
| except AttributeError: | ||
| logger.debug("rule set cache has incompatible structure (stale dev build): %s", path) | ||
| path.unlink() | ||
| return None | ||
|
|
||
| return cache.ruleset |
There was a problem hiding this comment.
can you open a separate issue and PR about this? i'm not sure this is the way we want to detect this, and don't want to merge it as part of this PR.
| rule = textwrap.dedent(""" | ||
| rule = textwrap.dedent( | ||
| """ |
There was a problem hiding this comment.
there's a lot of extra whitespace changes here. would you please revert them?
|
|
||
|
|
There was a problem hiding this comment.
these are good tests, thanks for including them
|
Hi @williballenthin ill address the issues shortly |
- remove bytes_rules from _RuleFeatureIndex; bytes_prefix_index is the only structure needed for candidate selection - build bytes_prefix_index directly in _index_rules_by_feature() instead of building bytes_rules then converting, removing one full pass - add if -1 in bytes_prefix_index guard to avoid temporary object creation for the short-pattern fallback (almost never taken) - remove assert isinstance(feature.value, bytes) checks in _match(); add Bytes.value: bytes class-level annotation so mypy narrows the type without the runtime check - remove cache structure compatibility block from cache.py per reviewer request to handle in a separate PR - update test assertions from bytes_rules to bytes_prefix_index
|
Hi @williballenthin I stripped those fields to compare semantic output, not run-specific noise.
The Json input which came out identical- Every rule that fired before fires after on the same set of addresses. |
Hey @williballenthin i have addressed all inline comments in the latest commit. Thank you ! |
|
thanks @devs6186 great work! |
Summary
Closes #2128.
Bytes feature matching currently scans all extracted
Bytesfeatures linearly for each rule pattern, callingfeature.value.startswith(pattern)on every one. A binary's instruction operands can produce hundreds of bytes features; with the default rule set, this means millions ofstartswithcomparisons on a moderately complex sample (see the evaluation counts cited in #2128).A bytes pattern of length L can only match extracted bytes whose length is ≥ L and whose first L bytes equal the pattern. By pre-indexing extracted bytes at a small set of fixed prefix lengths (4, 8, 16, 32, 64, 128, 256 bytes), we can reduce the comparison set from all extracted bytes to only those sharing the same fixed-length prefix — typically a handful of candidates.
What changed
capa/rules/__init__.pyModule-level constant
_BYTES_BUCKET_SIZES = (4, 8, 16, 32, 64, 128, 256)defines the prefix bucket sizes.In
RuleSet._match():bytes_featuresfrom the full feature set, a prefix indexbytes_prefix_index: dict[int, dict[bytes, list[...]]]is built in a single pass. Each extracted bytes feature of length ≥ N is registered in the bucket for N under its first-N-byte prefix.startswith, bringing the inner loop down from O(all_bytes) to O(candidates_with_same_prefix).evaluate()linear scan is used as a fallback so correctness is preserved for tiny patterns.Correctness
The prefix index is a refinement of the original linear scan — it can only reduce the candidate set, never miss a match, because two bytes sequences must share the same fixed-length prefix before
startswithis applied. Semantic behaviour is identical to the previous implementation.Performance note
The gain is greatest for patterns whose common prefix lengths (16, 32 bytes — GUIDs, S-boxes, public-key material) are well-separated from the bulk of extracted bytes features. The bucket index is rebuilt per scope call, so its overhead is bounded to O(n·7) construction plus O(1) lookups per pattern.
Tests
New test
test_bytes_prefix_index_correctnessintests/test_match.pycovers:startswithmatch (extracted bytes longer than pattern)Test plan
pytest tests/test_match.py -v— all 22 tests passpytest tests/test_rules.py -v— all existing tests passpre-commit run --all-files— isort, black, ruff all pass