Security and robustness fixes for the dataset/download path plus a

CHANGELOG.md

@@ -10,10 +10,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+* `l2l.data.utils.safe_extract` helper that guards archive extraction against path traversal (Zip-Slip). (@[Garthigan](https://github.com/Garthigan))
+
 ### Changed
 
+* Remove unused `distutils` import from `setup.py` so installation works on Python 3.12+ (distutils was removed from the stdlib in PEP 632). (@[Garthigan](https://github.com/Garthigan))
+* Add request timeouts to all dataset/file downloads. (@[Garthigan](https://github.com/Garthigan))
+
 ### Fixed
 
+* Guard all dataset archive extraction against path traversal (Zip-Slip) by routing `tarfile`/`zipfile` extraction through `safe_extract`. (@[Garthigan](https://github.com/Garthigan))
+* Load downloaded pretrained backbones with `torch.load(..., weights_only=True)` to avoid arbitrary code execution from tampered checkpoints. (@[Garthigan](https://github.com/Garthigan))
+* Close `zipfile.ZipFile` handles in `FC100` (and `tiered-ImageNet` tar handle) via context managers to avoid file-descriptor leaks. (@[Garthigan](https://github.com/Garthigan))
+
 
 ## v0.2.1
 

learn2learn/data/utils.py

@@ -1,16 +1,50 @@
 #!/usr/bin/env python3
 
+import os
+import tarfile
+import zipfile
+
 import torch
 import requests
 import tqdm
 
 CHUNK_SIZE = 1 * 1024 * 1024
+DOWNLOAD_TIMEOUT = 60
+
+
+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(archive, destination='.'):
+    """
+    Extracts an open TarFile or ZipFile while guarding against path traversal
+    (a.k.a. Zip-Slip), where a malicious archive writes files outside of
+    `destination` via entries such as `../../etc/passwd`.
+    """
+    if isinstance(archive, tarfile.TarFile):
+        names = [member.name for member in archive.getmembers()]
+    elif isinstance(archive, zipfile.ZipFile):
+        names = archive.namelist()
+    else:
+        raise TypeError(
+            'safe_extract expects a TarFile or ZipFile, got '
+            + type(archive).__name__
+        )
+    for name in names:
+        member_path = os.path.join(destination, name)
+        if not _is_within_directory(destination, member_path):
+            raise Exception('Attempted path traversal in archive: ' + name)
+    archive.extractall(destination)
 
 
 def download_file(source, destination, size=None):
     if size is None:
         size = 0
-    req = requests.get(source, stream=True)
+    req = requests.get(source, stream=True, timeout=DOWNLOAD_TIMEOUT)
     with open(destination, 'wb') as archive:
         for chunk in tqdm.tqdm(
             req.iter_content(chunk_size=CHUNK_SIZE),
@@ -24,11 +58,15 @@ def download_file(source, destination, size=None):
 def download_file_from_google_drive(id, destination):
     URL = "https://docs.google.com/uc?export=download"
     session = requests.Session()
-    response = session.get(URL, params={'id': id}, stream=True)
+    response = session.get(
+        URL, params={'id': id}, stream=True, timeout=DOWNLOAD_TIMEOUT,
+    )
     token = get_confirm_token(response)
     if token:
         params = {'id': id, 'confirm': token}
-        response = session.get(URL, params=params, stream=True)
+        response = session.get(
+            URL, params=params, stream=True, timeout=DOWNLOAD_TIMEOUT,
+        )
     save_response_content(response, destination)
 
 

learn2learn/text/datasets/news_classification.py

@@ -13,6 +13,8 @@
 
 from torch.utils.data import Dataset
 
+from learn2learn.data.utils import safe_extract
+
 
 class NewsClassification(Dataset):
     """
@@ -56,9 +58,9 @@ def __init__(self, root, train=True, transform=None, download=False):
             if train:
                 download_file_url = 'https://www.dropbox.com/s/o71z7fq7mydbznc/train_sample.csv.zip?dl=1'
 
-            r = requests.get(download_file_url)
-            z = zipfile.ZipFile(io.BytesIO(r.content))
-            z.extractall(path=root)
+            r = requests.get(download_file_url, timeout=60)
+            with zipfile.ZipFile(io.BytesIO(r.content)) as z:
+                safe_extract(z, root)
 
         if root:
 

learn2learn/vision/datasets/cifarfs.py

@@ -11,6 +11,7 @@
 from learn2learn.data.utils import (
     download_file_from_google_drive,
     download_file,
+    safe_extract,
 )
 
 
@@ -93,13 +94,13 @@ def _download(self):
                 destination=zip_file,
             )
             with zipfile.ZipFile(zip_file, 'r') as zfile:
-                zfile.extractall(self.raw_path)
+                safe_extract(zfile, self.raw_path)
             os.remove(zip_file)
         except Exception:
             download_file_from_google_drive('1pTsCCMDj45kzFYgrnO67BWVbKs48Q3NI',
                                             zip_file)
             with zipfile.ZipFile(zip_file, 'r') as zfile:
-                zfile.extractall(self.raw_path)
+                safe_extract(zfile, self.raw_path)
             os.remove(zip_file)
 
     def _process_zip(self):

learn2learn/vision/datasets/cu_birds200.py

@@ -8,6 +8,7 @@
 from learn2learn.data.utils import (
     download_file_from_google_drive,
     download_file,
+    safe_extract,
 )
 
 DATA_DIR = 'cubirds200'
@@ -360,13 +361,13 @@ def download(self):
         try:
             download_file(ZENODO_URL, tar_path)
             tar_file = tarfile.open(tar_path)
-            tar_file.extractall(data_path)
+            safe_extract(tar_file, data_path)
             tar_file.close()
             os.remove(tar_path)
         except Exception:
             download_file_from_google_drive(ARCHIVE_ID, tar_path)
             tar_file = tarfile.open(tar_path)
-            tar_file.extractall(data_path)
+            safe_extract(tar_file, data_path)
             tar_file.close()
             os.remove(tar_path)
 

learn2learn/vision/datasets/describable_textures.py

@@ -6,7 +6,7 @@
 from PIL import Image
 from torch.utils.data import Dataset
 
-from learn2learn.data.utils import download_file
+from learn2learn.data.utils import download_file, safe_extract
 from torchvision.datasets.folder import default_loader
 
 DATA_DIR = 'describable_textures'
@@ -144,7 +144,7 @@ def download(self):
         print('Downloading Describable Textures dataset (600Mb)')
         download_file(ARCHIVE_URL, tar_path)
         tar_file = tarfile.open(tar_path)
-        tar_file.extractall(data_path)
+        safe_extract(tar_file, data_path)
         tar_file.close()
         os.remove(tar_path)
 

learn2learn/vision/datasets/fc100.py

@@ -11,6 +11,7 @@
 from learn2learn.data.utils import (
     download_file_from_google_drive,
     download_file,
+    safe_extract,
 )
 
 
@@ -90,20 +91,20 @@ def download(self):
         print('Downloading FC100. (160Mb)')
         try:
             download_file(FC100.ZENODO_LINK, archive_path)
-            archive_file = zipfile.ZipFile(archive_path)
-            archive_file.extractall(self.root)
+            with zipfile.ZipFile(archive_path) as archive_file:
+                safe_extract(archive_file, self.root)
             os.remove(archive_path)
         except Exception:
             try:  # Download from Google Drive first
                 download_file_from_google_drive(FC100.GOOGLE_DRIVE_FILE_ID,
                                                 archive_path)
-                archive_file = zipfile.ZipFile(archive_path)
-                archive_file.extractall(self.root)
+                with zipfile.ZipFile(archive_path) as archive_file:
+                    safe_extract(archive_file, self.root)
                 os.remove(archive_path)
             except zipfile.BadZipFile:
                 download_file(FC100.DROPBOX_LINK, archive_path)
-                archive_file = zipfile.ZipFile(archive_path)
-                archive_file.extractall(self.root)
+                with zipfile.ZipFile(archive_path) as archive_file:
+                    safe_extract(archive_file, self.root)
                 os.remove(archive_path)
 
     def __getitem__(self, idx):

learn2learn/vision/datasets/fgvc_fungi.py

@@ -11,7 +11,7 @@
 from PIL import Image
 from torch.utils.data import Dataset
 
-from learn2learn.data.utils import download_file
+from learn2learn.data.utils import download_file, safe_extract
 
 DATA_DIR = 'fgvc_fungi'
 DATA_URL = 'https://labs.gbif.org/fgvcx/2018/fungi_train_val.tgz'
@@ -1494,13 +1494,13 @@ def download(self):
 
         # Extract data
         tar_file = tarfile.open(data_tar_path)
-        tar_file.extractall(data_path)
+        safe_extract(tar_file, data_path)
         tar_file.close()
         os.remove(data_tar_path)
 
         # Extract annotations
         tar_file = tarfile.open(annotations_tar_path)
-        tar_file.extractall(data_path)
+        safe_extract(tar_file, data_path)
         tar_file.close()
         os.remove(annotations_tar_path)
 

learn2learn/vision/datasets/tiered_imagenet.py

@@ -14,6 +14,7 @@
 from learn2learn.data.utils import (
     download_file_from_google_drive,
     download_file,
+    safe_extract,
 )
 
 
@@ -108,8 +109,8 @@ def download(self, file_id, destination):
         except Exception:
             archive_path = os.path.join(destination, 'tiered_imagenet.tar')
             download_file_from_google_drive(file_id, archive_path)
-            archive_file = tarfile.open(archive_path)
-            archive_file.extractall(destination)
+            with tarfile.open(archive_path) as archive_file:
+                safe_extract(archive_file, destination)
             os.remove(archive_path)
 
     def __getitem__(self, idx):

learn2learn/vision/datasets/vgg_flowers.py

@@ -8,11 +8,11 @@
 from PIL import Image
 from torch.utils.data import Dataset
 
-from learn2learn.data.utils import download_file
+from learn2learn.data.utils import download_file, safe_extract
 
 DATA_DIR = 'vgg_flower102'
-IMAGES_URL = 'http://www.robots.ox.ac.uk/~vgg/data/flowers/102/102flowers.tgz'
-LABELS_URL = 'http://www.robots.ox.ac.uk/~vgg/data/flowers/102/imagelabels.mat'
+IMAGES_URL = 'https://www.robots.ox.ac.uk/~vgg/data/flowers/102/102flowers.tgz'
+LABELS_URL = 'https://www.robots.ox.ac.uk/~vgg/data/flowers/102/imagelabels.mat'
 IMAGES_DIR = 'jpg'
 LABELS_PATH = 'imagelabels.mat'
 
@@ -92,12 +92,12 @@ def download(self):
         print('Downloading VGG Flower102 dataset (330Mb)')
         download_file(IMAGES_URL, tar_path)
         tar_file = tarfile.open(tar_path)
-        tar_file.extractall(data_path)
+        safe_extract(tar_file, data_path)
         tar_file.close()
         os.remove(tar_path)
 
         label_path = os.path.join(data_path, os.path.basename(LABELS_URL))
-        req = requests.get(LABELS_URL)
+        req = requests.get(LABELS_URL, timeout=60)
         with open(label_path, 'wb') as label_file:
             label_file.write(req.content)
 

learn2learn/vision/models/__init__.py

@@ -131,6 +131,6 @@ def get_pretrained_backbone(model, dataset, spec='default', root='~/data', downl
     elif model == 'wrn28':
         pretrained = WRN28Backbone()
 
-    weights = torch.load(destination, map_location='cpu')
+    weights = torch.load(destination, map_location='cpu', weights_only=True)
     pretrained.load_state_dict(weights)
     return pretrained

setup.py

@@ -3,7 +3,6 @@
 import sys
 import re
 
-from distutils.core import setup
 from setuptools import (
     setup as install,
     find_packages,

tests/unit/data/safe_extract_test.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+
+import os
+import shutil
+import tarfile
+import tempfile
+import unittest
+import zipfile
+
+from learn2learn.data.utils import safe_extract
+
+
+class SafeExtractTests(unittest.TestCase):
+
+    def setUp(self):
+        self.tmpdir = tempfile.mkdtemp()
+        self.source = os.path.join(self.tmpdir, 'source.txt')
+        with open(self.source, 'w') as f:
+            f.write('learn2learn')
+
+    def tearDown(self):
+        shutil.rmtree(self.tmpdir, ignore_errors=True)
+
+    def make_dir(self, name):
+        path = os.path.join(self.tmpdir, name)
+        os.makedirs(path)
+        return path
+
+    def test_benign_tar_extracts(self):
+        archive_path = os.path.join(self.tmpdir, 'benign.tar')
+        with tarfile.open(archive_path, 'w') as tar:
+            tar.add(self.source, arcname='sub/data.txt')
+        destination = self.make_dir('tar_out')
+        with tarfile.open(archive_path) as tar:
+            safe_extract(tar, destination)
+        self.assertTrue(
+            os.path.exists(os.path.join(destination, 'sub', 'data.txt'))
+        )
+
+    def test_benign_zip_extracts(self):
+        archive_path = os.path.join(self.tmpdir, 'benign.zip')
+        with zipfile.ZipFile(archive_path, 'w') as zf:
+            zf.writestr('sub/data.txt', 'learn2learn')
+        destination = self.make_dir('zip_out')
+        with zipfile.ZipFile(archive_path) as zf:
+            safe_extract(zf, destination)
+        self.assertTrue(
+            os.path.exists(os.path.join(destination, 'sub', 'data.txt'))
+        )
+
+    def test_malicious_tar_blocked(self):
+        archive_path = os.path.join(self.tmpdir, 'malicious.tar')
+        with tarfile.open(archive_path, 'w') as tar:
+            tar.add(self.source, arcname='../escaped.txt')
+        destination = self.make_dir('tar_mal_out')
+        with tarfile.open(archive_path) as tar:
+            self.assertRaises(Exception, safe_extract, tar, destination)
+        escaped = os.path.join(self.tmpdir, 'escaped.txt')
+        self.assertFalse(os.path.exists(escaped))
+
+    def test_malicious_zip_blocked(self):
+        archive_path = os.path.join(self.tmpdir, 'malicious.zip')
+        with zipfile.ZipFile(archive_path, 'w') as zf:
+            zf.writestr('../escaped.txt', 'pwned')
+        destination = self.make_dir('zip_mal_out')
+        with zipfile.ZipFile(archive_path) as zf:
+            self.assertRaises(Exception, safe_extract, zf, destination)
+        escaped = os.path.join(self.tmpdir, 'escaped.txt')
+        self.assertFalse(os.path.exists(escaped))
+
+    def test_rejects_unsupported_archive(self):
+        self.assertRaises(TypeError, safe_extract, object(), self.tmpdir)
+
+
+if __name__ == "__main__":
+    unittest.main()