fix: harden pickle source resolution precedence

Author: mldangelo-oai

packages/modelaudit-picklescan/src/modelaudit_picklescan/call_graph.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import ast
+import fnmatch
 import hashlib
 import os
 import sys
@@ -46,9 +47,7 @@
 _MAX_ASSIGNMENT_ALIASES = 128
 _MAX_ASSIGNMENT_ALIAS_PASSES = 256
 _MAX_FUNCTION_INSTANCE_ALIASES = 32
-_TRUSTED_META_PATH_FINDERS = tuple(sys.meta_path)
 _TRUSTED_PATH_HOOKS = tuple(sys.path_hooks)
-_TRUSTED_PATH_IMPORTERS = tuple(finder for finder in sys.path_importer_cache.values() if finder is not None)
 _MAX_CLASS_INSTANCE_ALIASES = 128
 _MAX_INHERITED_CLASS_METHODS = 128
 _MAX_WILDCARD_IMPORTS = 16
@@ -1047,17 +1046,17 @@ def _find_module_spec_without_imports(module_name: str) -> ModuleSpec | None:
     if isinstance(loaded_spec, ModuleSpec):
         return loaded_spec
 
-    if not _untrusted_meta_path_finder_precedes(BuiltinImporter):
+    if not _untrusted_meta_path_finder_precedes(BuiltinImporter, module_name):
         builtin_spec = BuiltinImporter.find_spec(module_name)
         if builtin_spec is not None:
             return builtin_spec
 
-    if not _untrusted_meta_path_finder_precedes(FrozenImporter):
+    if not _untrusted_meta_path_finder_precedes(FrozenImporter, module_name):
         frozen_spec = FrozenImporter.find_spec(module_name)
         if frozen_spec is not None:
             return frozen_spec
 
-    if _untrusted_meta_path_finder_precedes(PathFinder) or _has_untrusted_path_hook():
+    if _untrusted_meta_path_finder_precedes(PathFinder, module_name) or _has_untrusted_path_hook():
         return None
 
     return _find_standard_filesystem_spec(module_name)
@@ -1084,40 +1083,58 @@ def _find_standard_filesystem_spec(module_name: str) -> ModuleSpec | None:
     return spec
 
 
-def _untrusted_meta_path_finder_precedes(target: object) -> bool:
+def _matches_loaded_finder_type(finder: object, module_name: str, type_name: str) -> bool:
+    module = sys.modules.get(module_name)
+    finder_type = getattr(module, type_name, None) if module is not None else None
+    return isinstance(finder_type, type) and type(finder) is finder_type
+
+
+def _known_meta_path_finder_cannot_handle(finder: object, module_name: str) -> bool:
+    root_name = module_name.split(".", maxsplit=1)[0]
+    if _matches_loaded_finder_type(finder, "_distutils_hack", "DistutilsMetaFinder"):
+        return root_name not in {"distutils", "pip", "test"}
+
+    if _matches_loaded_finder_type(finder, "_virtualenv", "_Finder"):
+        virtualenv_module = sys.modules.get("_virtualenv")
+        patched_modules = getattr(virtualenv_module, "_DISTUTILS_PATCH", ()) if virtualenv_module is not None else ()
+        return module_name not in patched_modules
+
+    if _matches_loaded_finder_type(finder, "_pytest.assertion.rewrite", "AssertionRewritingHook"):
+        if module_name == "conftest":
+            return False
+        must_rewrite = getattr(finder, "_must_rewrite", ())
+        if any(module_name == name or module_name.startswith(f"{name}.") for name in must_rewrite):
+            return False
+        patterns = getattr(finder, "fnpats", ())
+        module_filename = f"{module_name.rsplit('.', maxsplit=1)[-1]}.py"
+        return all(not fnmatch.fnmatchcase(module_filename, pattern) for pattern in patterns)
+
+    return False
+
+
+def _untrusted_meta_path_finder_precedes(target: object, module_name: str) -> bool:
     for finder in sys.meta_path:
         if finder is target:
             return False
         if finder is BuiltinImporter or finder is FrozenImporter or finder is PathFinder:
             continue
-        if any(finder is trusted_finder for trusted_finder in _TRUSTED_META_PATH_FINDERS):
+        if _known_meta_path_finder_cannot_handle(finder, module_name):
             continue
         return True
     return True
 
 
 def _is_standard_path_hook(hook: object) -> bool:
