fix: reduce Keras ZIP custom-object false positives#716
Merged
Conversation
Avoid flagging safe built-in or allowlisted serialized layer objects as custom objects while preserving warnings for genuinely custom Keras layers, metrics, and losses. Co-Authored-By: Codex <noreply@openai.com>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (1)
WalkthroughAdds allowlist-based safety checks and known-safe layer recognition to the Keras ZIP scanner, updates KNOWN_SAFE_KERAS_LAYER_CLASSES with "NotEqual", updates the changelog, and adds tests to prevent false positives for built-in and allowlisted serialized objects. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
📝 Coding Plan
Comment |
Contributor
Author
Testing Metadata Update
Branch Validation
Broad Hugging Face Sweep (baseline on main)
Focused Rerun For This PR
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@modelaudit/scanners/keras_zip_scanner.py`:
- Around line 128-145: The current early return in
_should_flag_registered_object when
is_known_safe_keras_layer_class(registered_name) is True bypasses module-root
trust checks; change the logic so known-safe registered_name does not
short-circuit detection — instead, continue to evaluate module trust and
serialized-safety: call _layer_uses_allowlisted_module and
_is_known_safe_allowlisted_registered_object (and still keep the
_is_known_safe_serialized_layer short-circuit) so that only layers both
serialized-safe or coming from an allowlisted module/registered-object are
treated as safe; update the branch in _should_flag_registered_object (and the
equality branch comparing registered_name and layer_class) to require module
allowlist checks even for known-safe registered_name before returning False.
In `@tests/scanners/test_keras_zip_scanner.py`:
- Around line 540-623: Add a regression test to ensure trusted-looking names
from non-allowlisted modules still trigger detection: in
tests/scanners/test_keras_zip_scanner.py create a new test (or extend an
existing one) using KerasZipScanner that builds a model config where a layer has
class_name="Add", registered_name="Add", but module="my_custom_pkg.ops"
(non-allowlisted), call scanner.scan on create_configured_keras_zip(tmp_path,
config, file_name="untrusted_add.keras") and assert that result.checks contains
a check with name "Custom Object Detection" (and/or "Custom Layer Class
Detection") to lock in trust-root enforcement.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 3c00ef74-70af-4b9c-87f8-5b2d50e3676d
📒 Files selected for processing (3)
modelaudit/detectors/suspicious_symbols.pymodelaudit/scanners/keras_zip_scanner.pytests/scanners/test_keras_zip_scanner.py
Comment on lines
+128
to
+145
| def _should_flag_registered_object(self, layer: dict[str, Any]) -> bool: | ||
| registered_name = layer.get("registered_name") | ||
| if not isinstance(registered_name, str) or not registered_name.strip(): | ||
| return False | ||
|
|
||
| if is_known_safe_keras_layer_class(registered_name): | ||
| return False | ||
|
|
||
| layer_class = layer.get("class_name") | ||
| if isinstance(layer_class, str) and registered_name.strip() == layer_class.strip(): | ||
| if self._is_known_safe_serialized_layer(layer): | ||
| return False | ||
| return not ( | ||
| self._layer_uses_allowlisted_module(layer) | ||
| and self._is_known_safe_allowlisted_registered_object(layer_class) | ||
| ) | ||
|
|
||
| return True |
Contributor
There was a problem hiding this comment.
Known-safe registered_name currently bypasses trust-root validation.
At Line 133, _should_flag_registered_object exits early for any known-safe registered_name before module-root trust is considered. This allows non-allowlisted module metadata to avoid Custom Object Detection at Line 425.
🔧 Suggested fix
def _should_flag_registered_object(self, layer: dict[str, Any]) -> bool:
registered_name = layer.get("registered_name")
if not isinstance(registered_name, str) or not registered_name.strip():
return False
- if is_known_safe_keras_layer_class(registered_name):
- return False
+ if is_known_safe_keras_layer_class(registered_name):
+ module_values: list[str] = []
+ layer_config = layer.get("config", {})
+ if isinstance(layer_config, dict):
+ for key in ("module", "fn_module"):
+ value = layer_config.get(key)
+ if isinstance(value, str) and value.strip():
+ module_values.append(value.strip())
+ for key in ("module", "fn_module"):
+ value = layer.get(key)
+ if isinstance(value, str) and value.strip():
+ module_values.append(value.strip())
+
+ # Suppress only when module metadata is absent or fully allowlisted.
+ if not module_values or all(self._is_allowlisted_keras_module(v) for v in module_values):
+ return False
+ return TrueBased on learnings: Preserve or strengthen security detections; test both benign and malicious samples when adding scanner/feature changes.
Also applies to: 425-425
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@modelaudit/scanners/keras_zip_scanner.py` around lines 128 - 145, The current
early return in _should_flag_registered_object when
is_known_safe_keras_layer_class(registered_name) is True bypasses module-root
trust checks; change the logic so known-safe registered_name does not
short-circuit detection — instead, continue to evaluate module trust and
serialized-safety: call _layer_uses_allowlisted_module and
_is_known_safe_allowlisted_registered_object (and still keep the
_is_known_safe_serialized_layer short-circuit) so that only layers both
serialized-safe or coming from an allowlisted module/registered-object are
treated as safe; update the branch in _should_flag_registered_object (and the
equality branch comparing registered_name and layer_class) to require module
allowlist checks even for known-safe registered_name before returning False.
Comment on lines
+540
to
+623
| def test_registered_builtin_layer_does_not_false_positive(self, tmp_path: Path) -> None: | ||
| """Built-in layers with registered_name metadata should remain clean.""" | ||
| scanner = KerasZipScanner() | ||
| config = { | ||
| "class_name": "Functional", | ||
| "config": { | ||
| "layers": [ | ||
| { | ||
| "class_name": "InputLayer", | ||
| "name": "input_1", | ||
| "config": {"batch_shape": [None, 4]}, | ||
| }, | ||
| { | ||
| "class_name": "Add", | ||
| "name": "add_1", | ||
| "module": "keras.src.ops.numpy", | ||
| "registered_name": "Add", | ||
| "config": {}, | ||
| }, | ||
| ] | ||
| }, | ||
| } | ||
|
|
||
| result = scanner.scan(str(create_configured_keras_zip(tmp_path, config, file_name="builtin_registered.keras"))) | ||
|
|
||
| assert all(check.name != "Custom Layer Class Detection" for check in result.checks) | ||
| assert all(check.name != "Custom Object Detection" for check in result.checks) | ||
|
|
||
| def test_allowlisted_module_layer_does_not_false_positive(self, tmp_path: Path) -> None: | ||
| """Layers from allowlisted Keras modules should not be treated as custom objects.""" | ||
| scanner = KerasZipScanner() | ||
| config = { | ||
| "class_name": "Functional", | ||
| "config": { | ||
| "layers": [ | ||
| { | ||
| "class_name": "InputLayer", | ||
| "name": "input_1", | ||
| "config": {"batch_shape": [None, 4]}, | ||
| }, | ||
| { | ||
| "class_name": "NotEqual", | ||
| "name": "not_equal", | ||
| "module": "keras.src.ops.numpy", | ||
| "registered_name": "NotEqual", | ||
| "config": {}, | ||
| }, | ||
| ] | ||
| }, | ||
| } | ||
|
|
||
| result = scanner.scan(str(create_configured_keras_zip(tmp_path, config, file_name="allowlisted_module.keras"))) | ||
|
|
||
| assert all(check.name != "Custom Layer Class Detection" for check in result.checks) | ||
| assert all(check.name != "Custom Object Detection" for check in result.checks) | ||
|
|
||
| def test_allowlisted_module_does_not_suppress_unknown_custom_layer(self, tmp_path: Path) -> None: | ||
| """Unknown classes must still be flagged even if module metadata looks Keras-owned.""" | ||
| scanner = KerasZipScanner() | ||
| config = { | ||
| "class_name": "Functional", | ||
| "config": { | ||
| "layers": [ | ||
| { | ||
| "class_name": "InputLayer", | ||
| "name": "input_1", | ||
| "config": {"batch_shape": [None, 4]}, | ||
| }, | ||
| { | ||
| "class_name": "EvilLayer", | ||
| "name": "evil_layer", | ||
| "module": "keras.src.ops.numpy", | ||
| "registered_name": "EvilLayer", | ||
| "config": {}, | ||
| }, | ||
| ] | ||
| }, | ||
| } | ||
|
|
||
| result = scanner.scan(str(create_configured_keras_zip(tmp_path, config, file_name="evil_allowlisted.keras"))) | ||
|
|
||
| assert any(check.name == "Custom Layer Class Detection" for check in result.checks) | ||
| assert any(check.name == "Custom Object Detection" for check in result.checks) | ||
|
|
Contributor
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Add one regression for safe class names coming from non-allowlisted modules.
These tests are good, but please add a case like class_name="Add", registered_name="Add", module="my_custom_pkg.ops" and assert Custom Object Detection is raised. That locks in trust-root enforcement.
Based on learnings: Preserve or strengthen security detections; test both benign and malicious samples when adding scanner/feature changes.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/scanners/test_keras_zip_scanner.py` around lines 540 - 623, Add a
regression test to ensure trusted-looking names from non-allowlisted modules
still trigger detection: in tests/scanners/test_keras_zip_scanner.py create a
new test (or extend an existing one) using KerasZipScanner that builds a model
config where a layer has class_name="Add", registered_name="Add", but
module="my_custom_pkg.ops" (non-allowlisted), call scanner.scan on
create_configured_keras_zip(tmp_path, config, file_name="untrusted_add.keras")
and assert that result.checks contains a check with name "Custom Object
Detection" (and/or "Custom Layer Class Detection") to lock in trust-root
enforcement.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Validation
/Users/yashchhabria/projects/modelauditing/modelaudit/.venv/bin/ruff format modelaudit/ tests//Users/yashchhabria/projects/modelauditing/modelaudit/.venv/bin/ruff check --fix modelaudit/ tests//Users/yashchhabria/projects/modelauditing/modelaudit/.venv/bin/mypy modelaudit//Users/yashchhabria/projects/modelauditing/modelaudit/.venv/bin/pytest -n auto -m "not slow and not integration" --maxfail=1AddandNotEqualserialized objects while preserving remaining custom-layer and custom-metric/loss findingsSummary by CodeRabbit
Bug Fixes
Tests