fix: tighten dill MemoryError downgrade gating#710
Merged
Conversation
Contributor
|
Caution Review failedPull request was closed or merged during review WalkthroughReworks pickle scanner global-analysis: adds dill global detection, replaces marker-string heuristics with opcode/globals-based extension signals, introduces extension-based gating (including .dill and .joblib), and adjusts passes_global_gate and depth logic. Tests expanded for MemoryError, downgrade semantics, and joblib/dill sample generation. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 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 docstrings
🧪 Generate unit tests (beta)
📝 Coding Plan
Comment |
![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: 3
🤖 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/pickle_scanner.py`:
- Line 5666: The check that computes has_dill_globals only matches "dill" and
"dill.*", which misses legitimate dill internals like "_dill" and "dill._dill"
and can cause false positives in the later global gate (the code around
has_dill_globals and the global gate at the block referenced by lines
5711-5717). Update the detection logic that computes has_dill_globals in
pickle_scanner.py to also accept module names that equal "_dill" or start with
"_dill" and likewise handle "dill._dill" (e.g., treat any mod that is "dill",
startswith("dill."), is "_dill", or startswith("_dill.") as dill-related); then
ensure the related global gate checks (the 5711-5717 logic) use this expanded
predicate so benign dill internals are not rejected while preserving
malicious-dill detection.
- Around line 2901-2905: The current check in the pickle scanner returns True
purely based on ext_lower == ".dill" which trusts the filename extension; change
this to require evidence of Dill content before treating as legitimate: in the
function that performs the shallow pickle header check (the block that checks
ext_lower and currently returns True for ".dill"), replace the extension-only
pass with an evidence-based check — e.g., read additional bytes from the stream
and look for Dill-specific markers (module/import references like b"dill" or
Dill-specific persistent-type signatures) or, if available, attempt a
guarded/dry-run dill sniff (import dill and call a non-destructive detection
helper) and only return True when those markers/detection succeed; keep
ext_lower as a secondary hint, not a sole condition.
In `@tests/scanners/test_pickle_scanner.py`:
- Around line 2174-2178: The monkeypatched lambda for
PickleScanner._extract_globals_advanced has the wrong signature and will raise
TypeError when callers pass keyword args like multiple_pickles or
scan_start_time; update the patch to accept the same parameters used in real
calls (e.g., self, file_obj, multiple_pickles=True, scan_start_time=None) or
simply use **kwargs to absorb extras so the test stub matches the real method
signature in PickleScanner.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: ba029f34-b929-462b-9ba5-7bbaafcf36bb
📒 Files selected for processing (2)
modelaudit/scanners/pickle_scanner.pytests/scanners/test_pickle_scanner.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (3)
modelaudit/scanners/pickle_scanner.py (2)
5666-5666:⚠️ Potential issue | 🟠 MajorBroaden dill-global detection to include
_dillinternals.Line 5666 misses
_dill/_dill.*, so legitimate dill payloads can fail the dill-global gate.Suggested fix
- has_dill_globals = any(mod == "dill" or mod.startswith("dill.") for mod, _func, _opcode in advanced_globals) + has_dill_globals = any( + mod in {"dill", "_dill"} + or mod.startswith(("dill.", "_dill.", "dill._dill")) + for mod, _func, _opcode in advanced_globals + )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 `@modelaudit/scanners/pickle_scanner.py` at line 5666, The current dill detection only checks for "dill" and "dill.*" in the generator that sets has_dill_globals; update that check to also include "_dill" and "_dill.*" so internal dill modules are detected. Modify the comprehension that computes has_dill_globals (which iterates over advanced_globals) to treat mod equal to "_dill" or mod.startswith("_dill.") as a match in addition to the existing "dill" checks, and add unit tests covering benign and malicious payloads that use both "dill" and "_dill" module names.
2901-2905:⚠️ Potential issue | 🟠 MajorRequire evidence-based
.dilllegitimacy checks, not extension-only acceptance.Line 2904-2905 still returns
Truefor any.dillfile after a shallow header check, which keeps downgrade gating extension-trusting.Suggested fix
- # Dill can serialize plain pickle-compatible objects without - # embedding obvious dill markers near the front of the stream, so - # the extension remains the legitimacy signal here. if ext_lower == ".dill": - return True + f.seek(0) + sample = f.read(2048) + dill_indicators = [b"dill", b"_dill", b"dill._dill"] + return any(marker in sample for marker in dill_indicators)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 `@modelaudit/scanners/pickle_scanner.py` around lines 2901 - 2905, The current acceptance of ".dill" files is extension-only; change the logic around the ext_lower == ".dill" branch in the function handling pickle legitimacy (the block referencing ext_lower and returning True) to perform evidence-based checks instead of an unconditional return: read the file header and/or a small prefix of bytes and verify dill-specific markers (or attempt a controlled, sandboxed dill.loads that is tightly exception-caught), reject if deserialization yields code objects or REDUCE/BUILD opcodes, and only return True when those deeper checks pass; also add unit tests covering benign and malicious .dill samples to ensure the new checks prevent downgrade gating.tests/scanners/test_pickle_scanner.py (1)
2174-2178:⚠️ Potential issue | 🟠 MajorFix monkeypatch signature mismatch for
_extract_globals_advanced.Line 2177 patches
_extract_globals_advancedaslambda self, file_obj: set(), but the scanner calls this method with keyword args (e.g.,multiple_pickles,scan_start_time). This can raiseTypeErrorand invalidate this regression scenario.🐛 Proposed fix
monkeypatch.setattr( PickleScanner, "_extract_globals_advanced", - lambda self, file_obj: set(), + lambda self, file_obj, multiple_pickles=True, scan_start_time=None: set(), )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/scanners/test_pickle_scanner.py` around lines 2174 - 2178, The monkeypatch for PickleScanner._extract_globals_advanced uses a lambda with signature (self, file_obj) but the real scanner calls it with keyword arguments like multiple_pickles and scan_start_time; update the monkeypatch to accept the extra parameters (e.g., change the replacement to a callable matching PickleScanner._extract_globals_advanced by adding *args/**kwargs or explicit keyword params) so the patched function signature matches calls and returns set() as before.
🤖 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/pickle_scanner.py`:
- Around line 5712-5718: The current logic sets
has_extension_based_serialization_globals true for file_ext == ".joblib" without
requiring parsed joblib/sklearn/numpy globals; change the predicate so .joblib
only qualifies when advanced_globals are present and has_joblib_globals is true
(e.g. replace the .joblib branch with file_ext == ".joblib" and
bool(advanced_globals) and has_joblib_globals), keep the .dill branch gated by
has_dill_globals, and then leave has_legitimate_serialization_globals to combine
that result with the existing not has_dangerous_advanced_global check.
---
Duplicate comments:
In `@modelaudit/scanners/pickle_scanner.py`:
- Line 5666: The current dill detection only checks for "dill" and "dill.*" in
the generator that sets has_dill_globals; update that check to also include
"_dill" and "_dill.*" so internal dill modules are detected. Modify the
comprehension that computes has_dill_globals (which iterates over
advanced_globals) to treat mod equal to "_dill" or mod.startswith("_dill.") as a
match in addition to the existing "dill" checks, and add unit tests covering
benign and malicious payloads that use both "dill" and "_dill" module names.
- Around line 2901-2905: The current acceptance of ".dill" files is
extension-only; change the logic around the ext_lower == ".dill" branch in the
function handling pickle legitimacy (the block referencing ext_lower and
returning True) to perform evidence-based checks instead of an unconditional
return: read the file header and/or a small prefix of bytes and verify
dill-specific markers (or attempt a controlled, sandboxed dill.loads that is
tightly exception-caught), reject if deserialization yields code objects or
REDUCE/BUILD opcodes, and only return True when those deeper checks pass; also
add unit tests covering benign and malicious .dill samples to ensure the new
checks prevent downgrade gating.
In `@tests/scanners/test_pickle_scanner.py`:
- Around line 2174-2178: The monkeypatch for
PickleScanner._extract_globals_advanced uses a lambda with signature (self,
file_obj) but the real scanner calls it with keyword arguments like
multiple_pickles and scan_start_time; update the monkeypatch to accept the extra
parameters (e.g., change the replacement to a callable matching
PickleScanner._extract_globals_advanced by adding *args/**kwargs or explicit
keyword params) so the patched function signature matches calls and returns
set() as before.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: f51653da-8d4c-404c-85d8-377580683782
📒 Files selected for processing (2)
modelaudit/scanners/pickle_scanner.pytests/scanners/test_pickle_scanner.py
mldangelo
deleted the
codex/investigate-joblib/dill-memoryerror-security-issue
branch
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
1 participant
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.
Motivation
.dillfiles (and some serialization cases) as legitimate based primarily on extension or shallow markers, which can be abused by a crafted pickle to reduce severity to INFO and bypass warning-level gates.MemoryErroras a benign scanner limitation.Description
.dillfiles in_is_legitimate_serialization_file()instead of accepting the extension alone by checking forb"dill",b"_dill", orb"dill._dill"in the first 2KB of the file. (Editedmodelaudit/scanners/pickle_scanner.py.)dillglobals (has_dill_globals) and tighten the serialization global-gate so extension-validated files require positive advanced-global evidence (.joblib-> joblib/sklearn/numpy globals,.dill-> dill globals) before a MemoryError downgrade is allowed. (Editedmodelaudit/scanners/pickle_scanner.py.).dillfile that triggersMemoryErroris not downgraded to INFO when no dill globals are parsed. (Editedtests/scanners/test_pickle_scanner.py.)Testing
uv run ruff format modelaudit/ tests/successfully.uv run ruff check --fix modelaudit/ tests/successfully.uv run mypy modelaudit/with no type issues.uv run pytest -n auto -m "not slow and not integration" --maxfail=1which reported an existing unrelated test failure (tests/utils/helpers/test_secure_hasher.py::TestErrorHandling::test_hash_permission_denied); this failure is pre-existing and unrelated to the changes in this PR.uv run pytest tests/scanners/test_pickle_scanner.py -k "dill_memory_error or dangerous_globals_not_downgraded or legitimate_pytorch_bin_memory_error"and the selected pickle-scanner regression tests passed.Codex Task
Summary by CodeRabbit
New Features
Bug Fixes
Tests