diff --git a/src/mcore_bridge/utils/safetensors.py b/src/mcore_bridge/utils/safetensors.py index 0ad3ea3..1f4b561 100644 --- a/src/mcore_bridge/utils/safetensors.py +++ b/src/mcore_bridge/utils/safetensors.py @@ -30,6 +30,26 @@ def load_slice(self, slices): return self.loader()[slices] +def checkpoint_path(root: str, name: str) -> str: + """Join a checkpoint-relative name and reject paths that leave ``root``. + + ``os.path.join`` drops the root when ``name`` is absolute, and a weight-map + entry such as ``../../other.safetensors`` would otherwise open a file + outside the model directory. + """ + if os.path.isabs(name): + raise ValueError(f'Checkpoint file {name!r} is outside {root}') + base = os.path.realpath(root) + path = os.path.realpath(os.path.join(base, name)) + try: + inside = os.path.commonpath([base, path]) == base + except ValueError: + inside = False + if not inside: + raise ValueError(f'Checkpoint file {name!r} is outside {root}') + return path + + class SafetensorLazyLoader: def __init__(self, hf_model_dir: str, peft_format: bool = False): @@ -42,7 +62,7 @@ def __init__(self, hf_model_dir: str, peft_format: bool = False): def _open_file(self, filename: str): """Open a safetensors file if not already open.""" if filename not in self._file_handles: - file_path = os.path.join(self.hf_model_dir, filename) + file_path = checkpoint_path(self.hf_model_dir, filename) self._file_handles[filename] = safe_open(file_path, framework='pt') return self._file_handles[filename] diff --git a/tests/test_checkpoint_path.py b/tests/test_checkpoint_path.py new file mode 100644 index 0000000..bd0770e --- /dev/null +++ b/tests/test_checkpoint_path.py @@ -0,0 +1,35 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Weight-map filenames cannot leave the checkpoint directory.""" +import ast +import os +from pathlib import Path + + +def _join(): + path = Path(__file__).resolve().parents[1] / 'src/mcore_bridge/utils/safetensors.py' + tree = ast.parse(path.read_text()) + fn = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == 'checkpoint_path') + module = ast.Module(body=[fn], type_ignores=[]) + ast.fix_missing_locations(module) + namespace = {'os': os} + exec(compile(module, str(path), 'exec'), namespace) + return namespace['checkpoint_path'] + + +def test_relative_shard_stays_inside(tmp_path): + join = _join() + shard = tmp_path / 'model-00001.safetensors' + shard.write_bytes(b'x') + assert join(str(tmp_path), 'model-00001.safetensors') == os.path.realpath(shard) + + +def test_parent_and_absolute_names_are_rejected(tmp_path): + join = _join() + outside = tmp_path.parent / 'outside.safetensors' + outside.write_bytes(b'x') + for name in ('../outside.safetensors', str(outside)): + try: + join(str(tmp_path), name) + except ValueError: + continue + raise AssertionError(name)