orbax-checkpoint-uncontrolled-recursion-poc

orbax-checkpoint CWE-674 PoC

Target

Root Cause

_unchunk_array_leaves_in_place(d) in orbax/checkpoint/msgpack_utils.py recursively traverses a dict deserialized from a msgpack checkpoint file with no depth guard. A 1,001-level nested dict (3,009 bytes) triggers RecursionError.

def _unchunk_array_leaves_in_place(d):  # no depth guard
    for k, v in d.items():
        elif isinstance(v, dict):
            _unchunk_array_leaves_in_place(v)  # ← unbounded recursion

Trigger Path

victim: checkpointer.restore(malicious_dir/)
  → MsgpackHandler.deserialize(dir/checkpoint)
  → msgpack_restore(bytes)
  → _unchunk_array_leaves_in_place(state_dict)  ← RecursionError at depth 1001

Reproduce

pip install orbax-checkpoint msgpack
python3 poc_orbax.py

Expected output:

[*] Python recursion limit: 1000
[*] Payload: 3009 bytes, depth=1001
[+] CONFIRMED: RecursionError in _unchunk_array_leaves_in_place
[+] CONFIRMED via MsgpackHandler: maximum recursion depth exceeded

Fix

def _unchunk_array_leaves_in_place(d, _depth=0):
    if _depth > 500:
        raise ValueError('Checkpoint nesting depth exceeds limit')
    ...
    _unchunk_array_leaves_in_place(v, _depth + 1)

Files