diff --git a/scripts/rsi-training/evaluate_adapter.py b/scripts/rsi-training/evaluate_adapter.py index dc9a2d0..9de217b 100644 --- a/scripts/rsi-training/evaluate_adapter.py +++ b/scripts/rsi-training/evaluate_adapter.py @@ -41,7 +41,7 @@ def extract_adapter_archive(archive_path: pathlib.Path, destination: pathlib.Pat observed = set() with tarfile.open(fileobj=io.BytesIO(archive_path.read_bytes()), mode="r:") as archive: for member in archive.getmembers(): - if member.isdir() and member.name.rstrip("/") == expected_prefix.rstrip("/"): + if member.isdir(): continue if not member.isfile() or not member.name.startswith(expected_prefix): raise ValueError("adapter archive contains an unsafe path or non-regular file") diff --git a/scripts/rsi-training/test_adapter_archive.py b/scripts/rsi-training/test_adapter_archive.py index e00e173..73fe7b3 100644 --- a/scripts/rsi-training/test_adapter_archive.py +++ b/scripts/rsi-training/test_adapter_archive.py @@ -41,6 +41,23 @@ def test_extracts_only_the_two_expected_regular_files(self) -> None: self.assertEqual((destination / "adapter_config.json").read_bytes(), b"{}") self.assertEqual((destination / "adapter_model.safetensors").read_bytes(), b"weights") + def test_skips_leading_directory_entries(self) -> None: + parent = tarfile.TarInfo(".headlesscode-rsi-model-output/") + parent.type = tarfile.DIRTYPE + nested = tarfile.TarInfo(".headlesscode-rsi-model-output/adapter/") + nested.type = tarfile.DIRTYPE + with tempfile.TemporaryDirectory() as root: + archive = pathlib.Path(root) / "adapter.tar" + archive.write_bytes(tar_bytes([ + (".headlesscode-rsi-model-output/", b"", parent), + (".headlesscode-rsi-model-output/adapter/", b"", nested), + (".headlesscode-rsi-model-output/adapter/adapter_config.json", b"{}", None), + (".headlesscode-rsi-model-output/adapter/adapter_model.safetensors", b"weights", None), + ])) + destination = pathlib.Path(root) / "out" + MODULE.extract_adapter_archive(archive, destination) + self.assertEqual((destination / "adapter_config.json").read_bytes(), b"{}") + def test_rejects_path_escape_symlinks_and_extra_files(self) -> None: cases = [ [("../outside", b"x", None)],