diff --git a/studio/backend/tests/test_offline_embedding_minimal.py b/studio/backend/tests/test_offline_embedding_minimal.py index 7ff580c36d..ccc6b5f76a 100644 --- a/studio/backend/tests/test_offline_embedding_minimal.py +++ b/studio/backend/tests/test_offline_embedding_minimal.py @@ -51,6 +51,27 @@ def _modules_json(*paths): _COMMIT = "0123456789abcdef0123456789abcdef01234567" +def _fs_case_sensitive(root): + """Whether root's filesystem is case-sensitive (Linux yes; macOS/Windows usually no). The gate + mirrors the loader, whose file lookups follow the same rule, so some cases only exist on one.""" + probe = Path(root) / "_case_probe" + probe.write_text("x") + try: + return not (Path(root) / "_CASE_PROBE").exists() + finally: + probe.unlink() + + +def _requires_case_sensitive_fs(root): + if not _fs_case_sensitive(root): + pytest.skip("requires a case-sensitive filesystem") + + +def _requires_case_insensitive_fs(root): + if _fs_case_sensitive(root): + pytest.skip("requires a case-insensitive filesystem") + + def _make_cache( root, repo_id, @@ -382,6 +403,329 @@ def test_gate_blocks_sharded_pickle(hf_cache): assert _offline_decision("org/shard").blocked is True +def test_gate_blocks_indexed_pickle_shard_in_subdirectory(hf_cache): + # from_pretrained follows weight_map paths relative to the root index, so these nested shards + # are deserialized even though they are not direct children of the load root (iterdir misses + # them). The online gate blocks index-referenced subdir pickles; the offline gate must too. + _make_cache( + hf_cache, + "org/indexed-shard", + { + "pytorch_model.bin.index.json": ( + '{"weight_map": {"layer.weight": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/indexed-shard") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_blocks_indexed_pickle_shard_with_nonstandard_stem(hf_cache): + # The index tells the loader to deserialize this file, so a pickle EXTENSION is enough -- the + # shard's stem need not match the on-disk weight-name heuristic (which only guesses bare files). + _make_cache( + hf_cache, + "org/indexed-odd", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "shards/evil-00001-of-00001.bin"}}', + "shards/evil-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/indexed-odd") + assert decision.blocked is True + assert any(u["path"] == "shards/evil-00001-of-00001.bin" for u in decision.unsafe_files) + + +def test_gate_blocks_safetensors_index_pointing_to_pickle_shard(hf_cache): + # load_state_dict picks safetensors vs torch.load by each shard's own suffix, so a + # model.safetensors.index.json that maps a weight to a .bin shard still deserializes it. The + # index's own existence must not suppress the shard it names. + _make_cache( + hf_cache, + "org/st-index-pickle", + { + "model.safetensors.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/st-index-pickle") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_blocks_indexed_shard_with_no_pickle_extension(hf_cache): + # Transformers torch.loads any indexed shard not ending in .safetensors, so an unconventional + # extensionless name is still a deserialization target. + _make_cache( + hf_cache, + "org/indexed-noext", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "shards/payload"}}', + "shards/payload": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/indexed-noext") + assert decision.blocked is True + assert any(u["path"] == "shards/payload" for u in decision.unsafe_files) + + +_UPPER_INDEX_FILES = { + "PYTORCH_MODEL.BIN.INDEX.JSON": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", +} + + +def test_gate_blocks_uppercase_index_on_case_insensitive_fs(hf_cache): + # On a case-insensitive volume (Windows/macOS) from_pretrained opens an oddly-cased index when it + # requests the canonical lowercase name, so the loader-mirror lookup resolves it and blocks. + _requires_case_insensitive_fs(hf_cache) + _make_cache(hf_cache, "org/upper-index", _UPPER_INDEX_FILES) + with _no_network(): + decision = _offline_decision("org/upper-index") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_allows_uppercase_index_on_case_sensitive_fs(hf_cache): + # On a case-sensitive FS from_pretrained's os.path.isfile of the canonical lowercase name misses + # the uppercase artifact and never loads its shard, so the gate must not over-block it. + _requires_case_sensitive_fs(hf_cache) + _make_cache(hf_cache, "org/upper-index", _UPPER_INDEX_FILES) + with _no_network(): + assert _offline_decision("org/upper-index").blocked is False + + +def test_gate_blocks_indexed_shard_named_with_backslash(hf_cache): + # On POSIX a backslash is a literal filename char, so from_pretrained joins the raw weight_map + # value and deserializes a file actually named "dir\payload.bin"; the gate must probe it verbatim. + import os + + if os.sep != "/": + pytest.skip("backslash is a path separator off POSIX") + _make_cache( + hf_cache, + "org/backslash", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "dir\\\\payload.bin"}}', + "dir\\payload.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/backslash") + assert decision.blocked is True + assert any(u["path"] == "dir\\payload.bin" for u in decision.unsafe_files) + + +def test_gate_blocks_indexed_shard_with_uppercase_safetensors_suffix(hf_cache): + # load_state_dict's endswith(".safetensors") is case-sensitive, so a shard named payload.SAFETENSORS + # falls to torch.load. The gate must classify shard suffixes case-sensitively to match it. + _make_cache( + hf_cache, + "org/upper-suffix", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "shards/payload.SAFETENSORS"}}', + "shards/payload.SAFETENSORS": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/upper-suffix") + assert decision.blocked is True + assert any(u["path"] == "shards/payload.SAFETENSORS" for u in decision.unsafe_files) + + +def test_gate_allows_stale_safetensors_index_beside_direct_safetensors(hf_cache): + # A complete direct model.safetensors is selected before either index, so a stale + # model.safetensors.index.json referencing a .bin shard never deserializes -> must not block. + _make_cache( + hf_cache, + "org/direct-plus-stale-index", + { + "model.safetensors": "tensors", + "model.safetensors.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + assert _offline_decision("org/direct-plus-stale-index").blocked is False + + +def test_gate_blocks_pytorch_index_with_uppercase_safetensors_decoy(hf_cache): + # On a case-sensitive FS, from_pretrained asks for the canonical lowercase model.safetensors, does + # not find an uppercase decoy, and selects the pytorch index instead. The decoy must not suppress. + _requires_case_sensitive_fs(hf_cache) + _make_cache( + hf_cache, + "org/upper-decoy", + { + "MODEL.SAFETENSORS": "decoy", + "pytorch_model.bin.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/upper-decoy") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_blocks_direct_pickle_with_uppercase_safetensors_decoy(hf_cache): + # Same decoy against a direct pytorch_model.bin: the loader selects the pickle, so the uppercase + # safetensors must not suppress it on a case-sensitive FS. + _requires_case_sensitive_fs(hf_cache) + _make_cache( + hf_cache, + "org/upper-decoy-direct", + {"MODEL.SAFETENSORS": "decoy", "pytorch_model.bin": "pickle"}, + ) + with _no_network(): + decision = _offline_decision("org/upper-decoy-direct") + assert decision.blocked is True + assert any(u["path"] == "pytorch_model.bin" for u in decision.unsafe_files) + + +def test_gate_blocks_indexed_pickle_shard_in_module_subdir(hf_cache): + # A weight index inside a sentence-transformers module load root points at a nested pickle shard. + _make_cache( + hf_cache, + "org/mod-indexed", + { + "modules.json": _modules_json("0_Transformer"), + "0_Transformer/pytorch_model.bin.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "0_Transformer/shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/mod-indexed") + assert decision.blocked is True + assert any( + u["path"] == "0_Transformer/shards/pytorch_model-00001-of-00001.bin" + for u in decision.unsafe_files + ) + + +def test_gate_allows_indexed_pickle_shard_with_safetensors_sibling(hf_cache): + # A base model.safetensors makes the loader ignore the pickle index entirely, so it must not + # block (mirrors the direct-file safetensors-sibling suppression). + _make_cache( + hf_cache, + "org/indexed-both", + { + "pytorch_model.bin.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + "model.safetensors": "y", + }, + ) + with _no_network(): + assert _offline_decision("org/indexed-both").blocked is False + + +def test_gate_allows_indexed_safetensors_shard_in_subdirectory(hf_cache): + # A safetensors index lists inert shards -- following it must never block (guards against a + # scanner that flags every indexed shard regardless of format). + _make_cache( + hf_cache, + "org/st-indexed", + { + "model.safetensors.index.json": ( + '{"weight_map": {"w": "shards/model-00001-of-00001.safetensors"}}' + ), + "shards/model-00001-of-00001.safetensors": "tensors", + }, + ) + with _no_network(): + assert _offline_decision("org/st-indexed").blocked is False + + +def test_gate_blocks_on_index_path_traversal(hf_cache): + # A weight_map entry escaping the snapshot via ".." is abnormal/hostile -> fail closed. + _make_cache( + hf_cache, + "org/escape", + {"pytorch_model.bin.index.json": '{"weight_map": {"w": "../../../../etc/evil.bin"}}'}, + ) + with _no_network(): + assert _offline_decision("org/escape").blocked is True + + +def test_gate_allows_symlinked_sharded_safetensors(tmp_path, monkeypatch): + # Real HF caches store snapshot files as symlinks into blobs/. A resolve()-based containment + # check would escape the snapshot and false-block every sharded model; the lexical gate must not. + import hashlib + import os + + from huggingface_hub.file_download import repo_folder_name + + root = tmp_path / "hub" + root.mkdir() + monkeypatch.setenv("HF_HOME", str(tmp_path)) + monkeypatch.setenv("HF_HUB_CACHE", str(root)) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = root), + ) + repo_dir = root / repo_folder_name(repo_id = "org/sym", repo_type = "model") + (repo_dir / "refs").mkdir(parents = True) + (repo_dir / "refs" / "main").write_text(_COMMIT) + blobs = repo_dir / "blobs" + blobs.mkdir() + snapshot = repo_dir / "snapshots" / _COMMIT + (snapshot / "shards").mkdir(parents = True) + + def _blobbed(rel, content): + digest = hashlib.sha256(content.encode()).hexdigest() + (blobs / digest).write_text(content) + target = snapshot / rel + target.parent.mkdir(parents = True, exist_ok = True) + target.symlink_to(os.path.relpath(blobs / digest, target.parent)) + + _blobbed("config.json", "{}") + _blobbed( + "model.safetensors.index.json", + '{"weight_map": {"w": "shards/model-00001-of-00001.safetensors"}}', + ) + _blobbed("shards/model-00001-of-00001.safetensors", "tensors") + with _no_network(): + assert _offline_decision("org/sym").blocked is False + + +def test_gate_allows_index_without_weight_map(hf_cache): + # An index whose top-level JSON has no dict weight_map lets the loader resolve no shards, so it + # must not crash or block on its own (only inert safetensors are cached here). + _make_cache( + hf_cache, + "org/no-wm", + {"model.safetensors.index.json": "[]", "model.safetensors": "x"}, + ) + with _no_network(): + assert _offline_decision("org/no-wm").blocked is False + + def test_gate_allows_nothing_cached(hf_cache): with _no_network(): assert _offline_decision("org/missing").blocked is False diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 91d7ad8f0e..892f7862a9 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -46,17 +46,6 @@ _PICKLE_WEIGHT_RE = re.compile( r"\.(bin|pt|pth|ckpt|pkl|pickle)$", re.IGNORECASE, ) -# Base-model safetensors set: HF names the base pickle pytorch_model.bin but the safetensors -# model.safetensors (stems differ), so a base pickle is replaced only by these, not an adapter's. -_BASE_SAFETENSORS_RE = re.compile( - r"^(model(-\d+-of-\d+)?\.safetensors|model\.safetensors\.index\.json)$", - re.IGNORECASE, -) -# Adapter (PEFT) safetensors set: adapter_model.safetensors, its shards, or index. -_ADAPTER_SAFETENSORS_RE = re.compile( - r"^(adapter_model(-\d+-of-\d+)?\.safetensors|adapter_model\.safetensors\.index\.json)$", - re.IGNORECASE, -) # Non-blocking levels: clean or not-yet-finished. Anything else (unsafe/suspicious/ # malicious or a future label) blocks, so Hub schema drift fails CLOSED. @@ -94,6 +83,13 @@ _INERT_SUFFIXES = frozenset( _SOURCE_SUFFIXES = frozenset({".py", ".pyc", ".pyx", ".pyi"}) +# Torch-family weight indexes: from_pretrained feeds each shard they name to load_state_dict, which +# torch.load()s (pickle) any shard whose name does not end in .safetensors, whatever its stem. A +# pytorch index is superseded when a base safetensors is present (the loader prefers it); a +# safetensors index IS the chosen archive, so a non-safetensors target it names still loads. tf/flax +# indexes load via non-pickle loaders, so they are not a torch.load vector here. +_TORCH_INDEX_FILES = ("pytorch_model.bin.index.json", "model.safetensors.index.json") + # Root weight-index files. from_pretrained reads these to find sharded weights, so a # flagged subdir pickle is a load vector iff a root index references it. _TRANSFORMERS_INDEX_FILES = ( @@ -313,13 +309,72 @@ def _st_load_roots(snapshot: Path) -> list: return roots +def _indexed_pickle_shards(index_path: Path, root: Path, snapshot: Path) -> list: + """Shards a torch weight index points a ``from_pretrained`` load at that load_state_dict would + torch.load (pickle): every ``weight_map`` target NOT ending in ``.safetensors``, whatever its + stem (an arbitrary name like ``shards/payload`` still deserializes). Resolved relative to the + index dir (``root``) like the loader, so a shard in a nested dir is followed (iterdir misses it). + Lexical only, never ``Path.resolve()`` (HF snapshot files symlink into ``blobs/``, so resolving + escapes the snapshot and false-blocks every shard). Raises OSError -> caller fails CLOSED on an + unreadable/invalid index or a target escaping the snapshot.""" + import json + import os + + try: + # JSON is UTF-8 by spec; pin it so a non-ASCII index is not misdecoded (and needlessly + # blocked) under Windows' cp1252 default. + parsed = json.loads(index_path.read_text(encoding = "utf-8")) + except (OSError, ValueError) as exc: + raise OSError(f"unreadable weight index: {index_path}") from exc + weight_map = parsed.get("weight_map") if isinstance(parsed, dict) else None + if not isinstance(weight_map, dict): + return [] # no dict weight_map -> the loader resolves no shards from this index + snapshot_norm = os.path.normpath(str(snapshot)) + shards = [] + for shard in weight_map.values(): + raw = str(shard) + if not raw: + continue + # Join the RAW weight_map value like from_pretrained's os.path.join: on POSIX a backslash is a + # literal filename char (not a separator), so normalizing it would probe a different path than + # the loader opens. normpath + containment stay platform-aware (os.sep) to block "..". + joined = os.path.normpath(os.path.join(str(root), raw)) + if joined != snapshot_norm and not joined.startswith(snapshot_norm + os.sep): + raise OSError(f"weight index escapes the snapshot: {index_path}") + shard_path = Path(joined) + # Case-SENSITIVE, mirroring load_state_dict's own endswith(".safetensors"): a shard named + # payload.SAFETENSORS is not treated as safetensors by the loader and falls to torch.load. + if not shard_path.name.endswith(".safetensors") and shard_path.is_file(): + shards.append(shard_path) + return shards + + +def _loader_resolves(root: Path, name: str) -> bool: + """True iff from_pretrained would open ``name`` under ``root``. ``is_file()`` honors the platform + (case-sensitive on Linux, case-insensitive on Windows/macOS), so it mirrors the loader's own + lookup: an oddly-cased decoy counts as an alternative only where the loader would truly open it. + A name-fold instead would let an uppercase MODEL.SAFETENSORS suppress the scan on Linux while the + loader, asking for the canonical lowercase name, silently falls through to a pickle index.""" + return (root / name).is_file() + + def _cached_pickle_weight_files(snapshot: Path) -> list: - """Pickle weight files in snapshot's ST load roots, EXCLUDING those whose weight family also - ships an inert safetensors in the same dir (the loader prefers it): a base pickle is suppressed - only by a base model.safetensors, an adapter pickle only by adapter_model.safetensors -- an - unrelated safetensors is no substitute. Load roots only. Raises OSError if the snapshot root is - unreadable (caller blocks).""" + """Pickle weight files a SentenceTransformer/Transformers load deserializes from snapshot's ST + load roots, EXCLUDING those whose weight family also ships an inert safetensors in the same dir + (the loader prefers it): a base pickle is suppressed only by a base model.safetensors, an adapter + pickle only by adapter_model.safetensors -- an unrelated safetensors is no substitute. Covers + both direct-child pickles AND pickle shards referenced by a local weight index (which the loader + follows into nested dirs, matching the online gate). Raises OSError -- caller fails CLOSED -- if + the snapshot root or a weight index is unreadable, or an index reference escapes the snapshot.""" blocked = [] + seen = set() + + def _add(path: Path): + key = str(path) + if key not in seen: + seen.add(key) + blocked.append(path) + for root in _st_load_roots(snapshot): try: entries = [p for p in root.iterdir() if p.is_file()] @@ -327,15 +382,35 @@ def _cached_pickle_weight_files(snapshot: Path) -> list: if root == snapshot: raise # top-level unreadable -> fail closed continue # unreadable module subdir: nothing loadable to attest here - has_base_safetensors = any(_BASE_SAFETENSORS_RE.match(p.name) for p in entries) - has_adapter_safetensors = any(_ADAPTER_SAFETENSORS_RE.match(p.name) for p in entries) + # Safetensors alternatives the loader would actually resolve (never a bare name-fold, which + # fails OPEN: see _loader_resolves). A base pickle is replaced only by a base safetensors, an + # adapter pickle only by an adapter one. A single model.safetensors also outranks BOTH indexes. + has_direct_base_safetensors = _loader_resolves(root, "model.safetensors") + has_base_safetensors = has_direct_base_safetensors or _loader_resolves( + root, "model.safetensors.index.json" + ) + has_adapter_safetensors = _loader_resolves(root, "adapter_model.safetensors") for path in entries: if not _PICKLE_WEIGHT_RE.match(path.name): continue is_adapter = path.name.lower().startswith("adapter_model") has_alternative = has_adapter_safetensors if is_adapter else has_base_safetensors if not has_alternative: - blocked.append(path) + _add(path) + # A torch weight index makes from_pretrained load nested shards iterdir never sees; the loader + # torch.loads any not ending in .safetensors. Probe the canonical index name with the loader's + # own lookup (_loader_resolves), so an oddly-cased artifact it would never open does not block. + # A direct model.safetensors wins over BOTH indexes; failing that a base safetensors still + # outranks the pytorch index, while a safetensors index is itself the chosen archive. + for index_name in _TORCH_INDEX_FILES: + if not _loader_resolves(root, index_name): + continue + if has_direct_base_safetensors: + continue + if index_name == "pytorch_model.bin.index.json" and has_base_safetensors: + continue + for shard_path in _indexed_pickle_shards(root / index_name, root, snapshot): + _add(shard_path) return blocked