Damir2024/ArmNN-poc

Heap Buffer Overflow in Arm NN Deserializer via Integer Overflow

Summary

A critical heap buffer overflow vulnerability exists in the Arm NN library when deserializing maliciously crafted Flatbuffer models (.armnn). The vulnerability is caused by multiple integer overflows in the TensorInfo::GetNumBytes() and TensorShape::GetNumElements() functions, leading to undersized memory allocations and subsequent out-of-bounds writes.

Vulnerability Details

• Component: armnnDeserializer
• File: src/armnn/Tensor.cpp, src/armnnDeserializer/Deserializer.cpp
• Vulnerability Type: Heap-based Buffer Overflow (CWE-122) / Integer Overflow (CWE-190)
• Impact: Arbitrary Code Execution (ACE) or Denial of Service (DoS)

Root Cause Analysis

In src/armnn/Tensor.cpp, the GetNumElements() function calculates the total number of elements in a tensor by multiplying its dimensions:

unsigned int TensorShape::GetNumElements() const
{
    // ...
    unsigned int count = 1;
    for (unsigned int i = 0; i < m_NumDimensions; ++i)
    {
        if (m_DimensionsSpecificity[i])
        {
            count *= m_Dimensions[i]; // Potential Overflow
        }
    }
    return count;
}

And GetNumBytes() uses this count:

unsigned int TensorInfo::GetNumBytes() const
{
    return GetDataTypeSize(m_DataType) * GetNumElements(); // Potential Overflow
}

On systems where unsigned int is 32-bit, providing large dimensions (e.g., [65536, 65536]) causes count to overflow. For example, 65536 * 65536 = 4294967296, which overflows to 0.

Exploitation Path

In src/armnnDeserializer/Deserializer.cpp, several layer parsers (e.g., ParseFullyConnected, ParseConvolution2d) use GetNumBytes() to allocate buffers:

void DeserializerImpl::ParseFullyConnected(GraphPtr graph, unsigned int layerIndex)
{
    // ...
    TensorInfo weightsInfo = ToTensorInfo(weights->info());
    std::unique_ptr<unsigned char[]> permuteBuffer(new unsigned char[weightsInfo.GetNumBytes()]); // Undersized allocation
    
    // ...
    armnnUtils::Permute(weightsInfo.GetShape(), permutationVector,
                       weights.GetMemoryArea(), permuteBuffer.get(),
                       GetDataTypeSize(weightsInfo.GetDataType())); // OOB Write
}

If weightsInfo.GetNumBytes() overflows to a small value, permuteBuffer is allocated with insufficient space. The armnnUtils::Permute function then writes the actual weight data into this undersized buffer, causing a heap buffer overflow.

Proof of Concept (PoC)

A malicious .armnn file can be constructed with a FullyConnectedLayer containing a weights tensor with dimensions [65536, 65536].

Conceptual Schema Representation:

{
  "layers": [
    {
      "layer_type": "FullyConnected",
      "layer": {
        "base": { "index": 0, "layerName": "FC", ... },
        "weights": {
          "info": {
            "dimensions": [65536, 65536],
            "dataType": "Float32"
          },
          "data": { ... huge amount of data ... }
        }
      }
    }
  ]
}

When this model is loaded via IDeserializer::CreateNetworkFromBinary, the overflow is triggered, leading to a crash or potential code execution.

Impact

This vulnerability allows an attacker to execute arbitrary code or cause a denial of service by providing a malicious model file to an application using the Arm NN library. Given that Arm NN is used in edge devices and mobile applications for AI inference, the impact is significant.

Recommendations

1. Use 64-bit integers for size calculations: Change GetNumElements and GetNumBytes to return size_t and use uint64_t for internal calculations.
2. Implement Overflow Checks: Use safe math libraries or explicit checks (e.g., __builtin_mul_overflow) before performing multiplications in GetNumElements and GetNumBytes.
3. Validate Dimensions: Add sanity checks on tensor dimensions to ensure they do not exceed reasonable limits that could cause overflows.