-    if hook is zipimporter:
-        return True
-    return getattr(hook, "__module__", None) == "_frozen_importlib_external" and getattr(
-        hook, "__qualname__", ""
-    ).endswith("FileFinder.path_hook.<locals>.path_hook_for_FileFinder")
+    return hook is zipimporter or any(hook is trusted_hook for trusted_hook in _TRUSTED_PATH_HOOKS)
 
 
 def _has_untrusted_path_hook() -> bool:
-    if any(
-        not _is_standard_path_hook(hook) and not any(hook is trusted_hook for trusted_hook in _TRUSTED_PATH_HOOKS)
-        for hook in sys.path_hooks
-    ):
+    if any(not _is_standard_path_hook(hook) for hook in sys.path_hooks):
         return True
     for entry in sys.path:
         cache_key = entry or os.getcwd()
         finder = sys.path_importer_cache.get(cache_key)
-        if (
-            finder is not None
-            and not isinstance(finder, (FileFinder, zipimporter))
-            and not any(finder is trusted_finder for trusted_finder in _TRUSTED_PATH_IMPORTERS)
-        ):
+        if finder is not None and not isinstance(finder, (FileFinder, zipimporter)):
             return True
     return False
 
@@ -1130,6 +1147,13 @@ def _find_standard_path_spec(module_name: str, search_path: list[str]) -> Module
         (SourcelessFileLoader, BYTECODE_SUFFIXES),
     )
     for entry in search_path:
+        try:
+            zip_spec = zipimporter(entry).find_spec(module_name)
+        except ImportError:
+            zip_spec = None
+        if zip_spec is not None:
+            return zip_spec
+
         finder = FileFinder(entry, *loader_details)
         spec = finder.find_spec(module_name)
         if spec is None:
@@ -3577,7 +3601,7 @@ def _resolve_module_source(module_name: str) -> Path | None:
                 return loaded_source_path
         if loaded_spec.origin not in {"built-in", "frozen"}:
             return None
-    elif _untrusted_meta_path_finder_precedes(PathFinder) or _has_untrusted_path_hook():
+    elif _untrusted_meta_path_finder_precedes(PathFinder, module_name) or _has_untrusted_path_hook():
         return None
 
     spec = _find_standard_filesystem_spec(module_name)

packages/modelaudit-picklescan/tests/test_call_graph_import_statements.py

