From f43ff49c285cc5f6cf6517ca0a21dd5dae1acf99 Mon Sep 17 00:00:00 2001 From: fjankovi <161825881+fjankovi@users.noreply.github.com> Date: Tue, 25 Aug 2026 12:49:32 -0400 Subject: [PATCH] Harden LSUNClass key cache: store under dataset root and load with weights_only MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `LSUNClass.__init__` derived its key-cache filename from the ASCII letters of `root` and read/wrote it in the current working directory via `pickle.load` / `pickle.dump`: cache_file = "_cache_" + "".join(c for c in root if c in string.ascii_letters) if os.path.isfile(cache_file): self.keys = pickle.load(open(cache_file, "rb")) Two problems: - The path is resolved against the process CWD, not `root`, and the name is a predictable, lossy transform of `root`. Anyone able to write a file into the directory the user runs from (a shared scratch dir, a cloned repo, a CI workspace) can pre-plant `_cache_` and have it deserialized on the next `LSUN()` call. `pickle.load` on an attacker-controlled file is arbitrary code execution (CWE-502); there is no integrity check on the path. - Even setting that aside, the cache is a plain pickle with no restriction. Fix, matching the pattern already used by ImageNet/MNIST/PhotoTour in this package: - Store the cache next to the LMDB store as `root/_cache_keys.pt` (a trusted, per-class location; no CWD, no cross-root name collisions). - Load with `torch.load(..., weights_only=True)` so a cache file can never execute code on load. - Degrade gracefully when `root` is read-only (skip caching, re-enumerate). The cache only memoizes keys that are fully re-derivable from the trusted LMDB store, so old CWD caches are simply not found under the new path and are regenerated safely on first use — no migration needed. `torch.load` round-trips the raw `bytes` LMDB keys, including non-UTF-8 keys. Co-Authored-By: Claude Opus 5 (1M context) --- torchvision/datasets/lsun.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/torchvision/datasets/lsun.py b/torchvision/datasets/lsun.py index 6f6c7a5eb63..def01519bd0 100644 --- a/torchvision/datasets/lsun.py +++ b/torchvision/datasets/lsun.py @@ -1,11 +1,10 @@ import io import os.path -import pickle -import string from collections.abc import Iterable from pathlib import Path from typing import Any, Callable, cast, Optional, Union +import torch from PIL import Image from .utils import iterable_to_str, verify_str_arg @@ -23,13 +22,21 @@ def __init__( self.env = lmdb.open(root, max_readers=1, readonly=True, lock=False, readahead=False, meminit=False) with self.env.begin(write=False) as txn: self.length = txn.stat()["entries"] - cache_file = "_cache_" + "".join(c for c in root if c in string.ascii_letters) + # Cache the enumerated keys next to the LMDB store (a trusted, per-class + # location) rather than in the current working directory, and use the + # restricted ``weights_only`` unpickler so a cache file can never execute + # code on load. See https://github.com/pytorch/vision for context. + cache_file = os.path.join(root, "_cache_keys.pt") if os.path.isfile(cache_file): - self.keys = pickle.load(open(cache_file, "rb")) + self.keys = torch.load(cache_file, weights_only=True) else: with self.env.begin(write=False) as txn: self.keys = [key for key in txn.cursor().iternext(keys=True, values=False)] - pickle.dump(self.keys, open(cache_file, "wb")) + try: + torch.save(self.keys, cache_file) + except OSError: + # Read-only dataset directory: skip caching and re-enumerate next time. + pass def __getitem__(self, index: int) -> tuple[Any, Any]: img, target = None, None