trevdatastreams/armnn-tflite-validation-subgraph-index-oob-poc

Arm NN TFLite leading validation-subgraph index PoC

This package reproduces a deterministic heap-buffer-overflow read in Arm NN's TFLite parser. The parser explicitly accepts models containing validation subgraphs named with the VALIDATION: prefix and selects the sole non-validation subgraph by its original index. It then appends exactly one entry to m_SubgraphConnections and indexes that one-entry vector with the original subgraph index.

When a validation subgraph precedes the main subgraph, the selected main subgraph index is 1 while m_SubgraphConnections.size() is 1. RegisterConsumerOfTensor() consequently reads one vector object past the 24-byte heap allocation.

Tested revision

• Repository: ARM-software/armnn
• Commit: 2b61cecc9df7a43fca1463795062cf359e6be820
• Entry point: armnnTfLiteParser::ITfLiteParser::CreateNetworkFromBinaryFile()
• Sanitizers: AddressSanitizer and UndefinedBehaviorSanitizer

Minimal semantic differential

Both fixtures are 432 bytes and contain the same two subgraphs:

• A valid ABS main subgraph.
• An empty subgraph named VALIDATION:placeholder.

Only their order changes:

• Control: [main, VALIDATION:placeholder]; the selected main index is 0.
• Trigger: [VALIDATION:placeholder, main]; the selected main index is 1.

SHA-256:

7404e7e5b1c2c302738488d23445ac7c9dd1377ee149551a5a9d6123950d8974  control-main-first.tflite
f8fae3c4bcb69929816a4a09829755cbb5dbedae72cb465209c1c752897e276a  trigger-validation-first.tflite

The JSON source files make the order-only semantic differential explicit.

Reproduction

Build Arm NN with its TFLite parser and sanitizers enabled. Compile the small loader against the resulting libraries:

clang++ -std=c++17 -fsanitize=address,undefined -fno-omit-frame-pointer \
  -I /path/to/armnn/include \
  load-tflite.cpp \
  -L /path/to/armnn/build \
  -larmnnTfLiteParser -larmnn \
  -Wl,-rpath,/path/to/armnn/build \
  -o load-tflite

Run the control:

./load-tflite models/control-main-first.tflite

Expected:

parsed successfully

Run the trigger:

./load-tflite models/trigger-validation-first.tflite

AddressSanitizer reports:

ERROR: AddressSanitizer: heap-buffer-overflow
READ of size 8
#0 armnnTfLiteParser::TfLiteParserImpl::RegisterConsumerOfTensor(...)
#1 armnnTfLiteParser::TfLiteParserImpl::RegisterInputSlots(...)
#2 armnnTfLiteParser::TfLiteParserImpl::ParseElementwiseUnary(...)
#3 armnnTfLiteParser::TfLiteParserImpl::CreateNetworkFromModel()
#4 armnnTfLiteParser::TfLiteParserImpl::CreateNetworkFromBinaryFile(char const*)

The read occurs exactly zero bytes after the one-entry, 24-byte m_SubgraphConnections allocation. In three repeated trials, the control exited 0 three times and the trigger aborted with exit code 134 three times.

asan-output.txt contains the complete sanitizer report and repeated-results.txt contains the repetition summary.

Regenerating the fixtures

The fixtures were generated from TensorFlow's official tensorflow/compiler/mlir/lite/schema/schema.fbs using FlatBuffers flatc:

python3 generate_models.py \
  --flatc /path/to/flatbuffers/flatc \
  --schema /path/to/tensorflow/compiler/mlir/lite/schema/schema.fbs

Root cause and suggested fix

At the tested revision, TfLiteParser.cpp lines 1040–1091 retains the model's original subgraphIndex but allocates connection state with:

m_SubgraphConnections.emplace_back(subgraph->tensors.size());

Later, RegisterConsumerOfTensor() lines 1166–1173 uses:

TensorSlots& tensorSlots =
    m_SubgraphConnections[subgraphIndex][tensorIndex];

Allocate m_SubgraphConnections by the model's subgraph count and initialize the selected original index, or consistently map the selected subgraph to a local zero-based connection index. Add regression tests with validation subgraphs both before and after the main subgraph.

Impact

An attacker-controlled TFLite model can cause a deterministic native out-of-bounds read and process termination during model loading, before inference. The same invalid outer-vector reference is also used by RegisterProducerOfTensor() for output registration. This PoC establishes denial of service; it does not claim data disclosure or code execution.