Title

Heap buffer overflow in TFLite FullyConnectedPerChannel via WebNN quantized GEMM with wrong per-axis quantization dimension

Summary

The WebNN-to-TFLite translation layer fuses a quantized GEMM into TFLite’s FULLY_CONNECTED operator without validating that the weight tensor’s per-channel quantization axis is dimension 0 (output channels). An attacker can supply a weight tensor with per-channel quantization along dimension 1 (input channels) instead, causing PrepareImpl to allocate per-channel multiplier and shift arrays sized to the input channel count rather than the output channel count. When the kernel evaluates, it iterates over output channels and reads past the end of those arrays, producing a heap-buffer-overflow in the GPU process. Affected platforms: all where the TFLite WebNN backend is active (Linux, ChromeOS, Android by default; Windows and macOS as a fallback). No special GPU hardware is required.

Bisect

TFLite upstream — FullyConnectedPerChannel introduced

• Commit: 310635cc811ae3a64b3abc184885f6769c84f4d7
• Date: 2022-06-12
• Author: A. Unique TensorFlower
• Subject: TFLite FC Layer Per-channel quantization (Full-Integer i8xi8->i8 or i16xi8->i16)
• Piper: PiperOrigin-RevId: 454517886
• URL: https://github.com/tensorflow/tensorflow/commit/310635cc811ae3a64b3abc184885f6769c84f4d7

This commit introduced per-channel quantization for FULLY_CONNECTED in a single change: PrepareImpl allocates per_channel_output_multiplier and per_channel_output_shift arrays sized to filter->dims->data[quantized_dimension], with only a self-referential consistency check (scale->size == filter->dims->data[quantized_dimension]) and no enforcement that quantized_dimension == 0. The optimized FullyConnectedPerChannel kernel then iterates over filter->dims->data[0] (output channels) elements, causing a heap OOB read when quantized_dimension != 0 and dim 0 > dim quantized_dimension.

Became web-reachable

• Commit: 6e5458b43fdc3685cb17664e6709bd79a97dfb36
• Date: 2025-06-20
• Author: Junwei Fu (junwei.fu@intel.com)
• Subject: webnn: Fuse quantized GEMM into TFLite FULLY_CONNECTED
• Review: https://chromium-review.googlesource.com/c/chromium/src/+/6641102

This commit added quantized GEMM fusion in the WebNN TFLite graph builder without validating that the weight tensor’s per-channel quantization axis is dimension 0.

Root Cause

WebNN’s dequantizeLinear allows per-channel quantization where the scale tensor has a non-unit dimension indicating the quantization axis. When fusing a quantized GEMM into TFLite’s FULLY_CONNECTED, CanFuseQuantizeAndGetOutput(const mojom::Gemm&) checks only whether B’s scale has more than one element and whether bTranspose is true:

// services/webnn/tflite/graph_builder_tflite.cc:2072-2083
size_t number_of_b_scale =
    GetOperand(b_dequantize.scale_operand_id).descriptor.NumberOfElements();
const bool per_channel_quantization = number_of_b_scale != 1;
if (per_channel_quantization && !gemm.b_transpose) {
  return std::nullopt;
}

There is no validation that B’s quantization axis will map to TFLite FULLY_CONNECTED’s filter dimension 0 (output channels). If B has shape [N, K] and its scale has shape [1, K], SerializeQuantizeParams computes the quantization axis by scanning the scale shape for the non-unit dimension, arriving at axis 1:

// services/webnn/tflite/graph_builder_tflite.cc:7001-7009
for (size_t i = 0; i < scale_shape.size(); ++i) {
  if (scale_shape[i] != 1) {
    axis = (input_rank - scale_shape.size()) + i;   // yields axis = 1
  }
}

TFLite’s fully_connected::PrepareImpl validates only that the scale array length matches the filter size along the declared quantization dimension, which is a self-consistent but semantically wrong check:

// third_party/tflite/src/tensorflow/lite/kernels/fully_connected.cc:453-463
TF_LITE_ENSURE_EQ(context, affine_quantization->scale->size,
                  filter->dims->data[affine_quantization->quantized_dimension]);
// ...
data->per_channel_output_multiplier.resize(per_channel_quantization_size);
data->per_channel_output_shift.resize(per_channel_quantization_size);

For a filter of shape [1024, 2] with quantized_dimension = 1, the arrays are allocated with size 2. The optimized FullyConnectedPerChannel kernel then passes these 2-element arrays as per-row multipliers to the ruy GEMM backend, which expects 1024 entries (one per output row). The ruy kernel reads well past the allocated buffer.

Reproduce

Tested at commit 3ad31ba232d9a804b4de78d788e391f82b40a906. No source modifications required.

Build:

autoninja -C out/asan-release chrome

Windows

ASAN_OPTIONS=detect_odr_violation=0 out/asan-release/chrome.exe ^
  --no-sandbox ^
  --enable-features=WebMachineLearningNeuralNetwork ^
  --disable-features=WebNNOnnxRuntime ^
  --user-data-dir=%TEMP%\poc-tflite005 ^
  --enable-logging=stderr ^
  poc.html

Linux

out/asan-release/chrome \
  --no-sandbox \
  --enable-features=WebMachineLearningNeuralNetwork \
  --enable-logging=stderr \
  poc.html

The GPU process crashes within seconds. ASAN reports a heap-buffer-overflow READ of size 32 in ruy::Kernel8bitAvx2Impl, called from FullyConnectedPerChannel. The complete ASAN log is attached in issue_tflite005/asan.log.

==23132==ERROR: AddressSanitizer: heap-buffer-overflow on address 0x1254fc5da8af
READ of size 32 at 0x1254fc5da8af thread T23
    #0 ruy::Kernel8bitAvx2Impl<32>           kernel_avx2_fma.cc:359
    #5 GemmImplUsingRuy::Run                 cpu_backend_gemm_ruy.h:141
    #6 optimized_integer_ops::FullyConnectedPerChannel  fully_connected.h:95
    #7 fully_connected::EvalQuantized<1>     fully_connected.cc:1510
    #8 fully_connected::Eval<1>              fully_connected.cc:1795
    #12 webnn::tflite::GraphImplTflite::ComputeResources::DoDispatch  graph_impl_tflite.cc:328

0x1254fc5da8af is located 1 bytes before 8-byte region
allocated by thread T23 here:
    #1 std::vector<int>::resize              vector.h:1362
    #2 fully_connected::PrepareImpl          fully_connected.cc:463

SUMMARY: AddressSanitizer: heap-buffer-overflow ruy::Kernel8bitAvx2Impl

Credit

Please use c6eed09fc8b174b0f3eebedcceb1e792 as the credit for this vulnerability. Thank you.