Skip to content

Fix pickle complexity DoS: add memo_get node budget - #3688

Closed
AnkitNakhawa wants to merge 3 commits into
huggingface:mainfrom
AnkitNakhawa:main
Closed

AnkitNakhawa wants to merge 3 commits into
huggingface:mainfrom
AnkitNakhawa:main

Conversation

@AnkitNakhawa

@AnkitNakhawa AnkitNakhawa commented Jun 30, 2026 •

Copy link
Copy Markdown

Fix pickle memo bomb: add cumulative complexity limit to memo_get (closes #3620)

What

Adds a cumulative node-budget guard to Stack::memo_get in candle-core/src/pickle.rs to prevent the algorithmic-complexity DoS described in #3620.

Root cause

memo_get returned obj.clone() unconditionally — a deep clone of the entire Object tree on every BinGet. A crafted pickle that fetches the same memo slot twice, combines the copies, and writes the result back doubles the node count each cycle. After N cycles, parsing requires O(2^N) work and memory from an input whose size grows only linearly with N (the same "Billion Laughs" amplification principle). Two paths were confirmed:

  • Path A (CPU): BinGet + BinGet + Tuple2 + BinPut — N=25 stalls ~5.7s on a release build
  • Path B (memory): BinGet + BinGet + Build (dict-merge) + BinPut — ~92 GB RSS measured before OOM on a 128 GB host

Fix

memo_get now counts the nodes in the object being retrieved, checks the budget, and only clones if the budget is not exceeded. The cost is added to a running complexity counter on the Stack; if it exceeds max_complexity (default: 1,000,000 nodes),parsing is aborted before the expensive object-tree clone occurs. The node walk is iterative rather than recursive so a deeply-nested input cannot stack-overflow before the guard fires.

The 1,000,000-node default is a conservative initial bound — real checkpoint complexity is estimated well below this, but it has not been validated against large production checkpoints and can be adjusted based on maintainer guidance.

Stack::with_limit(n) is exposed for callers that need a custom budget. u64::MAX disables the check entirely.

Tests added

  • Exact PoC bytes from the issue report (both Path A and Path B) are rejected
  • Exponential tuple-doubling and dict-merge payloads exceeding the default limit are rejected
  • A deeply nested memo-expansion payload is rejected cleanly — node counting uses iterative traversal rather than recursion, so a deeply-nested input cannot stack-overflow before the guard fires
  • Normal small payloads and no-memo pickles continue to parse successfully
  • Custom and zero budgets behave as documented

Closes #3620

@astorise

Copy link
Copy Markdown
Contributor

Heads-up for triage: there's meaningful overlap with #3628 ("Bound the pickle VM's working set and nesting depth"). #3628's byte-based working-set cap also bounds the memo-replay amplification this node budget targets — each memo-expanded clone charges bytes against the cap, and #3628 even ships a rejects_memo_replay_amplification test for exactly that DoS. So the two guard largely the same surface via different metrics (cumulative node count here vs. a heap-bytes ceiling there). Might be worth deciding whether both are wanted or whether they should be unified, to avoid two overlapping limits on the same path.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Pickle memo bomb — CPU + memory exhaustion (algorithmic-complexity DoS) in candle-core pickle reader

2 participants