pytorch - 💡(How to fix) Fix [JIT] Uncontrolled memory allocation in Pickle Unpickler opcode handlers (setInput / readBytes)

Summary

torch::jit::Unpickler processes multiple opcodes where the size/length argument is read from the pickle stream and passed directly to a memory-allocating call with no validation. Two distinct opcode-handler sites are confirmed:

• Variant A (setInput): container-size opcode reads n, calls _M_default_append(n). With n = 28,257,437,280 the allocation request is ~226 GB.
• Variant B (readBytes): BINBYTES8/BINUNICODE8 reads 8-byte length, calls readBytes(length). With length = 4,194,304,000 the request is 4 GB.

Both trigger immediate OOM on any standard deployment. A 1.6 KB .pt file is sufficient. Bug fires before any tensor data is touched, so ZIP-layer hardening (2.11) does not protect against it.

Previously reported privately as GHSA-84cc-wfj8-w393; @malfet suggested filing as a public issue.

Root cause

// torch/csrc/jit/serialization/unpickler.cpp

// Variant A: no cap on attacker-controlled container size
void Unpickler::setInput(size_t n) {
    stack_.resize(n);  // n from pickle stream, unchecked
}

// Variant B: no cap on attacker-controlled blob length
std::string Unpickler::readBytes(size_t length) {
    // allocates `length` bytes directly from stream field
}

Crash signals

==PID== ERROR: libFuzzer: out-of-memory (malloc(28257437280))  # Variant A
==PID== ERROR: libFuzzer: out-of-memory (malloc(4194304000))   # Variant B

Reproduction

import torch
torch.jit.load("poc-059-pickle-uncontrolled-alloc-1613b.pt")
# → MemoryError / process killed by OOM

PoC files available on request (1,613 bytes and 2,156 bytes).

Suggested fix

Validate announced size before allocating at every opcode handler that reads an attacker-controlled size_t from the pickle stream:

namespace {
constexpr size_t kMaxPickleContainerElems = 1u << 30;  // 1B elements
constexpr size_t kMaxPickleBlobBytes      = 1u << 32;  // 4 GiB
}

void Unpickler::setInput(size_t n) {
    if (n > kMaxPickleContainerElems)
        AT_ERROR("Pickle container size out of bounds: ", n);
    stack_.resize(n);
}

std::string Unpickler::readBytes(size_t length) {
    if (length > kMaxPickleBlobBytes)
        AT_ERROR("Pickle blob size out of bounds: ", length);
    // ...
}

A full audit of Unpickler::readInstruction() for every (size_t n = read_uint64(); allocate(n)) pattern is recommended — the affected opcode set includes BINBYTES, BINBYTES8, BINUNICODE, BINUNICODE8, EMPTY_LIST, EMPTY_TUPLE, EMPTY_DICT.

Environment

• PyTorch 2.11.0+cu130 (confirmed); likely all versions
• Linux x86_64
• Crash is deterministic