AhoCorasick with regex literals (#3073) - #3145
Conversation
williballenthin
left a comment
williballenthin
left a comment
There was a problem hiding this comment.
this is cool!
do you have any benchmark results that you can share?
| # Change Log | ||
|
|
||
| ## master (unreleased) | ||
| - Aho-Corasick automaton generation and scan of regex literals to speed up rule matching corkami@google.com |
There was a problem hiding this comment.
reference original issue, too
| viv-utils==0.8.0 | ||
| vivisect==1.3.2 | ||
| msgspec==0.21.1 | ||
| pyahocorasick>=2.0 |
There was a problem hiding this comment.
requirements.txt should contain pinned versions, not ranges
There was a problem hiding this comment.
(we need to add this advice to AGENTS.md, too)
| try: | ||
| import re._parser as re_parser | ||
| except ImportError: | ||
| import sre_parse as re_parser # type: ignore |
There was a problem hiding this comment.
move imports to top of file, unless there's an import cycle or good reason for it (is there?).
also, what is sre_parse?
There was a problem hiding this comment.
ah, i guess maybe you use this as a sort of feature flag to signal when to use A-C matching. please use another mechanism for this. see later comment.
| if op == re_parser.LITERAL: | ||
| cur_run.append(chr(av)) | ||
| elif op == re_parser.SUBPATTERN: | ||
| sub_ast = av[-1] if isinstance(av, (tuple, list)) else av | ||
| sub_runs = walk_ast(sub_ast) | ||
| flush() | ||
| runs.update(sub_runs) | ||
| elif op in (re_parser.MAX_REPEAT, re_parser.MIN_REPEAT): | ||
| min_rep, _max_rep, sub_ast = av[0], av[1], av[2] | ||
| if min_rep >= 1 and len(sub_ast) == 1 and sub_ast[0][0] == re_parser.LITERAL: | ||
| cur_run.append(chr(sub_ast[0][1])) | ||
| else: | ||
| flush() | ||
| if min_rep >= 1: | ||
| sub_runs = walk_ast(sub_ast) | ||
| runs.update(sub_runs) | ||
| elif op == re_parser.BRANCH: |
There was a problem hiding this comment.
it would be helpful, but not required, to see examples of the state of the regex or something in code comments here, so readers can more easily understand the code.
like:
# regex: (aaa|bbb)
# ^ ^ pointing here, next step will be XYZor whatever makes sense to understand the parsing/walking
| def flush(): | ||
| nonlocal cur_run | ||
| if cur_run: | ||
| s = "".join(cur_run) | ||
| if len(s) >= _STRING_LITERAL_MIN: | ||
| runs.add(s) | ||
| cur_run = [] |
There was a problem hiding this comment.
note to self: walk_ast is probably not performance sensitive (O(#of rules), run upon initial rule parsing), so its probably ok to create function objects and closures here.
| else: | ||
| lit_set = req_lits | ||
| for lit in lit_set: | ||
| if len(lit) >= _STRING_LITERAL_MIN: |
There was a problem hiding this comment.
this is already part of the API contract for the _compute_best_required_literals so i think we can drop this check.
| if isinstance(wanted_feature, capa.features.common.Substring): | ||
| candidate_rule_names.add(rule_name) |
There was a problem hiding this comment.
do we have to check that the casing matches here? i think the substring literal as added as lowercase, though the found string may not have the same case?
There was a problem hiding this comment.
str(string_feature.value)[_end_index-len(wanted_feature.value):_end_index] == wanted_feature.valueor something
There was a problem hiding this comment.
worst case, we do a little too much work as a rule is evaluated that won't match. so the tradeoff is whether or not the case validation is easy/fast or not. already this filter is probably very effective.
There was a problem hiding this comment.
substring: WwW ends up FP-ing on string: www.google.com
| if wanted_feature.re.search(str(string_feature.value)): | ||
| candidate_rule_names.add(rule_name) |
There was a problem hiding this comment.
likewise, does it make sense to do any validation of the literal match casing before doing the full regex?
There was a problem hiding this comment.
but do we have access to the wanted literal here? maybe not, so maybe this isn't possible.
There was a problem hiding this comment.
regex: /WwW/ FPs on string: www.google.com
| assert capa.rules._required_literal("test[0-9]+") == "test" | ||
| assert capa.rules._required_literal("a[0-9]b") is None | ||
| assert capa.rules._required_literal("(test|b)") is None | ||
| assert capa.rules._required_literal("(test|example)") == {"test", "example"} |
There was a problem hiding this comment.
what about (test|(foo|bar)) i assume this is {test, foo, bar}. would like to confirm how nesting works.
There was a problem hiding this comment.
foo and bar are shorter than 4, so the whole outcome is None.
Nesting works as you expected with all literals being longer.
assert capa.rules._required_literal("(test|(foo|bar))") is None
assert capa.rules._required_literal("(test|(example|another))") == {"test", "example", "another"}| assert "test regex confirmation" not in matches3 | ||
|
|
||
|
|
||
| def test_differential_parity_prefilter_on_vs_off(monkeypatch): |
There was a problem hiding this comment.
i'd prefer to reset string_literal_index after the index construction, rather than use the presence of imports to signal feature flags. we don't use this latter pattern today, so it would be new and maybe confusing/inconsistent.
|
|
whoa that's outstanding. i didn't realize the impact would be so large. great work! |
|
i think we should add a pass to the rule linter to flag regexes for which we can't extract a literal. |
|
Wow, great performance improvements! |
3 participants
An Aho-Corasick automaton is generated from regex literals bigger than 3.
No performance hit should happen on smaller files.
This patch has been kept to a minimum of maintenance, while speeds up of 16-32x can be expected on files with a lot of detected strings.
Checklist