Add gluonts.util.safe_extract (#2606)

Co-authored-by: Jasper <[email protected]>
Co-authored-by: Lorenzo Stella <[email protected]>

src/gluonts/dataset/repository/_gp_copula_2019.py

@@ -32,6 +32,7 @@
 from gluonts.dataset.common import MetaData, TrainDatasets
 from gluonts.dataset.field_names import FieldName
 from gluonts.dataset.repository._util import metadata
+from gluonts.util import safe_extractall
 
 
 class GPCopulaDataset(NamedTuple):
@@ -140,7 +141,7 @@ def download_dataset(dataset_path: Path, ds_info: GPCopulaDataset):
     request.urlretrieve(ds_info.url, dataset_path / f"{ds_info.name}.tar.gz")
 
     with tarfile.open(dataset_path / f"{ds_info.name}.tar.gz") as tar:
-        tar.extractall(path=dataset_path)
+        safe_extractall(tar, path=dataset_path)
 
 
 def get_data(dataset_path: Path, ds_info: GPCopulaDataset):

src/gluonts/nursery/sagemaker_sdk/estimator.py

@@ -31,6 +31,7 @@
 from gluonts.dataset.repository import datasets
 from gluonts.model.estimator import Estimator
 from gluonts.model.predictor import Predictor
+from gluonts.util import safe_extractall
 
 from .defaults import (
     ENTRY_POINTS_FOLDER,
@@ -503,7 +504,7 @@ def _retrieve_model(self, locations):
         with self._s3fs.open(locations.model_archive, "rb") as stream:
             with tarfile.open(mode="r:gz", fileobj=stream) as archive:
                 with TemporaryDirectory() as temp_dir:
-                    archive.extractall(temp_dir)
+                    safe_extractall(archive, temp_dir)
                     predictor = Predictor.deserialize(Path(temp_dir))
 
         return predictor

src/gluonts/nursery/tsbench/src/cli/evaluations/download.py

@@ -19,6 +19,7 @@
 from typing import Any, cast, Dict, List, Optional
 import botocore
 import click
+from gluonts.util import safe_extract
 from tqdm.auto import tqdm
 from tqdm.contrib.concurrent import process_map
 from tsbench.analysis.utils import run_parallel
@@ -104,7 +105,7 @@ def _download_public_evaluations(
         file = Path(tmp) / "metrics.tar.gz"
         client.download_file(public_bucket, "metrics.tar.gz", str(file))
         with tarfile.open(file, mode="r:gz") as tar:
-            tar.extractall(evaluations_path)
+            safe_extractall(tar, evaluations_path)
 
     # Then, optionally download the forecasts
     if include_forecasts:

src/gluonts/shell/env.py

@@ -21,6 +21,7 @@
 
 from gluonts.dataset.common import Dataset, FileDataset, MetaData
 from gluonts.model import Predictor
+from gluonts.util import safe_extractall
 
 from . import sagemaker
 
@@ -41,7 +42,7 @@ def _load(self) -> Dict[str, Union[Dataset, Predictor]]:
         if "model" in self.channels:
             path = self.channels.pop("model")
             with tarfile.open(path / "model.tar.gz") as targz:
-                targz.extractall(path=path)
+                safe_extractall(targz, path)
             model = Predictor.deserialize(path)
 
         file_dataset = partial(FileDataset, freq=self.hyperparameters["freq"])

src/gluonts/util.py

@@ -12,7 +12,9 @@
 # permissions and limitations under the License.
 
 import copy
+import tarfile
 from functools import lru_cache
+from pathlib import Path
 from typing import TYPE_CHECKING, TypeVar
 
 T = TypeVar("T")
@@ -66,3 +68,36 @@ def my_property(self):
         introduced in Python 3.8.
         """
         return property(lru_cache(1)(method))
+
+
+def will_extractall_into(tar: tarfile.TarFile, path: Path) -> None:
+    """
+    Check that the content of ``tar`` will be extracted within ``path``
+    upon calling ``extractall``.
+
+    Raise a ``PermissionError`` if not.
+    """
+    path = Path(path).resolve()
+
+    for member in tar.getmembers():
+        member_path = (path / member.name).resolve()
+
+        try:
+            member_path.relative_to(path)
+        except ValueError:
+            raise PermissionError(f"'{member.name}' extracts out of target.")
+
+
+def safe_extractall(
+    tar: tarfile.TarFile,
+    path: Path = Path("."),
+    members=None,
+    *,
+    numeric_owner=False,
+):
+    """
+    Safe wrapper around ``TarFile.extractall`` that checks all destination
+    files to be strictly within the given ``path``.
+    """
+    will_extractall_into(tar, path)
+    tar.extractall(path, members, numeric_owner=numeric_owner)

test/test_util.py

@@ -11,7 +11,14 @@
 # express or implied. See the License for the specific language governing
 # permissions and limitations under the License.
 
-from gluonts.util import copy_with
+import tempfile
+import tarfile
+from pathlib import Path
+from typing import Optional
+
+import pytest
+
+from gluonts.util import copy_with, will_extractall_into
 
 
 def test_copy_with():
@@ -24,3 +31,33 @@ def __init__(self, value):
 
     assert a.value == 42
     assert b.value == 99
+
+
+@pytest.mark.parametrize(
+    "arcname, expect_failure",
+    [
+        (None, False),
+        ("./file.txt", False),
+        ("/a/../file.txt", False),
+        ("/a/../../file.txt", True),
+        ("../file.txt", True),
+    ],
+)
+def test_will_extractall_into(arcname: Optional[str], expect_failure: bool):
+    with tempfile.TemporaryDirectory() as tempdir:
+        file_path = Path(tempdir) / "a" / "file.txt"
+        file_path.parent.mkdir(parents=True)
+        file_path.touch()
+
+        with tarfile.open(Path(tempdir) / "archive.tar.gz", "w:gz") as tar:
+            tar.add(file_path, arcname=arcname)
+
+        if expect_failure:
+            with pytest.raises(PermissionError):
+                with tarfile.open(
+                    Path(tempdir) / "archive.tar.gz", "r:gz"
+                ) as tar:
+                    will_extractall_into(tar, Path(tempdir) / "b")
+        else:
+            with tarfile.open(Path(tempdir) / "archive.tar.gz", "r:gz") as tar:
+                will_extractall_into(tar, Path(tempdir) / "b")