EnigmaConsultant/flax-msgpack-dos-poc

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:

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

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

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}")