fix: tighten dill MemoryError downgrade gating

Includes follow-up review fixes to require stronger serialization evidence for MemoryError scanner-limitation downgrades while preserving legitimate dill handling.

Author: mldangelo

modelaudit/scanners/pickle_scanner.py

@@ -2854,6 +2854,76 @@ def _is_legitimate_serialization_file(path: str) -> bool:
     Validate that a file is a legitimate joblib or dill serialization file.
     This helps prevent security bypass by simply renaming malicious files.
     """
+
+    def _analyze_sample_globals(sample: bytes) -> tuple[bool, bool]:
+        if not sample:
+            return False, False
+
+        validation_context = {"is_ml_content": False, "overall_confidence": 0.0, "frameworks": {}}
+        has_dangerous_global = False
+        has_joblib_like_global = False
+
+        def _record_global(mod: str, func: str) -> None:
+            nonlocal has_dangerous_global, has_joblib_like_global
+            if mod in {"joblib", "sklearn", "numpy"} or mod.startswith(("joblib.", "sklearn.", "numpy.")):
+                has_joblib_like_global = True
+            if _is_actually_dangerous_global(mod, func, validation_context):
+                has_dangerous_global = True
+
+        # Raw protocol 0/1 GLOBAL parsing keeps this heuristic usable even when
+        # pickletools itself is the code path that hits MemoryError.
+        cursor = 0
+        max_global_len = 256
+        while cursor < len(sample):
+            global_pos = sample.find(b"c", cursor)
+            if global_pos == -1:
+                break
+
+            module_end = sample.find(b"\n", global_pos + 1, global_pos + 1 + max_global_len)
+            if module_end == -1:
+                cursor = global_pos + 1
+                continue
+
+            function_end = sample.find(b"\n", module_end + 1, module_end + 1 + max_global_len)
+            if function_end == -1:
+                cursor = global_pos + 1
+                continue
+
+            try:
+                module = sample[global_pos + 1 : module_end].decode("utf-8")
+                function = sample[module_end + 1 : function_end].decode("utf-8")
+            except UnicodeDecodeError:
+                cursor = global_pos + 1
+                continue
+
+            if module and function:
+                _record_global(module, function)
+            cursor = function_end + 1
+
+        try:
+            opcodes = list(_genops_with_fallback(io.BytesIO(sample), max_items=128))
+        except (_GenopsBudgetExceeded, ValueError, struct.error, UnicodeDecodeError, EOFError):
+            return has_dangerous_global, has_joblib_like_global
+        except Exception:
+            return has_dangerous_global, has_joblib_like_global
+
+        stack_global_refs, _callable_refs, _origin_refs, _origin_is_ext, _malformed = _simulate_symbolic_reference_maps(
+            opcodes
+        )
+
+        for idx, (opcode, arg, _pos) in enumerate(opcodes):
+            op_name = getattr(opcode, "name", "")
+            if op_name in {"GLOBAL", "INST"} and isinstance(arg, str):
+                parsed = _parse_module_function(arg)
+                if parsed:
+                    _record_global(parsed[0], parsed[1])
+            elif op_name == "STACK_GLOBAL":
+                stack_ref = stack_global_refs.get(idx)
+                if stack_ref:
+                    _record_global(stack_ref[0], stack_ref[1])
+
+        return has_dangerous_global, has_joblib_like_global
+
     try:
         with open(path, "rb") as f:
             # Read first few bytes to check for pickle magic
@@ -2874,34 +2944,24 @@ def _is_legitimate_serialization_file(path: str) -> bool:
                 # Common pickle opcode starts for protocols 0-1
                 return False
 
-            # For joblib files, look for joblib-specific patterns
-            # Also check extensionless files (e.g. HuggingFace cache blob hashes)
+            f.seek(0)
+            sample = f.read(64 * 1024)
+            has_dangerous_global, has_joblib_like_global = _analyze_sample_globals(sample)
+            if has_dangerous_global:
+                return False
+
+            # For joblib files and extensionless cache blobs, require opcode-level
+            # framework evidence instead of marker strings. Extension/substring
+            # checks alone are too easy to spoof.
             ext_lower = os.path.splitext(path)[1].lower()
             if ext_lower == ".joblib" or not ext_lower:
-                f.seek(0)
-                # Try to find joblib-specific markers in first 2KB
-                sample = f.read(2048)
-                # Look for joblib-specific indicators
-                joblib_indicators = [
-                    b"joblib",
-                    b"sklearn",
-                    b"numpy",
-                    b"_joblib",
-                    b"__main__",
-                    b"_pickle",
-                    b"NumpyArrayWrapper",
-                ]
-                if any(marker in sample for marker in joblib_indicators):
-                    return True
-                # For extensionless files, only return False if no indicators found
-                # (don't fall through to dill check)
-                if not ext_lower:
-                    return False
+                return bool(has_joblib_like_global)
 
-            # For dill files, they're usually just enhanced pickle
+            # Dill can serialize plain pickle-compatible objects without
+            # embedding obvious dill globals near the front of the stream, so
+            # a .dill extension remains a legitimacy signal after bounded
+            # dangerous-global rejection above.
             if ext_lower == ".dill":
-                # Dill files should contain standard pickle format
-                # Additional validation could check for dill-specific patterns
                 return True
 
         return False
@@ -5663,6 +5723,10 @@ def get_depth(x):
                 mod in {"joblib", "sklearn", "numpy"} or mod.startswith(("joblib.", "sklearn.", "numpy."))
                 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
+            )
             is_joblib_content = is_serialization_ext or (not file_ext and has_joblib_globals)
 
             # Check for recursion errors on legitimate ML model files
@@ -5703,13 +5767,17 @@ def get_depth(x):
                 and (has_pytorch_advanced_global or has_ordereddict_global)
                 and not has_dangerous_advanced_global
             )
-            # For serialization content, require positive evidence of legitimacy.
-            # Extension-validated files (.joblib/.dill) rely on extension +
-            # _is_legitimate_serialization_file() + no dangerous globals.
-            # Extensionless blobs require positive joblib/sklearn/numpy globals.
-            has_legitimate_serialization_globals = (is_serialization_ext and not has_dangerous_advanced_global) or (
-                not file_ext and bool(advanced_globals) and has_joblib_globals and not has_dangerous_advanced_global
-            )
+            # Require positive opcode-level framework globals for .joblib files;
+            # marker bytes and extensions alone are too weak. Dill stays more
+            # permissive because legitimate plain-object dill payloads may not
+            # expose dill globals before a resource limit hits.
+            has_extension_based_serialization_globals = (
+                file_ext == ".joblib" and bool(advanced_globals) and has_joblib_globals
+            ) or (file_ext == ".dill" and (has_dill_globals or not advanced_globals))
+            has_legitimate_serialization_globals = (
+                has_extension_based_serialization_globals
+                or (not file_ext and bool(advanced_globals) and has_joblib_globals)
+            ) and not has_dangerous_advanced_global
             passes_global_gate = (
                 has_legitimate_pytorch_globals
                 if file_ext == ".bin"

tests/scanners/test_pickle_scanner.py

@@ -2160,6 +2160,207 @@ def _raise_memory_error(*args: object, **kwargs: object) -> object:
     assert not any(issue.severity in {IssueSeverity.WARNING, IssueSeverity.CRITICAL} for issue in result.issues)
 
 
+def test_scan_dill_memory_error_without_dill_globals_not_downgraded(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Dangerous-looking dill prefixes must not qualify for the INFO downgrade."""
+    model_path = tmp_path / "suspicious.dill"
+    model_path.write_bytes(b"\x80\x04cdill\nloads\nq\x00." + b"dill" + b"\x00" * (256 * 1024))
+
+    def _raise_memory_error(*args: object, **kwargs: object) -> object:
+        raise MemoryError("simulated parser memory limit")
+
+    monkeypatch.setattr("modelaudit.scanners.pickle_scanner.pickletools.genops", _raise_memory_error)
+    monkeypatch.setattr(
+        PickleScanner,
+        "_extract_globals_advanced",
+        lambda self, file_obj, multiple_pickles=True, scan_start_time=None: set(),
+    )
+
+    result = PickleScanner().scan(str(model_path))
+
+    assert not any(check.name == "Pickle Parse Resource Limit" for check in result.checks)
+    format_validation_checks = [check for check in result.checks if check.name == "Pickle Format Validation"]
+    assert len(format_validation_checks) == 1
+    assert format_validation_checks[0].status == CheckStatus.FAILED
+    assert format_validation_checks[0].severity == IssueSeverity.WARNING
+    assert format_validation_checks[0].details["exception_type"] == "MemoryError"
+
+
+def test_scan_joblib_memory_error_requires_joblib_globals(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    """Only .joblib files with parsed framework globals should downgrade to INFO."""
+    model_path = tmp_path / "legitimate.joblib"
+    model_path.write_bytes(b"\x80\x04cjoblib.numpy_pickle\nNumpyArrayWrapper\nq\x00." + b"\x00" * (256 * 1024))
+
+    def _raise_memory_error(*args: object, **kwargs: object) -> object:
+        raise MemoryError("simulated parser memory limit")
+
+    monkeypatch.setattr("modelaudit.scanners.pickle_scanner.pickletools.genops", _raise_memory_error)
+    monkeypatch.setattr(
+        PickleScanner,
+        "_extract_globals_advanced",
+        lambda self, file_obj, multiple_pickles=True, scan_start_time=None: {
+            ("joblib.numpy_pickle", "NumpyArrayWrapper", "GLOBAL")
+        },
+    )
+
+    result = PickleScanner().scan(str(model_path))
+
+    resource_limit_checks = [check for check in result.checks if check.name == "Pickle Parse Resource Limit"]
+    assert len(resource_limit_checks) == 1
+    resource_limit_check = resource_limit_checks[0]
+    assert resource_limit_check.status == CheckStatus.FAILED
+    assert resource_limit_check.severity == IssueSeverity.INFO
+    assert resource_limit_check.details["reason"] == "memory_limit_on_legitimate_model"
+    assert resource_limit_check.details["exception_type"] == "MemoryError"
+    assert not any(
+        issue.severity in {IssueSeverity.WARNING, IssueSeverity.CRITICAL}
+        and "Unable to parse pickle file" in issue.message
+        for issue in result.issues
+    )
+
+
+def test_scan_joblib_memory_error_without_joblib_globals_not_downgraded(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Marker bytes alone must not qualify a .joblib file for INFO downgrade."""
+    model_path = tmp_path / "suspicious.joblib"
+    model_path.write_bytes(b"\x80\x04joblibsklearn" + b"\x00" * (256 * 1024))
+
+    def _raise_memory_error(*args: object, **kwargs: object) -> object:
+        raise MemoryError("simulated parser memory limit")
+
+    monkeypatch.setattr("modelaudit.scanners.pickle_scanner.pickletools.genops", _raise_memory_error)
+    monkeypatch.setattr(
+        PickleScanner,
+        "_extract_globals_advanced",
+        lambda self, file_obj, multiple_pickles=True, scan_start_time=None: set(),
+    )
+
+    result = PickleScanner().scan(str(model_path))
+
+    assert not any(check.name == "Pickle Parse Resource Limit" for check in result.checks)
+    format_validation_checks = [check for check in result.checks if check.name == "Pickle Format Validation"]
+    assert len(format_validation_checks) == 1
+    assert format_validation_checks[0].status == CheckStatus.FAILED
+    assert format_validation_checks[0].severity == IssueSeverity.WARNING
+    assert format_validation_checks[0].details["exception_type"] == "MemoryError"
+
+
+def test_scan_dill_memory_error_with_dill_globals_is_informational(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Legitimate dill globals should still allow the scanner-limitation downgrade."""
+    model_path = tmp_path / "legitimate.dill"
+    model_path.write_bytes(b"\x80\x04" + b"\x00" * (256 * 1024))
+
+    def _raise_memory_error(*args: object, **kwargs: object) -> object:
+        raise MemoryError("simulated parser memory limit")
+
+    monkeypatch.setattr("modelaudit.scanners.pickle_scanner.pickletools.genops", _raise_memory_error)
+    monkeypatch.setattr(
+        PickleScanner,
+        "_extract_globals_advanced",
+        lambda self, file_obj, multiple_pickles=True, scan_start_time=None: {("dill", "dump", "GLOBAL")},
+    )
+
+    result = PickleScanner().scan(str(model_path))
+
+    resource_limit_checks = [check for check in result.checks if check.name == "Pickle Parse Resource Limit"]
+    assert len(resource_limit_checks) == 1
+    resource_limit_check = resource_limit_checks[0]
+    assert resource_limit_check.status == CheckStatus.FAILED
+    assert resource_limit_check.severity == IssueSeverity.INFO
+    assert resource_limit_check.details["reason"] == "memory_limit_on_legitimate_model"
+    assert resource_limit_check.details["exception_type"] == "MemoryError"
+    assert not any(issue.severity in {IssueSeverity.WARNING, IssueSeverity.CRITICAL} for issue in result.issues)
+
+
+def test_scan_plain_dill_memory_error_without_globals_is_informational(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Plain-object dill files should keep the scanner-limitation downgrade path."""
+    model_path = tmp_path / "plain.dill"
+    model_path.write_bytes(dill.dumps([1, 2, 3]))
+
+    def _raise_memory_error(*args: object, **kwargs: object) -> object:
+        raise MemoryError("simulated parser memory limit")
+
+    monkeypatch.setattr("modelaudit.scanners.pickle_scanner.pickletools.genops", _raise_memory_error)
+    monkeypatch.setattr(
+        PickleScanner,
+        "_extract_globals_advanced",
+        lambda self, file_obj, multiple_pickles=True, scan_start_time=None: set(),
+    )
+
+    result = PickleScanner().scan(str(model_path))
+
+    resource_limit_checks = [check for check in result.checks if check.name == "Pickle Parse Resource Limit"]
+    assert len(resource_limit_checks) == 1
+    resource_limit_check = resource_limit_checks[0]
+    assert resource_limit_check.status == CheckStatus.FAILED
+    assert resource_limit_check.severity == IssueSeverity.INFO
+    assert resource_limit_check.details["reason"] == "memory_limit_on_legitimate_model"
+    assert resource_limit_check.details["exception_type"] == "MemoryError"
+    assert not any(issue.severity in {IssueSeverity.WARNING, IssueSeverity.CRITICAL} for issue in result.issues)
+
+
+def test_scan_dill_memory_error_with_internal_dill_globals_is_informational(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Internal dill globals should qualify for the scanner-limitation downgrade."""
+    model_path = tmp_path / "legitimate.dill"
+    model_path.write_bytes(b"\x80\x04" + b"\x00" * (256 * 1024))
+
+    def _raise_memory_error(*args: object, **kwargs: object) -> object:
+        raise MemoryError("simulated parser memory limit")
+
+    monkeypatch.setattr("modelaudit.scanners.pickle_scanner.pickletools.genops", _raise_memory_error)
+    monkeypatch.setattr(
+        PickleScanner,
+        "_extract_globals_advanced",
+        lambda self, file_obj, multiple_pickles=True, scan_start_time=None: {("_dill", "dump", "GLOBAL")},
+    )
+
+    result = PickleScanner().scan(str(model_path))
+
+    resource_limit_checks = [check for check in result.checks if check.name == "Pickle Parse Resource Limit"]
+    assert len(resource_limit_checks) == 1
+    resource_limit_check = resource_limit_checks[0]
+    assert resource_limit_check.status == CheckStatus.FAILED
+    assert resource_limit_check.severity == IssueSeverity.INFO
+    assert resource_limit_check.details["reason"] == "memory_limit_on_legitimate_model"
+    assert resource_limit_check.details["exception_type"] == "MemoryError"
+    assert not any(issue.severity in {IssueSeverity.WARNING, IssueSeverity.CRITICAL} for issue in result.issues)
+
+
+def test_scan_joblib_memory_error_with_dangerous_prefix_not_downgraded(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Marker bytes must not hide a dangerous pickle prefix in .joblib files."""
+    model_path = tmp_path / "suspicious.joblib"
+    model_path.write_bytes(b"\x80\x02cbuiltins\neval\nq\x00." + b"joblibsklearnnumpy" + b"\x00" * (256 * 1024))
+
+    def _raise_memory_error(*args: object, **kwargs: object) -> object:
+        raise MemoryError("simulated parser memory limit")
+
+    monkeypatch.setattr("modelaudit.scanners.pickle_scanner.pickletools.genops", _raise_memory_error)
+    monkeypatch.setattr(
+        PickleScanner,
+        "_extract_globals_advanced",
+        lambda self, file_obj, multiple_pickles=True, scan_start_time=None: set(),
+    )
+
+    result = PickleScanner().scan(str(model_path))
+
+    assert not any(check.name == "Pickle Parse Resource Limit" for check in result.checks)
+    format_validation_checks = [check for check in result.checks if check.name == "Pickle Format Validation"]
+    assert len(format_validation_checks) == 1
+    assert format_validation_checks[0].status == CheckStatus.FAILED
+    assert format_validation_checks[0].severity == IssueSeverity.WARNING
+    assert format_validation_checks[0].details["exception_type"] == "MemoryError"
+
+
 def test_scan_memory_error_with_dangerous_globals_not_downgraded(
     tmp_path: Path, monkeypatch: pytest.MonkeyPatch
 ) -> None:
@@ -2179,7 +2380,7 @@ def _raise_memory_error(*args: object, **kwargs: object) -> object:
     monkeypatch.setattr(
         PickleScanner,
         "_extract_globals_advanced",
-        lambda self, file_obj, multiple_pickles=True: {("builtins", "eval", "GLOBAL")},
+        lambda self, file_obj, multiple_pickles=True, scan_start_time=None: {("builtins", "eval", "GLOBAL")},
     )
 
     result = PickleScanner().scan(str(model_path))