trevdatastreams/armnn-tflite-splitv-count-mismatch-heap-overflow-poc

Arm NN TFLite SPLIT_V model-load heap-buffer-overflow write

Arm NN's TFLite ParseSplitV() implementation trusts two independent, model-controlled descriptions of the split count. It allocates a destination vector from SplitVOptions.num_splits, then copies the entire declared split_sizes tensor into that vector without checking that their element counts match.

Tested against Arm NN commit 2b61cecc9df7a43fca1463795062cf359e6be820.

Root cause

The parser obtains numSplits from the operator options and allocates that many integers:

numSplits = CHECKED_NON_NEGATIVE(options->num_splits);
std::vector<int> splitsData(numSplits);

It independently validates the source buffer against the tensor's declared shape, then copies all source bytes:

ValidateBuffer(splitsBufferPtr, splitsInfo, "splits");
::memcpy(splitsData.data(),
         splitsBufferPtr->data.data(),
         splitsInfo.GetNumBytes());

No check requires splitsInfo.GetNumElements() == numSplits.

Affected source:

https://github.com/ARM-software/armnn/blob/2b61cecc9df7a43fca1463795062cf359e6be820/src/armnnTfLiteParser/TfLiteParser.cpp#L4619-L4704

Differential proof

Both models are verifier-valid 608-byte FlatBuffers generated from TensorFlow's official TFLite schema. They differ at exactly one byte:

Artifact	SplitVOptions.num_splits	split_sizes	Result
control.tflite	2	shape [2], values [1, 3]	Loads; exit 0
poc.tflite	1	shape [2], values [1, 3]	ASan heap-buffer-overflow write

Binary comparison:

$ wc -c control.tflite poc.tflite
608 control.tflite
608 poc.tflite

$ cmp -l control.tflite poc.tflite
277   2   1

The PoC causes an eight-byte memcpy into a four-byte allocation. The attacker-controlled second split value is written immediately after the heap object. The fault reproduced in three independent executions.

PoC SHA-256: 68690b331dd3f6e4410bda79f5a8bd69cbe1c7ce956e1215949ebb6e0de5ef8f.

Control SHA-256: fa97d2808ce0ca95dfbab8db4336738f8907d8f991a3bf58381034a9ed4d8580.

Build and run

Build Arm NN with its TFLite parser and AddressSanitizer, then compile the included public-API harness:

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

Run:

ASAN_OPTIONS=detect_leaks=0:halt_on_error=1 \
UBSAN_OPTIONS=print_stacktrace=1:halt_on_error=1 \
./load-tflite control.tflite

ASAN_OPTIONS=detect_leaks=0:halt_on_error=1 \
UBSAN_OPTIONS=print_stacktrace=1:halt_on_error=1 \
./load-tflite poc.tflite

Impact

An application that loads an untrusted .tflite model through Arm NN's public parser API can suffer attacker-controlled heap corruption during parsing, before inference. The demonstrated impact is a deterministic denial of service. This report does not claim code execution.

Suggested fix

Before allocating or copying, require:

• split_sizes to be a one-dimensional INT32 tensor;
• splitsInfo.GetNumElements() == options->num_splits; and
• the number of operator outputs to equal the validated split count.

Reject mismatches with a clean ParseException and add a regression test for the supplied one-byte differential.

Novelty

The pre-work and post-reproduction scans found no matching Hugging Face model, GitHub issue, pull request, or other local candidate. The single post-run local match was this candidate's own initial gate record.