LOC 137 Parallel_Native

.py

#!/usr/bin/env python3
"""
Proof of Concept for Integer Overflow in PyTorch ParallelNative.cpp:137
Demonstrates the overflow behavior in Python (simulating the C++ logic)
"""

import sys

# Python 3 has arbitrary precision integers, so we need to simulate int64_t behavior
INT64_MAX = 2**63 - 1
INT64_MIN = -2**63

def to_int64(value):
    """Simulate int64_t overflow behavior"""
    # Wrap around if overflow occurs
    if value > INT64_MAX:
        return ((value - INT64_MIN) % (2**64)) + INT64_MIN
    elif value < INT64_MIN:
        return ((value - INT64_MIN) % (2**64)) + INT64_MIN
    return value

def divup(x, y):
    """Reproduce PyTorch's divup function with overflow simulation"""
    # This is the vulnerable implementation: (x + y - 1) / y
    result = to_int64(to_int64(x + y) - 1) // y
    return result

def calc_num_tasks_and_chunk_size_vulnerable(begin, end, grain_size, num_threads):
    """
    Reproduce the vulnerable function from ParallelNative.cpp:131-142
    """
    print(f"  Input: begin={begin}, end={end}, grain_size={grain_size}, threads={num_threads}")
    
    # Line 133: Check if (end - begin) < grain_size
    # This subtraction is also vulnerable to overflow!
    diff = to_int64(end - begin)
    print(f"  (end - begin) = {diff}")
    
    if diff < grain_size:
        print(f"  Early return: diff < grain_size")
        return (1, max(0, diff))
    
    # LINE 137: VULNERABLE CODE
    # int64_t chunk_size = divup((end - begin), get_num_threads());
    print(f"  Computing chunk_size via divup({diff}, {num_threads})...")
    
    # The divup expands to: ((end - begin) + num_threads - 1) / num_threads
    # This has multiple overflow points
    chunk_size = divup(diff, num_threads)
    print(f"  chunk_size = {chunk_size}")
    
    # Line 139
    chunk_size = max(grain_size, chunk_size)
    
    # Line 140: Another divup call
    num_tasks = divup(diff, chunk_size)
    print(f"  num_tasks = {num_tasks}")
    
    return (num_tasks, chunk_size)

def test_case_1():
    print("\n" + "="*70)
    print("Test Case 1: Large positive end, large negative begin")
    print("="*70)
    print("This tests overflow in (end - begin)")
    
    begin = -1000000000000  # Large negative
    end = INT64_MAX          # Maximum positive
    
    print(f"Expected: {end} - ({begin}) = {end - begin}")
    print(f"With int64_t overflow simulation:")
    
    result = calc_num_tasks_and_chunk_size_vulnerable(begin, end, 1000, 8)
    print(f"✓ Result: num_tasks={result[0]}, chunk_size={result[1]}")
    
    if result[0] < 0 or result[1] < 0:
        print("⚠️  OVERFLOW DETECTED: Negative result values!")
        return True
    return False

def test_case_2():
    print("\n" + "="*70)
    print("Test Case 2: Both extremes (maximum overflow scenario)")
    print("="*70)
    
    begin = INT64_MIN + 1000   # Near minimum
    end = INT64_MAX - 1000      # Near maximum
    
    print(f"begin near INT64_MIN, end near INT64_MAX")
    print(f"Expected difference: {end - begin}")
    
    result = calc_num_tasks_and_chunk_size_vulnerable(begin, end, 1, 8)
    print(f"✓ Result: num_tasks={result[0]}, chunk_size={result[1]}")
    
    # The actual result should be positive and huge
    # If we get negative or small values, overflow occurred
    expected_diff = end - begin
    if result[1] < expected_diff // 16:  # Should be at least diff/threads ballpark
        print("⚠️  OVERFLOW DETECTED: Result is suspiciously small!")
        return True
    return False

def test_case_3():
    print("\n" + "="*70)
    print("Test Case 3: Overflow in divup addition")
    print("="*70)
    print("When (end - begin) is near INT64_MAX, adding num_threads causes overflow")
    
    begin = 0
    end = INT64_MAX - 10      # Very close to max
    
    result = calc_num_tasks_and_chunk_size_vulnerable(begin, end, 1, 16)
    print(f"✓ Result: num_tasks={result[0]}, chunk_size={result[1]}")
    
    if result[1] < 0:
        print("⚠️  OVERFLOW DETECTED: Negative chunk_size!")
        return True
    return False

def test_case_4():
    print("\n" + "="*70)
    print("Test Case 4: Normal case (control - should work correctly)")
    print("="*70)
    
    begin = 0
    end = 1000000
    
    result = calc_num_tasks_and_chunk_size_vulnerable(begin, end, 1000, 8)
    print(f"✓ Result: num_tasks={result[0]}, chunk_size={result[1]}")
    
    # Should be reasonable values
    if result[0] > 0 and result[1] > 0 and result[1] <= end - begin:
        print("✓ Normal case works as expected")
        return False
    return True

def demonstrate_overflow_directly():
    print("\n" + "="*70)
    print("Direct Demonstration of Overflow in divup()")
    print("="*70)
    
    x = INT64_MAX - 10
    y = 16
    
    print(f"divup({x}, {y}) expands to:")
    print(f"  ({x} + {y} - 1) / {y}")
    print(f"  Step 1: {x} + {y} = {x + y}")
    
    if (x + y) > INT64_MAX:
        print(f"  ⚠️  OVERFLOW: Result {x + y} > INT64_MAX ({INT64_MAX})")
        print(f"  With int64_t wrapping: {to_int64(x + y)}")
    
    result = divup(x, y)
    print(f"  Final result: {result}")
    
    # Correct result without overflow
    correct_result = (x + y - 1) // y
    print(f"  Correct result (no overflow): {correct_result}")
    
    if result != correct_result:
        print(f"  ⚠️  MISMATCH: Got {result}, expected {correct_result}")
        return True
    return False

if __name__ == "__main__":
    print("="*70)
    print("PyTorch Integer Overflow PoC")
    print("Target: aten/src/ATen/ParallelNative.cpp:137")
    print("="*70)
    print(f"INT64_MAX = {INT64_MAX}")
    print(f"INT64_MIN = {INT64_MIN}")
    
    overflow_detected = False
    
    # Run tests
    overflow_detected |= test_case_4()  # Control test first
    overflow_detected |= test_case_1()
    overflow_detected |= test_case_2()
    overflow_detected |= test_case_3()
    overflow_detected |= demonstrate_overflow_directly()
    
    print("\n" + "="*70)
    print("SUMMARY")
    print("="*70)
    if overflow_detected:
        print("⚠️  INTEGER OVERFLOW VULNERABILITY CONFIRMED")
        print("The vulnerable code in ParallelNative.cpp:137 can overflow")
        print("when given extreme input values.")
    else:
        print("Tests completed. Check output for anomalies.")
    
    print("\nNote: In actual C++ code with int64_t, this overflow causes")
    print("undefined behavior which may crash, produce wrong results,")
    print("or appear to work depending on compiler optimization.")
    print("="*70)