README.md

# Flax msgpack DoS — Unbounded Chunk Allocation PoC

## Vulnerability

`flax.serialization.msgpack_restore()` calls `_unchunk()` which calls
`np.concatenate(_dict_to_tuple(data["chunks"]))` with **no bound** on chunk count or
cumulative allocation size. `MAX_CHUNK_SIZE=2**30` is an encoding-time guard only
and is never consulted during deserialization.

Root cause in `flax/serialization.py`:
```python
def _unchunk(data):
    chunks = data["chunks"]
    ...
    return np.concatenate(_dict_to_tuple(chunks))  # NO size/count limit
```

## Attack Vector

Attacker publishes a malicious `flax_model.msgpack` on HuggingFace.
Victim calls `FlaxPreTrainedModel.from_pretrained(repo)` which calls `msgpack_restore()`.
Result: process attempts O(N * chunk_size) RAM allocation and is OOM-killed.

## PoC File

`flax_model.msgpack` — 10 chunks x 1 MB = 10 MB demonstration file.
Scale to 10,000 chunks x 1 MB to exhaust 10 GB RAM on victim machine.

## Reproduce

```bash
pip install flax
python -c "import flax.serialization; flax.serialization.msgpack_restore(open('flax_model.msgpack','rb').read())"
```

## Fix

Add a cumulative size guard in `_unchunk()`:
```python
MAX_TOTAL = 2 ** 31  # 2 GiB hard cap
total = sum(len(v) if isinstance(v, (bytes, bytearray)) else getattr(v, 'nbytes', 0) for v in chunks.values())
if total > MAX_TOTAL:
    raise ValueError(f"Chunked array exceeds maximum safe size: {total}")
```