Arbitrary File Overwrite in download_model_with_test_data in onnx/onnx

@@ -10,7 +10,6 @@
 import re
 import shutil
 import sys
-import tarfile
 import tempfile
 import time
 import unittest
@@ -242,8 +241,7 @@ def download_model(
             )
             urlretrieve(model_test.url, download_file.name)
             print("Done")
-            with tarfile.open(download_file.name) as t:
-                t.extractall(models_dir)
+            onnx.utils._extract_model_safe(download_file.name, models_dir)
         except Exception as e:
             print(f"Failed to prepare data for model {model_test.model_name}: {e}")
             raise

@@ -11,7 +11,6 @@
 import json
 import os
 import sys
-import tarfile
 from io import BytesIO
 from os.path import join
 from typing import IO, Any, Dict, List, cast
@@ -290,35 +289,6 @@ def load(
     return onnx.load(cast(IO[bytes], BytesIO(model_bytes)))
 
 
-def _tar_members_filter(tar: tarfile.TarFile, base: str) -> list[tarfile.TarInfo]:
-    """Check that the content of ``tar`` will be extracted safely
-
-    Args:
-        tar: The tarball file
-        base: The directory where the tarball will be extracted
-
-    Returns:
-        list of tarball members
-    """
-    result = []
-    for member in tar:
-        member_path = os.path.join(base, member.name)
-        abs_base = os.path.abspath(base)
-        abs_member = os.path.abspath(member_path)
-        if not abs_member.startswith(abs_base):
-            raise RuntimeError(
-                f"The tarball member {member_path} in downloading model contains "
-                f"directory traversal sequence which may contain harmful payload."
-            )
-        elif member.issym() or member.islnk():
-            raise RuntimeError(
-                f"The tarball member {member_path} in downloading model contains "
-                f"symbolic links which may contain harmful payload."
-            )
-        result.append(member)
-    return result
-
-
 def download_model_with_test_data(
     model: str,
     repo: str = "onnx/models:main",
@@ -393,23 +363,14 @@ def download_model_with_test_data(
                 "download the model from the model hub."
             )
 
-    with tarfile.open(local_model_with_data_path) as model_with_data_zipped:
-        # FIXME: Avoid index manipulation with magic numbers
-        local_model_with_data_dir_path = local_model_with_data_path[
-            0 : len(local_model_with_data_path) - 7
-        ]
-        # Mitigate tarball directory traversal risks
-        if hasattr(tarfile, "data_filter"):
-            model_with_data_zipped.extractall(
-                path=local_model_with_data_dir_path, filter="data"
-            )
-        else:
-            model_with_data_zipped.extractall(
-                path=local_model_with_data_dir_path,
-                members=_tar_members_filter(
-                    model_with_data_zipped, local_model_with_data_dir_path
-                ),
-            )
+    # FIXME: Avoid index manipulation with magic numbers,
+    # remove ".tar.gz"
+    local_model_with_data_dir_path = local_model_with_data_path[
+        0 : len(local_model_with_data_path) - 7
+    ]
+    onnx.utils._extract_model_safe(
+        local_model_with_data_path, local_model_with_data_dir_path
+    )
     model_with_data_path = (
         local_model_with_data_dir_path
         + "/"

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import os
+import tarfile
 
 import onnx.checker
 import onnx.helper
@@ -232,3 +233,65 @@ def extract_model(
     onnx.save(extracted, output_path)
     if check_model:
         onnx.checker.check_model(output_path)
+
+
+def _tar_members_filter(
+    tar: tarfile.TarFile, base: str | os.PathLike
+) -> list[tarfile.TarInfo]:
+    """Check that the content of ``tar`` will be extracted safely
+
+    Args:
+        tar: The tarball file
+        base: The directory where the tarball will be extracted
+
+    Returns:
+        list of tarball members
+    """
+    result = []
+    for member in tar:
+        member_path = os.path.join(base, member.name)
+        abs_base = os.path.abspath(base)
+        abs_member = os.path.abspath(member_path)
+        if not abs_member.startswith(abs_base):
+            raise RuntimeError(
+                f"The tarball member {member_path} in downloading model contains "
+                f"directory traversal sequence which may contain harmful payload."
+            )
+        elif member.issym() or member.islnk():
+            raise RuntimeError(
+                f"The tarball member {member_path} in downloading model contains "
+                f"symbolic links which may contain harmful payload."
+            )
+        result.append(member)
+    return result
+
+
+def _extract_model_safe(
+    model_tar_path: str | os.PathLike, local_model_with_data_dir_path: str | os.PathLike
+) -> None:
+    """Safely extracts a tar file to a specified directory.
+
+    This function ensures that the extraction process mitigates against
+    directory traversal vulnerabilities by validating or sanitizing paths
+    within the tar file. It also provides compatibility for different versions
+    of the tarfile module by checking for the availability of certain attributes
+    or methods before invoking them.
+
+    Args:
+        model_tar_path: The path to the tar file to be extracted.
+        local_model_with_data_dir_path: The directory path where the tar file
+      contents will be extracted to.
+    """
+    with tarfile.open(model_tar_path) as model_with_data_zipped:
+        # Mitigate tarball directory traversal risks
+        if hasattr(tarfile, "data_filter"):
+            model_with_data_zipped.extractall(
+                path=local_model_with_data_dir_path, filter="data"
+            )
+        else:
+            model_with_data_zipped.extractall(
+                path=local_model_with_data_dir_path,
+                members=_tar_members_filter(
+                    model_with_data_zipped, local_model_with_data_dir_path
+                ),
+            )

@@ -290,6 +290,35 @@ def load(
     return onnx.load(cast(IO[bytes], BytesIO(model_bytes)))
 
 
+def _tar_members_filter(tar: tarfile.TarFile, base: str) -> list[tarfile.TarInfo]:
+    """Check that the content of ``tar`` will be extracted safely
+
+    Args:
+        tar: The tarball file
+        base: The directory where the tarball will be extracted
+
+    Returns:
+        list of tarball members
+    """
+    result = []
+    for member in tar:
+        member_path = os.path.join(base, member.name)
+        abs_base = os.path.abspath(base)
+        abs_member = os.path.abspath(member_path)
+        if not abs_member.startswith(abs_base):
+            raise RuntimeError(
+                f"The tarball member {member_path} in downloading model contains "
+                f"directory traversal sequence which may contain harmful payload."
+            )
+        elif member.issym() or member.islnk():
+            raise RuntimeError(
+                f"The tarball member {member_path} in downloading model contains "
+                f"symbolic links which may contain harmful payload."
+            )
+        result.append(member)
+    return result
+
+
 def download_model_with_test_data(
     model: str,
     repo: str = "onnx/models:main",
@@ -298,6 +327,7 @@ def download_model_with_test_data(
     silent: bool = False,
 ) -> str | None:
     """Downloads a model along with test data by name from the onnx model hub and returns the directory to which the files have been extracted.
+    Users are responsible for making sure the model comes from a trusted source, and the data is safe to be extracted.
 
     Args:
         model: The name of the onnx model in the manifest. This field is
@@ -368,7 +398,18 @@ def download_model_with_test_data(
         local_model_with_data_dir_path = local_model_with_data_path[
             0 : len(local_model_with_data_path) - 7
         ]
-        model_with_data_zipped.extractall(local_model_with_data_dir_path)
+        # Mitigate tarball directory traversal risks
+        if hasattr(tarfile, "data_filter"):
+            model_with_data_zipped.extractall(
+                path=local_model_with_data_dir_path, filter="data"
+            )
+        else:
+            model_with_data_zipped.extractall(
+                path=local_model_with_data_dir_path,
+                members=_tar_members_filter(
+                    model_with_data_zipped, local_model_with_data_dir_path
+                ),
+            )
     model_with_data_path = (
         local_model_with_data_dir_path
         + "/"