fix(modeling): reject sharded-checkpoint weight_map entries that escape the checkpoint folder (#4067)

[Anai-Guo]

What this fixes

Closes #4067 (path-traversal half of the report).

load_checkpoint_in_model builds the list of shard files to load directly from the untrusted weight_map of a sharded-checkpoint index (*.index.json):

checkpoint_files = sorted(list(set(index.values())))
checkpoint_files = [os.path.join(checkpoint_folder, f) for f in checkpoint_files]

The values come straight from the model author. Because os.path.join resolves .. segments and discards the prefix entirely when given an absolute path, a malicious or corrupted index can point a "shard" anywhere on the filesystem:

{"weight_map": {"w": "../../../../etc/passwd"}}      // escapes via traversal
{"weight_map": {"w": "/etc/shadow"}}                  // absolute path wins in os.path.join

The file is then opened and parsed as a checkpoint shard, which is a path-traversal / arbitrary-file-read primitive (and an easy DoS vector when pointed at a large or special file).

The fix

After resolving each entry against checkpoint_folder, verify it still lives inside that folder and raise ValueError otherwise. Legitimate sharded checkpoints (HF save_pretrained output) always reference plain sibling filenames, so they are unaffected.

checkpoint_folder_abs = os.path.abspath(checkpoint_folder)
for shard_file in checkpoint_files:
    resolved = os.path.abspath(os.path.join(checkpoint_folder, shard_file))
    if resolved != checkpoint_folder_abs and not resolved.startswith(checkpoint_folder_abs + os.sep):
        raise ValueError(...)

The containment check (startswith(folder + os.sep)) also rejects absolute paths and, on Windows, paths on a different drive.

Scope note

This targets the path-traversal vector specifically. The "malicious-model DoS" angle in the title (unbounded shard size / count) is broader and policy-dependent — happy to follow up separately if maintainers want bounds there too.

🤖 Generated with Claude Code