Skip to content

Commit b8fe2a3

Browse files
committed
fix(injection_detection): clear error for inline yara_rules missing a requested name
_load_rules built the compile sources by indexing yara_rules for every requested name, so a name not present inline raised an opaque KeyError. Validate up front and raise a clear ValueError (fail-closed). +test. Signed-off-by: WatchTree-19 <119982314+WatchTree-19@users.noreply.github.com>
1 parent afda55d commit b8fe2a3

2 files changed

Lines changed: 59 additions & 2 deletions

File tree

nemoguardrails/library/injection_detection/actions.py

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -168,10 +168,25 @@ def _load_rules(
168168
log.warning("Injection config was provided but no modules were specified. Returning None.")
169169
return None
170170

171+
if yara_rules:
172+
# When inline `yara_rules` are provided, only those rules are available
173+
# (built-in rules are not loaded in this mode). Every requested rule name
174+
# in `injections` must therefore be present in `yara_rules`; otherwise we
175+
# fail with a clear, actionable error instead of an opaque KeyError.
176+
missing_rules = [name for name in rule_names if name not in yara_rules]
177+
if missing_rules:
178+
msg = (
179+
"Provided set of `injections` in injection config contains rule names "
180+
"not present in the inline `yara_rules`: %r. Available inline rules are %r. "
181+
"Note that built-in rules are not loaded when `yara_rules` is provided."
182+
) % (missing_rules, sorted(yara_rules))
183+
log.error(msg)
184+
raise ValueError(msg)
185+
171186
try:
172187
if yara_rules:
173-
rules_source = {name: rule for name, rule in yara_rules.items() if name in rule_names}
174-
rules = yara.compile(sources={rule_name: rules_source[rule_name] for rule_name in rule_names})
188+
rules_source = {rule_name: yara_rules[rule_name] for rule_name in rule_names}
189+
rules = yara.compile(sources=rules_source)
175190
else:
176191
rules_to_load = {rule_name: str(yara_path.joinpath(f"{rule_name}.yara")) for rule_name in rule_names}
177192
rules = yara.compile(filepaths=rules_to_load)

tests/test_injection_detection.py

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -251,6 +251,48 @@ def test_load_inline_yara_rules():
251251
assert matches[0].namespace == inline_rule_name
252252

253253

254+
def test_inline_yara_rules_missing_injection_name_raises_value_error():
255+
"""Requested `injections` names absent from inline `yara_rules` must raise a
256+
clear ValueError instead of an opaque KeyError.
257+
258+
Regression test: `_load_rules` previously indexed the inline rules dict with
259+
every name in `injections`, so listing a rule name that was not defined in
260+
`yara_rules` (e.g. mixing a built-in rule name with an inline custom rule)
261+
crashed with `KeyError` inside the YARA compile call, and neither
262+
`_validate_injection_config` nor `_extract_injection_config` caught it.
263+
"""
264+
inline_rule_name = "custom_rule"
265+
inline_rule_content = 'rule custom_rule { strings: $a = "foo" condition: $a }'
266+
267+
config = RailsConfig.from_content(
268+
yaml_content=f"""
269+
models: []
270+
rails:
271+
config:
272+
injection_detection:
273+
injections:
274+
- sqli
275+
- {inline_rule_name}
276+
action:
277+
reject
278+
yara_rules:
279+
{inline_rule_name}: |-
280+
{inline_rule_content}
281+
""",
282+
colang_content="",
283+
)
284+
285+
_validate_injection_config(config)
286+
action_option, yara_path, rule_names, yara_rules = _extract_injection_config(config)
287+
288+
# "sqli" is requested via `injections` but is not defined in `yara_rules`.
289+
assert "sqli" in rule_names
290+
assert "sqli" not in yara_rules
291+
292+
with pytest.raises(ValueError, match="sqli"):
293+
_load_rules(yara_path, rule_names, yara_rules)
294+
295+
254296
@pytest.mark.asyncio
255297
async def test_omit_injection_action():
256298
"""Test the omit action for injection detection."""

0 commit comments

Comments
 (0)