You can deploy this model and start earning money today!

shiyamganesh/
tensorrt-qkv-runnerstate-poc

Name: tensorrt-qkv-runnerstate-poc
Availability: InStock

Views

No views yet

TensorRT QKV runnerStateBuffer PluginField length validation PoC package

Status

This repository is a triager-side source-level PoC package, not a final mutated TensorRT .engine file.

The bug candidate is a validation gap in NVIDIA TensorRT OSS QKV plugin runtime deserialization:

runnerStateBuffer is serialized as PluginFieldType::kUNKNOWN with mRunnerStateBuffer.size().
Runtime plugin creators store only fc->fields[i].data.
They do not validate fc->fields[i].type or fc->fields[i].length.
The constructor later calls dispatcher deserialization using dispatcher->getSerializationSize(), not the actual serialized field length.

The included ASan harness models the exact plugin-level validation gap and demonstrates an out-of-bounds read if a short runnerStateBuffer reaches the runtime creator.

Affected source snapshot

Tested TensorRT source snapshot / commit:

5302b288c5603895508711bdf2df11dc29c8ad92

Primary files:

text

plugin/bertQKVToContextPlugin/qkvToContextPlugin.cpp
plugin/bertQKVToContextPlugin/mhaRunner.cu
plugin/common/plugin.cpp
include/NvInferPluginBase.h
include/NvInferRuntime.h

Files in this package

text

README.md
qkv_runnerstate_harness.cpp
qkv_runnerstate_fixed_harness.cpp
expected_vulnerable_asan_output.txt
expected_fixed_output.txt
qkv_runnerstate_proposed_patch.diff
tensorrt_qkv_key_lines.txt
description_for_huntr.md
runtime_poc_requirements.md
NO_REAL_ENGINE_FILE_INCLUDED.txt

Reproduce the plugin-level issue

On Linux with clang or g++ and AddressSanitizer:

bash

clang++ -std=c++17 -fsanitize=address -g -O1 qkv_runnerstate_harness.cpp -o qkv_runnerstate_harness
./qkv_runnerstate_harness

Expected vulnerable result:

text

ERROR: AddressSanitizer: heap-buffer-overflow
READ of size 4

The expected output is also included in:

expected_vulnerable_asan_output.txt

Reproduce the fixed behavior

bash

clang++ -std=c++17 -fsanitize=address -g -O1 qkv_runnerstate_fixed_harness.cpp -o qkv_runnerstate_fixed_harness
./qkv_runnerstate_fixed_harness

Expected fixed behavior:

Rejected safely: invalid runnerStateBuffer length: rejected before deserialize

Important limitation

This package does not include a malformed TensorRT .engine/.trt/.mytrtfile.

The remaining triage question is whether TensorRT core allows a malformed serialized plugin field with a short runnerStateBuffer.length to reach:

createPlugin(..., TensorRTPhase::kRUNTIME)

If TensorRT core does not validate that field length before invoking the plugin creator, then the QKV plugin code can overread because it ignores the actual PluginField.length.

Suggested triager-side runtime test

NVIDIA/TensorRT triage can test the candidate by creating or mutating a serialized engine containing a QKV plugin runtime field collection where:

text

field.name   = "runnerStateBuffer"
field.type   = PluginFieldType::kUNKNOWN
field.data   = non-null
field.length = shorter than dispatcher->getSerializationSize()

Then deserialize the engine and verify whether the plugin creator reaches the runtime constructor and whether the dispatcher deserializer reads beyond the supplied field data.

Suggested fix

Validate the runtime runnerStateBuffer field metadata before deserialization:

cpp

PLUGIN_VALIDATE(fc->fields[i].type == PluginFieldType::kUNKNOWN,
    "invalid runnerStateBuffer field type");
runnerStateBufferLength = fc->fields[i].length;
runnerStateBuffer = static_cast<void const*>(fc->fields[i].data);

Before constructing/deserializing:

cpp

PLUGIN_VALIDATE(runnerStateBufferLength == expectedSerializationSize,
    "invalid runnerStateBuffer field length");

Prefer passing the validated actual length into the constructor/deserializer instead of recomputing an expected length independently of the serialized field metadata.

shiyamganesh/tensorrt-qkv-runnerstate-poc

TensorRT QKV runnerStateBuffer PluginField length validation PoC package

Status

Affected source snapshot

Files in this package

Reproduce the plugin-level issue

Reproduce the fixed behavior

Important limitation

Suggested triager-side runtime test

Suggested fix

shiyamganesh/
tensorrt-qkv-runnerstate-poc