# 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):

```python
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):

```python
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:

```python
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.