ericblackgachara/flax-uncontrolled-recursion-poc

flax — CWE-674 Uncontrolled Recursion PoC

Target

Field	Value
Repository	google/flax
Version	v0.12.7 (latest)
Commit	0f128b5141a7ab40df430d280bc86d3929ccb1ce
Platform	huntr.com
CWE	CWE-674 (Uncontrolled Recursion)
CVSS	7.5 High

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.