[Fix] Validate shard filenames in sharded checkpoint index files (#4033)

* Validate shard filenames in sharded checkpoint index files

Reject shard references with path traversal or mismatched extensions
to prevent a crafted safetensors index from loading pickle payloads.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* remove unrelated tests

* Update src/huggingface_hub/serialization/_torch.py

Co-authored-by: célina <hanouticelina@gmail.com>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: célina <hanouticelina@gmail.com>

Author: 3 people

src/huggingface_hub/serialization/_torch.py

@@ -510,14 +510,32 @@ def _load_sharded_checkpoint(
     with open(index_file, encoding="utf-8") as f:
         index = json.load(f)
 
-    # 2. Validate keys if in strict mode
+    # 2. Validate shard filenames from the index
+    # This prevents path traversal attacks and extension confusion attacks
+    # (e.g. a safetensors index referencing .bin pickle files)
+    expected_extension = Path(filename_pattern.format(suffix="")).suffix  # e.g. ".safetensors"
+    shard_files = list(set(index["weight_map"].values()))
+    for shard_file in shard_files:
+        # Reject path traversal (e.g. "../malicious.bin", absolute paths)
+        if os.path.isabs(shard_file) or ".." in Path(shard_file).parts:
+            raise ValueError(
+                f"Invalid shard filename '{shard_file}' in index file '{index_file}'. "
+                "Shard filenames must be relative paths without '..' components."
+            )
+        # Reject extension mismatch (e.g. .bin shard in a .safetensors index)
+        if not shard_file.endswith(expected_extension):
+            raise ValueError(
+                f"Invalid shard filename '{shard_file}' in index file '{index_file}'. "
+                f"Expected '{expected_extension}' extension to match the index format."
+            )
+
+    # 3. Validate keys if in strict mode
     # This is done before loading any shards to fail fast
     if strict:
         _validate_keys_for_strict_loading(model, index["weight_map"].keys())
 
-    # 3. Load each shard using `load_state_dict`
+    # 4. Load each shard using `load_state_dict`
     # Get unique shard files (multiple parameters can be in same shard)
-    shard_files = list(set(index["weight_map"].values()))
     for shard_file in shard_files:
         # Load shard into memory
         shard_path = os.path.join(save_directory, shard_file)
@@ -531,7 +549,7 @@ def _load_sharded_checkpoint(
         # Explicitly remove the state dict from memory
         del state_dict
 
-    # 4. Return compatibility info
+    # 5. Return compatibility info
     loaded_keys = set(index["weight_map"].keys())
     model_keys = set(model.state_dict().keys())
     return _IncompatibleKeys(

tests/test_serialization.py

@@ -828,3 +828,53 @@ def __init__(self):
     load_torch_model(model, tmp_path, safe=safe, filename_pattern=filename_pattern)
     mock_load.assert_called_once()
     assert mock_load.call_args.kwargs["filename_pattern"] == expected_filename_pattern
+
+
+class TestShardedCheckpointValidation:
+    """Regression tests for shard filename validation in sharded checkpoint loading.
+
+    See https://github.com/huggingface/hackerone/issues/141 for more details.
+
+    Ensures that crafted index files cannot trick the loader into deserializing
+    unsafe pickle payloads or accessing files outside the checkpoint directory.
+    """
+
+    def test_safetensors_index_rejects_bin_shard(self, tmp_path):
+        """A safetensors index file referencing a .bin shard must be rejected."""
+        index = {
+            "metadata": {"total_size": 100},
+            "weight_map": {
+                "layer_1": "model-00001-of-00001.bin",
+            },
+        }
+        (tmp_path / "model.safetensors.index.json").write_text(json.dumps(index))
+        (tmp_path / "model-00001-of-00001.bin").touch()
+
+        with pytest.raises(ValueError, match="Invalid shard filename.*Expected '.safetensors' extension"):
+            load_torch_model(Mock(), tmp_path)
+
+    def test_safetensors_index_rejects_path_traversal(self, tmp_path):
+        """A shard filename with '..' path traversal must be rejected."""
+        index = {
+            "metadata": {"total_size": 100},
+            "weight_map": {
+                "layer_1": "../malicious.safetensors",
+            },
+        }
+        (tmp_path / "model.safetensors.index.json").write_text(json.dumps(index))
+
+        with pytest.raises(ValueError, match="Invalid shard filename.*without '..' components"):
+            load_torch_model(Mock(), tmp_path)
+
+    def test_safetensors_index_rejects_absolute_path(self, tmp_path):
+        """A shard filename with an absolute path must be rejected."""
+        index = {
+            "metadata": {"total_size": 100},
+            "weight_map": {
+                "layer_1": "/tmp/malicious.safetensors",
+            },
+        }
+        (tmp_path / "model.safetensors.index.json").write_text(json.dumps(index))
+
+        with pytest.raises(ValueError, match="Invalid shard filename.*without '..' components"):
+            load_torch_model(Mock(), tmp_path)