EnigmaConsultant/huntr-poc-jax-deserialize-dtype-enum-keyerror-dos

Unvalidated DType enum in jax.export.deserialize causes uncaught KeyError DoS

Summary

The public jax.export.deserialize() entry point parses an attacker-controlled flatbuffer (Exported blob) without validating an enum field. The very first thing done during deserialization is building the version-10/11 unique_avals table; each aval is passed to _deserialize_aval(), which looks up the aval's Dtype enum in a Python dict with _dtype_kind_to_dtype[aval.Dtype()]. aval.Dtype() is a signed int8 read straight from the untrusted buffer; only values 0..29 are valid keys. Any other byte (e.g. 30, or a negative byte like 0xFF = -1) raises an uncaught KeyError that propagates out of the public API and terminates the consuming process — a denial of service for any service that deserializes exported JAX functions from untrusted input.

Target

• Package: jax / jaxlib (google/jax) — jax.export serialization
• Version tested: jax==0.11.0, jaxlib==0.11.0 (unmodified, from PyPI)
• Dependencies: flatbuffers==25.12.19, CPython 3.13.12, Linux x86-64
• Affected file: jax/_src/export/serialization.py
• Entry point: jax.export.deserialize(blob)

Root cause

_deserialize_exported() iterates the unique_avals table before any other field is parsed (serialization.py lines 284–286):

for i in range(exp.UniqueAvalsLength()):
    scope = ...
    _deserialize_aval(exp.UniqueAvals(i), scope=scope, sharding=None)

_deserialize_aval() then does an unchecked dict subscript on the raw enum (serialization.py line 825):

def _deserialize_aval(aval: ser_flatbuf.AbstractValue, *,
                      scope: shape_poly.SymbolicScope,
                      sharding: named_sharding.NamedSharding | None,
                      ) -> core.ShapedArray:
  dtype = _dtype_kind_to_dtype[aval.Dtype()]   # <-- line 825: no validation

aval.Dtype() returns the int8 DType enum straight from the flatbuffer. _dtype_kind_to_dtype only maps enum values 0..29; any other value is not a key and raises KeyError. Because the lookup fires for every aval before any other field of the blob is touched, it is the most reachable instance of this pattern.

Sibling instances of the same unvalidated-enum → dict-lookup pattern

The same shape exists (less reachable) at:

• line 833: mem_space = _memory_space_from_enum[ser_mem_space]
• line 707: _axis_type_from_enum[ser_mesh.AxisTypes(i)]

A complete fix should validate all three (e.g. .get() with a raised ValueError/typed deserialization error instead of a bare KeyError).

Proof of concept

Two independent reproductions are included, both against unmodified PyPI jax==0.11.0.

1. build_bad_dtype.py — crafts a minimal Exported flatbuffer from scratch whose UniqueAvals[0].Dtype = 30, then jax.export.deserialize raises KeyError: 30 at line 825.

2. mutate_real_blob.py — the strong differential PoC. It exports and serializes a genuine function (export.export(jax.jit(f))(jnp.ones((3,), f32), jnp.ones((3,), f32)) → serialize(), 1064 bytes), locates the single Dtype byte (absolute offset 219, value 10 = float32) via the flatbuffer vtable, and flips only that one byte to 30. This is a single-byte differential.

   • Negative control: the unmutated 1064-byte blob deserializes cleanly to in_avals = (ShapedArray(float32[3]), ShapedArray(float32[3])).
   • Crash: the byte-flipped copy raises the identical KeyError: 30.

   The negative control isolates the crash to the single unvalidated Dtype enum byte and nothing else.

victim.py is a realistic consumer that simply calls export.deserialize(blob) with no try/except.

Files

• build_bad_dtype.py — synthetic minimal-blob PoC
• mutate_real_blob.py — real-blob single-byte differential PoC + negative control
• victim.py — minimal consumer that crashes
• real_mutated_dtype30.bin — the 1064-byte real blob with only the Dtype byte flipped to 30
• bad_dtype.bin — the synthetic minimal malicious blob

Captured evidence (verbatim)

Building the differential + negative control:

$ python mutate_real_blob.py
serialized bytes: 1064
UniqueAvalsLength: 1
UniqueAvals[0].Dtype (valid) = 10
Dtype byte at abs offset 219 current value 10
NEG CONTROL unmutated deserialize OK, in_avals = (ShapedArray(float32[3]), ShapedArray(float32[3]))
wrote real_mutated_dtype30.bin (only the Dtype byte changed to 30)

Feeding the single-byte-mutated blob to a realistic consumer:

$ python victim.py real_mutated_dtype30.bin
Traceback (most recent call last):
  File ".../jax-5thbug/victim.py", line 5, in <module>
    exp = export.deserialize(blob)   # realistic consumer, no try/except
  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 285, in _deserialize_exported
    _deserialize_aval(exp.UniqueAvals(i), scope=scope, sharding=None)
  File ".../jax/_src/export/serialization.py", line 825, in _deserialize_aval
    dtype = _dtype_kind_to_dtype[aval.Dtype()]
            ~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^
KeyError: 30
EXIT=1

Environment confirmation:

jax 0.11.0
jaxlib 0.11.0
flatbuffers 25.12.19
py 3.13.12

Impact

Uncaught exception / denial of service. Any process that deserializes an Exported blob from an untrusted or semi-trusted source (model registries, caches, RPC payloads, uploaded artifacts) can be crashed with a single malformed byte. deserialize() is documented to only be safe for trusted input for RCE reasons, but a plain KeyError from a one-byte enum corruption is an availability bug distinct from and cheaper than any code-execution concern, and is trivially triggered by fuzzing or truncation/corruption in transit.

Suggested fix

Validate the enum before the lookup, raising a typed deserialization error:

kind = aval.Dtype()
try:
    dtype = _dtype_kind_to_dtype[kind]
except KeyError:
    raise ValueError(f"Unknown DType enum value in serialized aval: {kind}")

Apply the same guard to _memory_space_from_enum (line 833) and _axis_type_from_enum (line 707).

Deduplication

• No CVE is currently assigned to this specific unvalidated-Dtype-enum KeyError in jax.export's deserializer.
• Distinct from prior JAX deserialization findings in this research batch, which target different code paths: pytree recursion DoS, the index-table DoS, and the shape-poly power OOM. This finding is the _deserialize_aval DType enum dict-lookup specifically, reachable before any of those paths, and is a plain uncaught KeyError rather than resource exhaustion.