CVE-2007-4559 Patch

python/arachne_runtime/__init__.py

@@ -118,7 +118,29 @@ def init(
                     env_file = m.name
                 else:
                     model_file = m.name
-            tar.extractall(".")
+
+            import os
+
+            def is_within_directory(directory, target):
+
+                abs_directory = os.path.abspath(directory)
+                abs_target = os.path.abspath(target)
+
+                prefix = os.path.commonprefix([abs_directory, abs_target])
+
+                return prefix == abs_directory
+
+            def safe_extract(tar, path=".", members=None, *, numeric_owner=False):
+
+                for member in tar.getmembers():
+                    member_path = os.path.join(path, member.name)
+                    if not is_within_directory(path, member_path):
+                        raise Exception("Attempted Path Traversal in Tar File")
+
+                tar.extractall(path, members, numeric_owner=numeric_owner) 
+
+
+            safe_extract(tar, ".")
 
     assert model_file is not None
 

python/arachne_runtime/module/tvm.py

@@ -18,7 +18,26 @@ def _open_module_file(file: str) -> Tuple[Optional[str], Optional[bytearray], An
         tvm = importlib.import_module("tvm")
     with tempfile.TemporaryDirectory() as tmp_dir:
         with tarfile.open(file) as t:
-            t.extractall(tmp_dir)
+            def is_within_directory(directory, target):
+
+                abs_directory = os.path.abspath(directory)
+                abs_target = os.path.abspath(target)
+
+                prefix = os.path.commonprefix([abs_directory, abs_target])
+
+                return prefix == abs_directory
+
+            def safe_extract(tar, path=".", members=None, *, numeric_owner=False):
+
+                for member in tar.getmembers():
+                    member_path = os.path.join(path, member.name)
+                    if not is_within_directory(path, member_path):
+                        raise Exception("Attempted Path Traversal in Tar File")
+
+                tar.extractall(path, members, numeric_owner=numeric_owner) 
+
+
+            safe_extract(t, tmp_dir)
         graph = None
         params = None
         graph_path = os.path.join(tmp_dir, "mod.json")

python/arachne_runtime/rpc/servicer/servicer.py

@@ -38,7 +38,26 @@ def Init(self, request, context):
             path = args["model_dir"]
             tmpdir = tempfile.mkdtemp()
             with tarfile.open(path, "r") as f:
-                f.extractall(tmpdir)
+                def is_within_directory(directory, target):
+
+                    abs_directory = os.path.abspath(directory)
+                    abs_target = os.path.abspath(target)
+
+                    prefix = os.path.commonprefix([abs_directory, abs_target])
+
+                    return prefix == abs_directory
+
+                def safe_extract(tar, path=".", members=None, *, numeric_owner=False):
+
+                    for member in tar.getmembers():
+                        member_path = os.path.join(path, member.name)
+                        if not is_within_directory(path, member_path):
+                            raise Exception("Attempted Path Traversal in Tar File")
+
+                    tar.extractall(path, members, numeric_owner=numeric_owner) 
+
+
+                safe_extract(f, tmpdir)
             args["model_dir"] = tmpdir
 
         if path is None:

tests/rpc/test_run.py

@@ -67,7 +67,26 @@ def test_tvm_runtime_rpc2(rpc_port=5052):
                     env_file = m.name
                 else:
                     model_file = m.name
-            tar.extractall(".")
+            def is_within_directory(directory, target):
+
+                abs_directory = os.path.abspath(directory)
+                abs_target = os.path.abspath(target)
+
+                prefix = os.path.commonprefix([abs_directory, abs_target])
+
+                return prefix == abs_directory
+
+            def safe_extract(tar, path=".", members=None, *, numeric_owner=False):
+
+                for member in tar.getmembers():
+                    member_path = os.path.join(path, member.name)
+                    if not is_within_directory(path, member_path):
+                        raise Exception("Attempted Path Traversal in Tar File")
+
+                tar.extractall(path, members, numeric_owner=numeric_owner) 
+
+
+            safe_extract(tar, ".")
         assert model_file is not None
         assert env_file is not None
         # local run

tests/test_tvm_graph_executor.py

@@ -21,7 +21,26 @@ def _test_tvm_runtime(device):
 
         os.chdir(tmp_dir)
         with tarfile.open(tvm_package_path, "r:gz") as tar:
-            tar.extractall(".")
+            def is_within_directory(directory, target):
+
+                abs_directory = os.path.abspath(directory)
+                abs_target = os.path.abspath(target)
+
+                prefix = os.path.commonprefix([abs_directory, abs_target])
+
+                return prefix == abs_directory
+
+            def safe_extract(tar, path=".", members=None, *, numeric_owner=False):
+
+                for member in tar.getmembers():
+                    member_path = os.path.join(path, member.name)
+                    if not is_within_directory(path, member_path):
+                        raise Exception("Attempted Path Traversal in Tar File")
+
+                tar.extractall(path, members, numeric_owner=numeric_owner) 
+
+
+            safe_extract(tar, ".")
         tvm_model_path = tmp_dir + "/tvm_package_0.tar"
 
         input_data = np.array(np.random.random_sample([1, 224, 224, 3]), dtype=np.float32)  # type: ignore