Studio: fail closed on index-referenced nested pickle shards in the offline embedding gate#7366
Conversation
…ffline embedding gate The offline embedding security gate (HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE) only scanned the direct files of each SentenceTransformer load root and never parsed local weight indexes, so a cached snapshot whose pytorch_model.bin.index.json maps a weight to a nested shard (e.g. shards/pytorch_model-00001-of-00001.bin) was treated as inert and allowed. The loader then follows the index into the subdir and unpickles the shard. The online gate already blocks index-referenced subdir pickles, so the offline path was strictly weaker. Parse each local weight index in a load root and follow weight_map into nested dirs, flagging any referenced pickle-extension shard. Paths resolve lexically (normpath), never Path.resolve(), since HF cache snapshot files symlink into blobs/ and resolving would leave the snapshot dir and false-block every sharded model offline. An absolute path, a .. traversal that escapes the snapshot, or an unreadable/invalid index fails closed. The existing safetensors-sibling suppression is kept.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 22a8e16636
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if not has_base_safetensors: | ||
| _add(shard_path) |
There was a problem hiding this comment.
When an offline snapshot contains model.safetensors.index.json whose weight_map points to a .bin/.pt shard, has_base_safetensors is true solely because that index exists, so this condition drops the shard and the gate allows the load. Transformers selects that safetensors-named index and then loads the shard paths it names based on their actual suffixes, so a hostile index can route the supposedly safe load to a pickle shard and bypass this fail-closed check. Only suppress shards from the PyTorch index when a valid safetensors alternative is selected, or require .safetensors targets for the safetensors index.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in ef15858: a safetensors index no longer suppresses shards it names. load_state_dict chooses safetensors vs torch.load per shard by the shard suffix, so only the pytorch index is superseded by a base safetensors now; a model.safetensors.index.json is the chosen archive, and any non-safetensors target it references is flagged.
| if shard_path.suffix.lower() in _PICKLE_WEIGHT_SUFFIXES and shard_path.is_file(): | ||
| shards.append(shard_path) |
There was a problem hiding this comment.
A pytorch_model.bin.index.json can legally map weights to filenames without one of these suffixes (for example shards/weights.data or shards/payload), and Transformers passes every indexed filename to its PyTorch loader unless it ends in .safetensors. This suffix allowlist therefore returns no pickle for an offline cached malicious shard with an unconventional name and allows the load, despite the index making that file a deserialization target. Classify indexed shards by the selected index format (or treat every non-safetensors target of a PyTorch index as unsafe) rather than requiring a conventional pickle extension.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in ef15858: indexed shards are now flagged when the target name does not end in .safetensors (what the loader torch.loads), instead of a pickle-extension allowlist, so arbitrary names like shards/payload or weights.data are caught.
…e extension load_state_dict picks safetensors vs torch.load per shard by the shard's own suffix, so two offline-gate gaps remained: - A model.safetensors.index.json whose weight_map points at a .bin shard was suppressed by has_base_safetensors (the index file itself matches the base safetensors regex), yet Transformers still torch.loads that shard. Only the pytorch index is superseded by a base safetensors now; a safetensors index is the chosen archive, so its non-safetensors targets are always flagged. - A pytorch index can map weights to arbitrary names (shards/payload, weights.data); the loader torch.loads any target not ending in .safetensors. Flag indexed shards by that rule instead of a pickle-extension allowlist. Restrict the scan to the two torch-family indexes (tf/flax load via non-pickle loaders). Add regression tests for both cases.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ef15858479
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| for index_path in entries: | ||
| if index_path.name not in _TORCH_INDEX_FILES: | ||
| continue |
There was a problem hiding this comment.
On Windows and case-insensitive macOS volumes, from_pretrained can open a cache file such as PYTORCH_MODEL.BIN.INDEX.JSON when it requests the lowercase canonical filename, but this exact-case comparison skips parsing it. A malicious nested .bin shard referenced by that index is consequently allowed through the offline gate and then deserialized; compare a normalized lowercase name here, as the surrounding weight/safetensors checks already do.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 56faf42: the index-name match now lower-cases the on-disk filename before comparing, matching the case-insensitive weight and safetensors checks, so an oddly-cased index that a case-insensitive volume would open is parsed and its shards flagged.
The index-name check compared the on-disk filename exactly, while the surrounding weight and safetensors matches use case-insensitive rules. On a case-insensitive volume (Windows or macOS) from_pretrained opens an oddly-cased cache file such as PYTORCH_MODEL.BIN.INDEX.JSON when it requests the canonical lowercase name, so the exact-case check skipped it and a nested pickle shard it referenced was allowed through. Lower-case the index name before matching, as the rest of the gate does, and add a regression test.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 56faf42593
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| name = index_path.name.lower() | ||
| if name not in _TORCH_INDEX_FILES: | ||
| continue | ||
| if name == "pytorch_model.bin.index.json" and has_base_safetensors: |
There was a problem hiding this comment.
When a load root contains a complete model.safetensors, from_pretrained selects that direct file before either sharded index, but this condition suppresses only the PyTorch index. Consequently, a cache containing model.safetensors plus a stale model.safetensors.index.json that references a .bin shard is blocked offline even though that shard is never deserialized, defeating the existing safetensors-sibling allowance for an otherwise loadable model. Distinguish a direct safetensors file from a safetensors index and skip both indexes when the direct file wins.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in fec6699: a direct model.safetensors now suppresses BOTH indexes, since from_pretrained selects it first; a stale model.safetensors.index.json referencing a .bin shard beside a complete model.safetensors is no longer over-blocked.
| 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) | ||
| if not shard_path.name.lower().endswith(".safetensors") and shard_path.is_file(): |
There was a problem hiding this comment.
transformers.modeling_utils.load_state_dict uses a case-sensitive endswith(".safetensors") check; a shard such as shards/payload.SAFETENSORS therefore falls through to torch.load. This lowercase comparison instead treats that shard as inert, so a canonical weight index can reference a pickle stored under an uppercase suffix and the offline gate will allow the deserialization.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in fec6699: shard suffixes are now classified case-sensitively (shard_path.name.endswith('.safetensors')), matching load_state_dict, so an uppercase .SAFETENSORS shard the loader would torch.load is flagged.
… index scan
Two edge cases in the offline weight-index scan:
- load_state_dict decides safetensors vs torch.load with a case-sensitive
endswith(".safetensors"), so a shard named payload.SAFETENSORS still
deserializes via torch.load. Classify indexed shard suffixes case-sensitively
to match, instead of lower-casing (which treated such a shard as inert).
- A complete direct model.safetensors is selected before either sharded index,
so a stale model.safetensors.index.json referencing a .bin shard never loads.
Skip both indexes when a direct model.safetensors is present, so an otherwise
loadable model is not over-blocked.
Add regression tests for both.
|
@codex review |
|
Codex Review: Didn't find any major issues. Keep it up! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
Path.read_text() uses the locale default, which is cp1252 on Windows, so a UTF-8 weight index with non-ASCII bytes raised UnicodeDecodeError and the gate blocked an otherwise loadable model. JSON is UTF-8 by spec (and how the loader reads it), so pin the encoding.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 01900ceb88
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| 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) | ||
| # A complete direct model.safetensors outranks BOTH sharded indexes in from_pretrained. | ||
| has_direct_base_safetensors = any(p.name.lower() == "model.safetensors" for p in entries) |
There was a problem hiding this comment.
On Linux and other case-sensitive filesystems, a snapshot containing MODEL.SAFETENSORS plus a canonical pytorch_model.bin.index.json whose shard is a pickle will be allowed here: this case-insensitive check suppresses index scanning. Transformers looks up the canonical lowercase model.safetensors, does not find the uppercase file, then selects the PyTorch index and deserializes its shard. Restrict this suppression to the exact filename on case-sensitive platforms (or otherwise mirror the loader's filename resolution) so an uppercase decoy cannot bypass the offline security gate.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in d9d979c: safetensors alternatives are now probed with (root / name).is_file() instead of a case-folded listing, mirroring from_pretrained's own lookup. is_file() honors the platform's case rules, so an uppercase MODEL.SAFETENSORS decoy suppresses the pickle scan only where the loader would actually open it; on a case-sensitive FS the pytorch pickle is now correctly flagged.
…e lookup The offline gate decided a safetensors alternative existed by case-folding the directory listing. On a case-sensitive filesystem that let an uppercase decoy such as MODEL.SAFETENSORS suppress the pickle scan, yet from_pretrained asks for the canonical lowercase model.safetensors, does not find the decoy, and selects the pickle (a direct pytorch_model.bin or the pytorch index) and deserializes it. Probe each alternative with (root / name).is_file() instead, mirroring the loader: is_file() honors the platform's case rules, so a decoy suppresses only where the loader would truly open it. Suppression must never fail open; detection stays case-insensitive (fail closed). Add regression tests for the direct and indexed pickle decoys (skipped on case-insensitive volumes, where no bypass exists).
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d9d979c559
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| name = index_path.name.lower() | ||
| if name not in _TORCH_INDEX_FILES: |
There was a problem hiding this comment.
On case-sensitive filesystems, lowercasing here makes an artifact such as PYTORCH_MODEL.BIN.INDEX.JSON look like a loadable index and blocks any referenced .bin shard, even though from_pretrained checks the canonical lowercase filename with os.path.isfile and will not load that uppercase file. This causes offline embedding loads to be rejected for snapshots containing an oddly-cased, unused index; only treat such a file as an index when the canonical path resolves on the current filesystem.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 96e8b0c: index detection now probes the canonical name with the loader's own is_file lookup (_loader_resolves) rather than case-folding the listing, so an uppercase artifact the loader would never open no longer blocks on a case-sensitive filesystem, while a canonical/odd-cased index the loader does open on a case-insensitive volume still blocks.
| for shard in weight_map.values(): | ||
| rel = _normalize_repo_path(str(shard)) |
There was a problem hiding this comment.
On POSIX, a backslash is a valid filename character, not a path separator. An index can therefore name shards\\payload.bin and the cached snapshot can contain that exact pickle filename; Transformers joins the raw weight_map value and deserializes it, while this normalization changes it to shards/payload.bin, fails is_file(), and allows the offline load. Keep the raw value for the filesystem lookup (and apply containment checks using the platform's path semantics) so this cannot bypass the fail-closed gate.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 96e8b0c: the shard lookup now joins the raw weight_map value with os.path.join instead of normalizing backslashes, so on POSIX a literal name like dir\payload.bin is probed exactly as from_pretrained opens it. Containment stays platform-aware via normpath/os.sep.
Two more loader-fidelity gaps in the offline index scan: - Shard lookup normalized backslashes to forward slashes. On POSIX a backslash is a literal filename character, so an index naming dir\payload.bin matches a real pickle of that exact name that Transformers joins and deserializes, while the normalized dir/payload.bin missed it. Join the raw weight_map value with os.path.join so the probe mirrors the loader on each platform. - Index detection case-folded the directory listing, so on a case-sensitive filesystem an uppercase PYTORCH_MODEL.BIN.INDEX.JSON artifact the loader never opens was treated as live and its shard blocked. Probe the canonical name with the loader's own is_file lookup instead, so an index counts only where from_pretrained would actually load it. Update the uppercase-index tests to assert the correct per-filesystem behavior and add a POSIX backslash-shard regression test.
|
@codex review |
|
Codex Review: Didn't find any major issues. More of your lovely PRs please. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
2 participants
Summary
The offline embedding security gate (
HF_HUB_OFFLINE/TRANSFORMERS_OFFLINE) missed pickle shards that a local weight index references from a nested directory, so they could be deserialized by SentenceTransformer/Transformers without a scan. The online gate already blocks these, so this brings the offline path in line.Details
evaluate_file_security(..., local_only_load=True)routes to_evaluate_local_onlythen_cached_pickle_weight_files, which only iterated the direct files of each SentenceTransformer load root (the snapshot root plusmodules.jsonmodule dirs). It never parsed localpytorch_model.bin.index.json, so a cached snapshot whose root index maps a weight toshards/pytorch_model-00001-of-00001.binwas treated as inert and allowed.from_pretrainedthen follows the index into the subdir and unpickles the shard.The online gate (
_indexed_shard_pathsplus themaybe_shardhandling) already treats index-referenced subdir pickles as load-path vectors, so the offline path was strictly weaker.Fix
_indexed_pickle_shards: parse each local weight index in a load root and followweight_mapinto nested dirs, flagging any referenced shard with a pickle extension (its stem need not match the on-disk name heuristic, since the index tells the loader to deserialize it).os.path.normpath, neverPath.resolve(). HF cache snapshot files are symlinks intoblobs/, soresolve()would leave the snapshot dir and false-block every sharded model offline. An absolute path, a..traversal that escapes the snapshot, or an unreadable/invalid index fails closed (blocks).Tests
Added cases to
test_offline_embedding_minimal.py:..traversal inweight_mapweight_mapThe four blocking tests fail against the pre-fix scanner. The offline embedding, file security, and security gate suites all pass (122).