fix(transformers): validate archive extraction paths

onnxruntime/python/tools/transformers/convert_tf_models_to_pytorch.py

@@ -5,6 +5,8 @@
 
 import glob
 import os
+import tarfile
+import zipfile
 
 import requests
 
@@ -63,6 +65,54 @@ def download_compressed_file(tf_ckpt_url, ckpt_dir):
     return compressed_file_dir
 
 
+def _is_safe_archive_member(extract_dir, member_name):
+    # Normalize separators so that backslash-based traversal entries are detected
+    # on POSIX as well as Windows (zip members may use either separator).
+    normalized_parts = [part for part in member_name.replace("\\", "/").split("/") if part not in ("", ".")]
+    extract_dir = os.path.realpath(extract_dir)
+    candidate_path = os.path.realpath(os.path.join(extract_dir, *normalized_parts))
+    try:
+        return os.path.commonpath([extract_dir, candidate_path]) == extract_dir
+    except ValueError:
+        # os.path.commonpath raises ValueError for mixed drives or absolute/relative
+        # mixes (e.g. on Windows). Treat any such case as unsafe.
+        return False
+
+
+def safe_extract_archive(archive_path, extract_dir):
+    archive_path = os.path.realpath(archive_path)
+    extract_dir = os.path.realpath(extract_dir)
+
+    if tarfile.is_tarfile(archive_path):
+        with tarfile.open(archive_path, "r:*") as tar_ref:
+            for member in tar_ref.getmembers():
+                if not _is_safe_archive_member(extract_dir, member.name):
+                    raise ValueError(f"Archive member '{member.name}' resolves outside '{extract_dir}'")
+                if member.issym() or member.islnk():
+                    raise ValueError(f"Archive member '{member.name}' is a link, which is not allowed")
+            try:
+                tar_ref.extractall(extract_dir, filter="data")
+            except TypeError as exc:
+                # The "data" extraction filter (Python 3.9+/backports) rejects symlinks,
+                # hardlinks, device files and absolute paths. Without it, extraction is
+                # not safe, so refuse rather than falling back to an unfiltered extract.
+                raise RuntimeError(
+                    "Safe archive extraction requires the tarfile 'data' filter, "
+                    "which is unavailable in this Python runtime."
+                ) from exc
+        return
+
+    if zipfile.is_zipfile(archive_path):
+        with zipfile.ZipFile(archive_path, "r") as zip_ref:
+            for member_name in zip_ref.namelist():
+                if not _is_safe_archive_member(extract_dir, member_name):
+                    raise ValueError(f"Archive member '{member_name}' resolves outside '{extract_dir}'")
+            zip_ref.extractall(extract_dir)
+        return
+
+    raise ValueError(f"Unsupported archive format: {archive_path}")
+
+
 def get_ckpt_prefix_path(ckpt_dir):
     # get prefix
     sub_folder_dir = None
@@ -94,23 +144,17 @@ def download_tf_checkpoint(model_name, tf_models_dir="tf_models"):
         zip_dir = download_compressed_file(tf_ckpt_url, ckpt_dir)
 
         # unzip file
-        import zipfile  # noqa: PLC0415
-
-        with zipfile.ZipFile(zip_dir, "r") as zip_ref:
-            zip_ref.extractall(ckpt_dir)
-            os.remove(zip_dir)
+        safe_extract_archive(zip_dir, ckpt_dir)
+        os.remove(zip_dir)
 
         return get_ckpt_prefix_path(ckpt_dir)
 
     elif re.search(".tar.gz$", tf_ckpt_url) is not None:
         tar_dir = download_compressed_file(tf_ckpt_url, ckpt_dir)
 
         # untar file
-        import tarfile  # noqa: PLC0415
-
-        with tarfile.open(tar_dir, "r") as tar_ref:
-            tar_ref.extractall(ckpt_dir)
-            os.remove(tar_dir)
+        safe_extract_archive(tar_dir, ckpt_dir)
+        os.remove(tar_dir)
 
         return get_ckpt_prefix_path(ckpt_dir)
 

onnxruntime/test/python/transformers/test_convert_tf_models_to_pytorch.py

@@ -0,0 +1,63 @@
+import importlib.util
+import io
+import tarfile
+import zipfile
+from pathlib import Path
+
+import pytest
+
+# Resolve the module path relative to this file so the test works regardless of CWD.
+MODULE_PATH = (
+    Path(__file__).resolve().parents[4]
+    / "onnxruntime"
+    / "python"
+    / "tools"
+    / "transformers"
+    / "convert_tf_models_to_pytorch.py"
+)
+SPEC = importlib.util.spec_from_file_location("convert_tf_models_to_pytorch", MODULE_PATH)
+assert SPEC is not None and SPEC.loader is not None
+convert_tf_models_to_pytorch = importlib.util.module_from_spec(SPEC)
+SPEC.loader.exec_module(convert_tf_models_to_pytorch)
+
+
+def test_safe_extract_archive_rejects_tar_path_traversal(tmp_path):
+    archive_path = tmp_path / "malicious.tar.gz"
+    with tarfile.open(archive_path, "w:gz") as tar_ref:
+        data = b"malicious"
+        member = tarfile.TarInfo(name="../escape.txt")
+        member.size = len(data)
+        tar_ref.addfile(member, io.BytesIO(data))
+
+    with pytest.raises(ValueError, match="outside"):
+        convert_tf_models_to_pytorch.safe_extract_archive(archive_path, tmp_path)
+
+
+def test_safe_extract_archive_rejects_zip_path_traversal(tmp_path):
+    archive_path = tmp_path / "malicious.zip"
+    with zipfile.ZipFile(archive_path, "w") as zip_ref:
+        zip_ref.writestr("../escape.txt", "malicious")
+
+    with pytest.raises(ValueError, match="outside"):
+        convert_tf_models_to_pytorch.safe_extract_archive(archive_path, tmp_path)
+
+
+def test_safe_extract_archive_rejects_zip_backslash_traversal(tmp_path):
+    archive_path = tmp_path / "malicious_backslash.zip"
+    with zipfile.ZipFile(archive_path, "w") as zip_ref:
+        zip_ref.writestr("..\\escape.txt", "malicious")
+
+    with pytest.raises(ValueError, match="outside"):
+        convert_tf_models_to_pytorch.safe_extract_archive(archive_path, tmp_path)
+
+
+def test_safe_extract_archive_rejects_tar_symlink(tmp_path):
+    archive_path = tmp_path / "malicious_link.tar.gz"
+    with tarfile.open(archive_path, "w:gz") as tar_ref:
+        member = tarfile.TarInfo(name="link")
+        member.type = tarfile.SYMTYPE
+        member.linkname = "/etc/passwd"
+        tar_ref.addfile(member)
+
+    with pytest.raises(ValueError, match="link"):
+        convert_tf_models_to_pytorch.safe_extract_archive(archive_path, tmp_path)