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,38 @@
     return compressed_file_dir
 
 
+def _is_safe_archive_member(extract_dir, member_name):
+    extract_dir = os.path.realpath(extract_dir)
+    candidate_path = os.path.realpath(os.path.join(extract_dir, member_name))
+    return os.path.commonpath([extract_dir, candidate_path]) == extract_dir
+
+
+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}'")
+            try:
+                tar_ref.extractall(extract_dir, filter="data")
+            except TypeError:
+                tar_ref.extractall(extract_dir)
+        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 +128,17 @@
         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)
 

test_convert_tf_models_to_pytorch.py

@@ -0,0 +1,34 @@
+import importlib.util
+import io
+import tarfile
+import zipfile
+from pathlib import Path
+
+import pytest
+
+MODULE_PATH = Path("onnxruntime/python/tools/transformers/convert_tf_models_to_pytorch.py")
+SPEC = importlib.util.spec_from_file_location("convert_tf_models_to_pytorch", MODULE_PATH)
+convert_tf_models_to_pytorch = importlib.util.module_from_spec(SPEC)
+assert SPEC.loader is not None
+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)