@@ -2048,6 +2048,58 @@ def find_spec(
     assert not marker.exists()
 
 
+def test_call_graph_fails_closed_for_meta_path_finder_installed_before_import(tmp_path: Path) -> None:
+    module_dir = tmp_path / "modules"
+    module_dir.mkdir()
+    module_name = "modelaudit_tp_preimport_meta_path_probe"
+    (module_dir / f"{module_name}.py").write_text(
+        "def invoke(command):\n    return command\n",
+        encoding="utf-8",
+    )
+    marker = tmp_path / "preimport_meta_path_called"
+    child_code = """
+import sys
+from importlib.machinery import ModuleSpec
+from pathlib import Path
+
+module_dir = Path(sys.argv[1])
+marker = Path(sys.argv[2])
+module_name = sys.argv[3]
+
+class PreexistingMetaPathFinder:
+    @staticmethod
+    def find_spec(fullname, path=None, target=None):
+        del path, target
+        if fullname == module_name:
+            marker.write_text(fullname, encoding="utf-8")
+            return ModuleSpec(fullname, loader=None, origin="custom://module")
+        return None
+
+sys.path.insert(0, str(module_dir))
+sys.meta_path.insert(0, PreexistingMetaPathFinder())
+from modelaudit_picklescan.call_graph import _call_graph_source_unavailable_reason
+
+reason = _call_graph_source_unavailable_reason(module_name)
+if reason != "source_unavailable":
+    raise SystemExit(f"unexpected source reason: {reason!r}")
+if marker.exists():
+    raise SystemExit("pre-existing finder was invoked")
+"""
+
+    result = subprocess.run(
+        [sys.executable, "-c", child_code, str(module_dir), str(marker), module_name],
+        cwd=str(tmp_path),
+        env=dict(os.environ),
+        check=False,
+        capture_output=True,
+        text=True,
+        timeout=30,
+    )
+
+    assert result.returncode == 0, result.stderr
+    assert not marker.exists()
+
+
 def test_scan_bytes_analyzes_stdlib_source_modules_without_custom_meta_path_finders(
     tmp_path: Path,
     monkeypatch: pytest.MonkeyPatch,
@@ -2305,6 +2357,84 @@ def custom_path_hook(path: str) -> object:
     assert not marker.exists()
 
 
+def test_call_graph_rejects_path_hook_with_spoofed_standard_metadata(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    module_name = "modelaudit_tp_spoofed_path_hook_probe"
+    module_dir = tmp_path / "modules"
+    module_dir.mkdir()
+    (module_dir / f"{module_name}.py").write_text("def invoke():\n    return None\n", encoding="utf-8")
+    marker = tmp_path / "spoofed_path_hook_called"
+
+    def spoofed_path_hook(path: str) -> object:
+        marker.write_text(path, encoding="utf-8")
+        raise ImportError
+
+    spoofed_path_hook.__module__ = "_frozen_importlib_external"
+    spoofed_path_hook.__qualname__ = "FileFinder.path_hook.<locals>.path_hook_for_FileFinder"
+    monkeypatch.syspath_prepend(str(module_dir))
+    monkeypatch.setattr(sys, "path_hooks", [spoofed_path_hook, *sys.path_hooks])
+    sys.path_importer_cache.pop(str(module_dir), None)
+    _clear_call_graph_caches()
+
+    try:
+        assert call_graph._call_graph_source_unavailable_reason(module_name) == "source_unavailable"
+    finally:
+        _clear_call_graph_caches()
+
+    assert not marker.exists()
+
+
+def test_call_graph_honors_zipimport_precedence_over_later_extension(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    module_name = "modelaudit_tp_zip_shadowing_extension"
+    zip_path = tmp_path / "modules.zip"
+    extension_dir = tmp_path / "extension"
+    extension_dir.mkdir()
+    with zipfile.ZipFile(zip_path, "w") as archive:
+        archive.writestr(f"{module_name}.py", "def invoke(command):\n    return command\n")
+    (extension_dir / f"{module_name}{EXTENSION_SUFFIXES[0]}").write_bytes(b"not imported")
+    monkeypatch.setattr(sys, "path", [str(zip_path), str(extension_dir), *sys.path])
+    importlib.invalidate_caches()
+    _clear_call_graph_caches()
+
+    try:
+        report = scan_bytes(
+            _global_call_payload(module_name, "invoke", _unicode_operand("echo hidden")),
+            source="zip-shadowing-extension.pkl",
+        )
+    finally:
+        _clear_call_graph_caches()
+
+    assert report.status == ScanStatus.INCONCLUSIVE
+    assert report.verdict == SafetyVerdict.UNKNOWN
+    assert _has_call_graph_source_unavailable_notice(report, module_name, "invoke", "source_unavailable")
+
+
+def test_call_graph_keeps_extension_precedence_over_later_zipimport(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    module_name = "modelaudit_tp_extension_shadowing_zip"
+    extension_dir = tmp_path / "extension"
+    extension_dir.mkdir()
+    zip_path = tmp_path / "modules.zip"
+    (extension_dir / f"{module_name}{EXTENSION_SUFFIXES[0]}").write_bytes(b"not imported")
+    with zipfile.ZipFile(zip_path, "w") as archive:
+        archive.writestr(f"{module_name}.py", "def invoke(command):\n    return command\n")
+    monkeypatch.setattr(sys, "path", [str(extension_dir), str(zip_path), *sys.path])
+    importlib.invalidate_caches()
+    _clear_call_graph_caches()
+
+    try:
+        assert call_graph._call_graph_source_unavailable_reason(module_name) is None
+    finally:
+        _clear_call_graph_caches()
+
+
 def test_scan_bytes_marks_bytecode_only_invoked_call_graph_source_unavailable(
     tmp_path: Path,
     monkeypatch: pytest.MonkeyPatch,
@@ -2614,6 +2744,7 @@ def test_scan_bytes_analyzes_shadowed_torch_extension_callable_invocation(
         "import os\n\ndef device(command):\n    os.system(command)\n    return command\n",
         encoding="utf-8",
     )
+    monkeypatch.delitem(sys.modules, "torch", raising=False)
     monkeypatch.syspath_prepend(str(module_dir))
     importlib.invalidate_caches()
     _clear_call_graph_caches()
@@ -2658,6 +2789,7 @@ def __call__(self, command):
 """,
         encoding="utf-8",
     )
+    monkeypatch.delitem(sys.modules, "torch", raising=False)
     monkeypatch.syspath_prepend(str(module_dir))
     importlib.invalidate_caches()
     _clear_call_graph_caches()
@@ -2745,6 +2877,7 @@ def test_call_graph_analyzes_shadowed_torch_storage_persistent_id_reference(
         "import os\n\ndef __getattr__(name):\n    os.system(name)\n    raise AttributeError(name)\n",
         encoding="utf-8",
     )
+    monkeypatch.delitem(sys.modules, "torch", raising=False)
     monkeypatch.syspath_prepend(str(module_dir))
     importlib.invalidate_caches()
     _clear_call_graph_caches()