ericblackgachara/flax-uncontrolled-recursion-poc

flax — CWE-674 Uncontrolled Recursion PoC

Target

Root Cause

_unchunk_array_leaves_in_place() in flax/serialization.py recurses without a depth limit. msgpack allows up to 1024 nesting levels (C-layer) but Python's recursion limit is 1000. A checkpoint with depth 1001–1024 decodes fine but crashes Python on traversal.

Vulnerable Code

# flax/serialization.py
def _unchunk_array_leaves_in_place(d):
    if isinstance(d, dict):
        if '__msgpack_chunked_array__' in d:
            return _unchunk(d)
        else:
            for k, v in d.items():
                elif isinstance(v, dict):
                    _unchunk_array_leaves_in_place(v)  # ← NO DEPTH GUARD

Trigger Path

load checkpoint file
→ serialization.msgpack_restore()
→ _unchunk_array_leaves_in_place()  ← RecursionError here

Reproduction

pip install flax msgpack
python3 poc_flax.py
# Expected: RecursionError in _unchunk_array_leaves_in_place

Payload

• —File size: 3,008 bytes (2.9 KB)
• —Depth: 1001 levels (Python limit = 1000, msgpack limit = 1024)
• —Trigger: Any call to msgpack_restore(), from_bytes(), or restore_checkpoint()

Fix

Add a depth guard parameter: def _unchunk_array_leaves_in_place(d, _depth=0): Raise ValueError when _depth > 500. Or convert to iterative traversal.