trevdatastreams/armnn-reduce-axis-shape-oob-poc

Arm NN native Reduce axis heap out-of-bounds write

Summary

A native .armnn model can give a ReduceLayer an axis that is outside the input tensor rank. The deserializer copies the serialized axis vector into ReduceDescriptor::m_vAxis without validating its elements.

During normal armnn::Optimize() shape inference, ReduceLayer subtracts the number of serialized axes from the input rank to size an output-dimension vector. It then copies every input dimension whose numeric index is not in the axis vector. An out-of-range axis therefore reduces the allocation without removing any input dimension from the copy loop. The last copy writes past the heap allocation.

The included differential models are both 580 bytes. They differ by one byte:

• control: axis = [2]
• trigger: axis = [100]

The control loads and optimizes successfully in three of three runs. The trigger produces an AddressSanitizer heap-buffer-overflow write in three of three runs.

Target

• Project: Arm NN
• Commit: 2b61cecc9df7a43fca1463795062cf359e6be820
• Public entry points:
  • armnnDeserializer::IDeserializer::CreateNetworkFromBinary()
  • armnn::Optimize()
• Backend used for optimization: CpuRef
• Sanitizers: AddressSanitizer and UndefinedBehaviorSanitizer
• GPU: not required

Reproduce

From the workspace root, with Arm NN already built with ASan/UBSan:

./cyber/huntr-mfv/candidates/armnn-flatbuffers-reduce-axis-shape-oob-write/reproduce.sh "$PWD"

Expected control output:

loaded and optimized 580 bytes

Expected trigger result:

ERROR: AddressSanitizer: heap-buffer-overflow
WRITE of size 4
...
armnn::ReduceLayer::InferOutputShapes(...)
armnn::ReduceLayer::ValidateTensorShapesFromInputs()
armnn::Graph::InferTensorInfos()
armnn::Optimize(...)

The captured control-{1,2,3}.txt and trigger-{1,2,3}.txt files contain three independent runs of each model. Each trigger exits with status 134.

Root cause

The native deserializer copies every serialized axis value without checking it against the input rank:

descriptor.m_vAxis =
    std::vector<unsigned int>(flatBufferAxis->begin(), flatBufferAxis->end());

ReduceLayer::InferOutputShapes() then derives:

outputRank = input.GetNumDimensions() -
             static_cast<unsigned int>(m_Param.m_vAxis.size());
std::vector<unsigned int> dimSizes(outputRank, 1);

For the trigger, the input rank is four and the one-element axis vector is [100], so dimSizes has three elements. The following loop finds that none of the valid input dimensions 0..3 equals 100, then writes all four input dimensions into that three-element vector:

for (unsigned int i = 0; i < input.GetNumDimensions(); ++i)
{
    if (std::find(m_Param.m_vAxis.begin(),
                  m_Param.m_vAxis.end(),
                  i) == m_Param.m_vAxis.end())
    {
        dimSizes[outputIndex] = input[i];
        ++outputIndex;
    }
}

ASan reports the fourth four-byte write immediately after the 12-byte heap region allocated for the three-element vector.

Fixture integrity

e2ccb82ef498586447658446fd23bdaac77b29d77f316d4b6bb77b57f2b25de2  control-axis-2.armnn
274f197d458ff0f911e3080eab2a4e38fc9e29d69200a2cc3eca22c21f4aa16e  trigger-axis-100.armnn

cmp -l reports a single differing byte at offset 249: octal 002 in the control and octal 144 (100 decimal) in the trigger.

Security impact

An application that accepts an untrusted native Arm NN model and prepares it for execution can be terminated by a deterministic heap out-of-bounds write during graph optimization. The model passes FlatBuffers structural verification and deserializes before the write occurs. No inference inputs, large dimensions, or GPU are required.

This report demonstrates native heap corruption and denial of service. It does not claim code execution.

Prior-art review

Fresh searches of Hugging Face, Samsung/Arm public issue trackers, GitHub, and the local report corpus found no public report for the native ReduceDescriptor.axis to ReduceLayer::InferOutputShapes() path.

The finding is distinct from public Arm NN issues involving:

• TensorInfo dimension-vector mismatches
• tensor element-count integer overflow
• Splitter descriptor-count mismatches
• optional-field null dereferences
• input/output slot index resolution
• ScatterNd runtime shape metadata

Suggested fix

Before adding the layer, reject any Reduce axis that is greater than or equal to the input rank and reject duplicate axes. InferOutputShapes() should also validate the descriptor independently and use checked writes rather than assuming the axis vector contains a unique subset of input dimensions.