fix: harden import-only pickle global detection#691
Conversation
WalkthroughAdds import-only GLOBAL/STACK_GLOBAL detection to the pickle scanner: simulates symbolic reference maps, tracks import origins, classifies import targets (safe/unknown/dangerous), and surfaces import-only metadata in findings to reduce false positives. Changes
Sequence Diagram(s)sequenceDiagram
participant Input as Pickle Bytecode
participant Scanner as PickleScanner
participant Simulator as SymbolicRefSimulator
participant Classifier as ImportClassifier
participant Reporter as FindingReporter
Input->>Scanner: supply pickle stream
Scanner->>Simulator: parse opcodes & simulate refs
Simulator-->>Scanner: stack_global_refs, callable_refs, callable_origin_refs
Scanner->>Classifier: resolve & classify import targets
Classifier-->>Scanner: classification (safe/unknown/dangerous) + confidence
Scanner->>Reporter: produce findings enriched with import_only, classification, ml_context_confidence
Reporter-->>Scanner: suppress or emit final findings
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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/pickle_scanner.py`:
- Around line 1804-1808: The current _is_safe_import_only_global() lets
import-only allowlist entries like torch.load pass before checking
actual-dangery status; change its logic to first call
_is_actually_dangerous_global(mod, func) and if that returns True immediately
treat it as unsafe (return False), otherwise proceed to the existing allowlist
check against IMPORT_ONLY_SAFE_GLOBALS or _is_safe_ml_global; update references
in the same module to use this revised behavior. Also add regression tests
covering import-only payload "GLOBAL torch\nload\n" and stack/global payload
"STACK_GLOBAL torch.load" (both benign and malicious variants) to ensure
dangerous detections are preserved or strengthened. Ensure tests reference the
scanner functions (_is_safe_import_only_global and
_is_actually_dangerous_global) so failures point to the updated logic.
In `@tests/scanners/test_pickle_scanner.py`:
- Around line 1214-1252: Update the three tests
(test_import_only_global_malicious_is_flagged,
test_import_only_global_mixed_case_module_is_flagged,
test_import_only_stack_global_is_flagged) to assert that the positive detections
include import-only metadata: after you compute failed_checks and pick the
matching check(s) (variable c), add assertions that c.details["import_only"] is
True and that c.details.get("import_reference") is present (or equals the
expected value like "evilpkg.thing"/"EvilPkg.thing" for the GLOBAL cases);
ensure you check these fields for both the "Global Module Reference Check" and
the "STACK_GLOBAL Module Check" positive cases so the tests fail if a detection
is no longer classified as import-only.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: c4f1c4e2-cd2b-4b34-805b-e9d0e51f696c
📒 Files selected for processing (3)
CHANGELOG.mdmodelaudit/scanners/pickle_scanner.pytests/scanners/test_pickle_scanner.py
mldangelo
force-pushed
the
feat/pickle-import-only-findings
branch
from
4742262 to
fd05e62
Compare
|
Rebased onto current Changes in this update:
Validation passed:
|
|
Rebased onto current , addressed the two CodeRabbit findings, and reran validation locally.\n\nChanges in this update:\n- keep dangerous import-only refs like \ out of the import-only safe path\n- add \ to failed \ import-only findings\n- harden the positive import-only tests to assert \ and exact \ values\n- add explicit regressions for import-only \ and \n\nValidation passed:\n- \331 files left unchanged\n- \All checks passed!\n- \All checks passed!\n- \331 files already formatted\n- \Success: no issues found in 191 source files\n- ============================= test session starts ============================== ........................................................................ [ 3%] |
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
modelaudit/scanners/pickle_scanner.py (1)
2048-2053:⚠️ Potential issue | 🟠 MajorAdd
_pop_to_mark()call to INST handler to consume the marked frame.The INST opcode consumes a marked frame of constructor arguments (per pickletools spec:
stack_before=[mark, stackslice], stack_after=[any]), but the current handler leaves it on the stack. This breaks subsequent stack simulation for all later opcodes, causing_pop_to_mark()calls in TUPLE, LIST, OBJ, and other handlers to consume stale frames instead, corruptingcallable_origin_refstracking.Align INST with the pattern used for TUPLE, LIST, and OBJ (which all call
_pop_to_mark()for opcodes with the same stack contract).if name == "INST" and isinstance(arg, str): parsed = _parse_module_function(arg) if parsed: callable_refs[i] = parsed callable_origin_refs[i] = i + _pop_to_mark() stack.append(unknown) continue🤖 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 2048 - 2053, The INST opcode handler currently pushes unknown without consuming the marked frame; modify the INST handler in pickle_scanner.py to call _pop_to_mark() (like the TUPLE, LIST, and OBJ handlers) before appending unknown so the marked frame is removed from stack and callable_origin_refs/callable_refs tracking remains correct; ensure the call occurs before stack.append(unknown) and that any returned slice/value from _pop_to_mark() is discarded if unused.
🤖 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 1842-1863: The classifier _classify_import_reference currently
only matches denylisted refs with exact case; update it to perform a
case-normalized denylist check (e.g., normalize mod and func to a canonical
case) before falling back to "unknown_third_party" so mixed-case variants of
denylisted modules/functions are classified as dangerous; do this by checking
the normalized pair against the existing denylist/blocked set (or via a new
helper used by _is_actually_dangerous_global) while preserving the existing
calls to _is_resolved_import_target, _is_safe_import_only_global, and
_is_plausible_import_only_module and reusing WARNING_SEVERITY_MODULES to
determine severity.
- Around line 1805-1813: The helper _is_safe_import_only_global currently treats
entries from _is_safe_ml_global (ML_SAFE_GLOBALS) as import-only-safe, which
lets deserializers like dill.load, dill.loads, and joblib._pickle_load bypass
checks; remove that widening by not calling _is_safe_ml_global here and only
consulting IMPORT_ONLY_SAFE_GLOBALS (or explicit constructor/data symbols). Fix
options: (A) change _is_safe_import_only_global to return False if
_is_actually_dangerous_global(...) else return func in
IMPORT_ONLY_SAFE_GLOBALS.get(mod, frozenset()) (remove the _is_safe_ml_global
branch), and move any true constructor/data symbols from ML_SAFE_GLOBALS into
IMPORT_ONLY_SAFE_GLOBALS; or (B) add deserializer APIs (e.g., dill.load,
dill.loads, joblib._pickle_load) to ALWAYS_DANGEROUS_FUNCTIONS so they are never
treated as import-only-safe. Ensure you update ML_SAFE_GLOBALS and tests
accordingly.
---
Outside diff comments:
In `@modelaudit/scanners/pickle_scanner.py`:
- Around line 2048-2053: The INST opcode handler currently pushes unknown
without consuming the marked frame; modify the INST handler in pickle_scanner.py
to call _pop_to_mark() (like the TUPLE, LIST, and OBJ handlers) before appending
unknown so the marked frame is removed from stack and
callable_origin_refs/callable_refs tracking remains correct; ensure the call
occurs before stack.append(unknown) and that any returned slice/value from
_pop_to_mark() is discarded if unused.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 36309ccd-9218-4324-a4b7-2f96c627cf3b
📒 Files selected for processing (3)
CHANGELOG.mdmodelaudit/scanners/pickle_scanner.pytests/scanners/test_pickle_scanner.py
| def _is_safe_import_only_global(mod: str, func: str, ml_context: dict[str, Any] | None = None) -> bool: | ||
| """Return True when an import-only target is explicitly safe to treat as benign.""" | ||
| if _is_actually_dangerous_global(mod, func, ml_context or {}): | ||
| return False | ||
|
|
||
| if _is_safe_ml_global(mod, func): | ||
| return True | ||
|
|
||
| return func in IMPORT_ONLY_SAFE_GLOBALS.get(mod, frozenset()) |
There was a problem hiding this comment.
Don't let deserializers inherit the import-only safe list.
_is_safe_import_only_global() now treats every ML_SAFE_GLOBALS entry as benign once _is_actually_dangerous_global() returns false. That currently whitelists loader gadgets such as dill.load, dill.loads, and joblib._pickle_load, so a standalone reference to them becomes a passing safe_allowlisted check. Keep this helper limited to constructor/data symbols, or move these APIs into ALWAYS_DANGEROUS_FUNCTIONS. 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 1805 - 1813, The helper
_is_safe_import_only_global currently treats entries from _is_safe_ml_global
(ML_SAFE_GLOBALS) as import-only-safe, which lets deserializers like dill.load,
dill.loads, and joblib._pickle_load bypass checks; remove that widening by not
calling _is_safe_ml_global here and only consulting IMPORT_ONLY_SAFE_GLOBALS (or
explicit constructor/data symbols). Fix options: (A) change
_is_safe_import_only_global to return False if
_is_actually_dangerous_global(...) else return func in
IMPORT_ONLY_SAFE_GLOBALS.get(mod, frozenset()) (remove the _is_safe_ml_global
branch), and move any true constructor/data symbols from ML_SAFE_GLOBALS into
IMPORT_ONLY_SAFE_GLOBALS; or (B) add deserializer APIs (e.g., dill.load,
dill.loads, joblib._pickle_load) to ALWAYS_DANGEROUS_FUNCTIONS so they are never
treated as import-only-safe. Ensure you update ML_SAFE_GLOBALS and tests
accordingly.
| def _classify_import_reference( | ||
| mod: str, func: str, ml_context: dict[str, Any] | ||
| ) -> tuple[bool, IssueSeverity | None, str]: | ||
| """Classify a resolved GLOBAL/STACK_GLOBAL import target. | ||
|
|
||
| Returns (is_failure, severity, classification) where classification is one of | ||
| safe_allowlisted, dangerous, unknown_third_party, or unresolved. | ||
| """ | ||
| if not _is_resolved_import_target(mod, func): | ||
| return False, None, "unresolved" | ||
|
|
||
| if _is_actually_dangerous_global(mod, func, ml_context): | ||
| base_sev = IssueSeverity.WARNING if mod in WARNING_SEVERITY_MODULES else IssueSeverity.CRITICAL | ||
| return True, base_sev, "dangerous" | ||
|
|
||
| if _is_safe_import_only_global(mod, func, ml_context): | ||
| return False, None, "safe_allowlisted" | ||
|
|
||
| if not _is_plausible_import_only_module(mod): | ||
| return False, None, "implausible" | ||
|
|
||
| return True, IssueSeverity.WARNING, "unknown_third_party" |
There was a problem hiding this comment.
Case-normalize denylist matching before falling back to unknown_third_party.
This classifier still relies on exact-case blocklist hits. Mixed-case variants of denylisted refs therefore fall through to unknown_third_party, which weakens the new import-only path for malicious payloads that only change casing. 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 1842 - 1863, The
classifier _classify_import_reference currently only matches denylisted refs
with exact case; update it to perform a case-normalized denylist check (e.g.,
normalize mod and func to a canonical case) before falling back to
"unknown_third_party" so mixed-case variants of denylisted modules/functions are
classified as dangerous; do this by checking the normalized pair against the
existing denylist/blocked set (or via a new helper used by
_is_actually_dangerous_global) while preserving the existing calls to
_is_resolved_import_target, _is_safe_import_only_global, and
_is_plausible_import_only_module and reusing WARNING_SEVERITY_MODULES to
determine severity.
1 participant
Summary
import_only=trueValidation
uv run ruff format modelaudit/ tests/uv run ruff check --fix modelaudit/ tests/uv run ruff check modelaudit/ tests/uv run ruff format --check modelaudit/ tests/uv run mypy modelaudit/uv run pytest tests/scanners/test_pickle_scanner.py -quv run pytest -n auto -m "not slow and not integration" --maxfail=1Notes
Summary by CodeRabbit
Bug Fixes
Tests