From 46adb06da38688145cbb4f8c808ef398b7547c81 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adil=20Burak=20=C5=9Een?= <56400880+adilburaksen@users.noreply.github.com> Date: Fri, 19 Jun 2026 01:57:34 +0300 Subject: [PATCH] Validate num_channels against tensor sizes in FILTER_BANK kernel FilterBankInit reads num_channels from the op's init flexbuffer and allocates a work area of (num_channels + 1) * sizeof(uint64_t); FilterbankAccumulateChannels then loops num_channels + 1 times, indexing channel_frequency_starts[i], channel_weight_starts[i] and channel_widths[i] and writing num_channels output elements. FilterBankPrepare validated only the rank and type of each tensor, never that num_channels is consistent with the per-channel metadata tensor lengths or the output length, so a model whose flexbuffer declares a num_channels larger than those tensors caused an out-of-bounds read. On targets where size_t is 32 bits, a large num_channels also overflows the work-area size computation into a tiny allocation that the accumulation loop then writes past (out-of-bounds write). Both were confirmed with AddressSanitizer. Reject an out-of-range num_channels in FilterBankInit, and validate in FilterBankPrepare that num_channels > 0, that the metadata tensors each have num_channels + 1 elements, and that the output has num_channels elements (mirroring the recent fix for the sibling FilterBankSpectralSubtraction kernel), plus a regression test. --- signal/micro/kernels/filter_bank.cc | 29 ++++++++++ signal/micro/kernels/filter_bank_test.cc | 69 ++++++++++++++++++++++++ 2 files changed, 98 insertions(+) diff --git a/signal/micro/kernels/filter_bank.cc b/signal/micro/kernels/filter_bank.cc index 1cf08d22ce3..0a014d57cfa 100644 --- a/signal/micro/kernels/filter_bank.cc +++ b/signal/micro/kernels/filter_bank.cc @@ -61,6 +61,16 @@ void* FilterBankInit(TfLiteContext* context, const char* buffer, length); params->config.num_channels = fbw.ElementAsInt32(kNumChannelsIndex); + // Reject a num_channels that is non-positive or large enough to overflow the + // work-area size computation below. Without this, a crafted model can wrap + // the multiplication (notably where size_t is 32 bits) into a tiny + // allocation, which FilterbankAccumulateChannels then writes past. + if (params->config.num_channels <= 0 || + static_cast(params->config.num_channels) > + SIZE_MAX / sizeof(uint64_t) - 1) { + return nullptr; + } + params->work_area = static_cast(context->AllocatePersistentBuffer( context, (params->config.num_channels + 1) * sizeof(uint64_t))); @@ -99,18 +109,21 @@ TfLiteStatus FilterBankPrepare(TfLiteContext* context, TfLiteNode* node) { TF_LITE_ENSURE(context, input != nullptr); TF_LITE_ENSURE_EQ(context, NumDimensions(input), 1); TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteInt16); + const int channel_frequency_starts_size = ElementCount(*input->dims); micro_context->DeallocateTempTfLiteTensor(input); input = micro_context->AllocateTempInputTensor(node, kChWeightStartsTensor); TF_LITE_ENSURE(context, input != nullptr); TF_LITE_ENSURE_EQ(context, NumDimensions(input), 1); TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteInt16); + const int channel_weight_starts_size = ElementCount(*input->dims); micro_context->DeallocateTempTfLiteTensor(input); input = micro_context->AllocateTempInputTensor(node, kChannelWidthsTensor); TF_LITE_ENSURE(context, input != nullptr); TF_LITE_ENSURE_EQ(context, NumDimensions(input), 1); TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteInt16); + const int channel_widths_size = ElementCount(*input->dims); micro_context->DeallocateTempTfLiteTensor(input); TfLiteTensor* output = @@ -118,8 +131,24 @@ TfLiteStatus FilterBankPrepare(TfLiteContext* context, TfLiteNode* node) { TF_LITE_ENSURE(context, output != nullptr); TF_LITE_ENSURE_EQ(context, NumDimensions(output), 1); TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteUInt64); + const int output_size = ElementCount(*output->dims); micro_context->DeallocateTempTfLiteTensor(output); + // Validate that num_channels (from the init flexbuffer) is consistent with + // the per-channel metadata tensors and the output, to prevent out-of-bounds + // access in FilterbankAccumulateChannels. That loop reads indices + // [0, num_channels] (num_channels + 1 elements) of channel_frequency_starts, + // channel_weight_starts and channel_widths, and writes num_channels output + // elements (work_area[1 .. num_channels]). + auto* params = reinterpret_cast(node->user_data); + TF_LITE_ENSURE(context, params != nullptr); + const int num_channels = params->config.num_channels; + TF_LITE_ENSURE(context, num_channels > 0); + TF_LITE_ENSURE_EQ(context, channel_frequency_starts_size, num_channels + 1); + TF_LITE_ENSURE_EQ(context, channel_weight_starts_size, num_channels + 1); + TF_LITE_ENSURE_EQ(context, channel_widths_size, num_channels + 1); + TF_LITE_ENSURE_EQ(context, output_size, num_channels); + return kTfLiteOk; } diff --git a/signal/micro/kernels/filter_bank_test.cc b/signal/micro/kernels/filter_bank_test.cc index 5203ad0b0bc..e08c7fa33d3 100644 --- a/signal/micro/kernels/filter_bank_test.cc +++ b/signal/micro/kernels/filter_bank_test.cc @@ -256,4 +256,73 @@ TEST(FilterBankTest, FilterBankTest16Channel) { g_gen_data_size_filter_bank_16_channel, output)); } +// PoC: 17-element (16-channel) side tensors but a flexbuffer declaring +// num_channels = 32. FilterBankPrepare never validates num_channels against +// the tensor element counts, so FilterbankAccumulateChannels loops +// num_channels + 1 = 33 times and reads channel_weight_starts[i] / +// channel_widths[i] past the end of the 17-element buffers (heap OOB read). +TEST(FilterBankTest, FilterBankNumChannelsOverflowPoc) { + int input1_shape[] = {1, 129}; + int input2_shape[] = {1, 59}; + int input3_shape[] = {1, 59}; + int input4_shape[] = {1, 17}; + int input5_shape[] = {1, 17}; + int input6_shape[] = {1, 17}; + int output_shape[] = {1, 16}; + + uint64_t output[16]; + + const uint32_t input1[] = { + 645050, 4644, 3653, 24262, 56660, 43260, 50584, 57902, 31702, 5401, + 45555, 34852, 8518, 43556, 13358, 19350, 40221, 18017, 27284, 64491, + 60099, 17863, 11001, 29076, 32666, 65268, 50947, 28694, 32377, 30014, + 25607, 22547, 45086, 10654, 46797, 8622, 47348, 43085, 5747, 51544, + 50364, 6208, 20696, 59782, 14429, 60125, 37079, 32673, 63457, 60142, + 34042, 11280, 1874, 33734, 62118, 13766, 54398, 47818, 50976, 46930, + 25906, 59441, 25958, 59136, 1756, 18652, 29213, 13379, 51845, 1207, + 55626, 27108, 43771, 35236, 3374, 40959, 47707, 41540, 34282, 27094, + 36329, 13593, 65257, 47006, 46857, 1114, 37106, 18738, 25969, 15461, + 2842, 36470, 32489, 61622, 23613, 29624, 32820, 30438, 9543, 6767, + 23037, 52896, 12059, 32264, 11575, 42400, 43344, 27511, 16712, 6877, + 4910, 50047, 61569, 57237, 48558, 2310, 22192, 7874, 46141, 64056, + 61997, 7298, 31372, 25316, 683, 58940, 18755, 17898, 19196}; + + const int16_t input2[] = { + -2210, 1711, 3237, 1247, 2507, 61, 1019, 899, 206, 146, 2849, 2756, + 1260, 1280, 1951, 213, 617, 2047, 211, 347, 2821, 3747, 150, 1924, + 3962, 942, 1430, 2678, 993, 308, 3364, 2491, 954, 1308, 879, 3950, + 1, 3556, 3628, 2104, 78, 1298, 1080, 342, 1337, 1639, 2352, 829, + 1358, 2498, 1647, 2507, 3816, 3767, 3735, 1155, 2221, 2196, 1160}; + + const int16_t input3[] = { + 408, 3574, 1880, 2561, 2011, 3394, 1019, 445, 3901, 343, 1874, 3846, + 3566, 1830, 327, 111, 623, 1037, 2803, 1947, 1518, 661, 3239, 2351, + 1257, 269, 1574, 3431, 3972, 2487, 2181, 1458, 552, 717, 679, 1031, + 1738, 1782, 128, 2242, 353, 1460, 3305, 1424, 3813, 2895, 164, 272, + 3886, 3135, 141, 747, 3233, 1478, 2612, 3837, 3271, 73, 1746}; + + const int16_t input4[] = {5, 6, 7, 9, 11, 12, 14, 16, 18, + 20, 22, 25, 27, 30, 32, 35, 33}; + + const int16_t input5[] = {0, 1, 2, 4, 6, 7, 9, 11, 13, + 15, 17, 20, 22, 25, 27, 30, 33}; + + const int16_t input6[] = {1, 1, 2, 2, 1, 2, 2, 2, 2, 2, 3, 2, 3, 2, 3, 3, 3}; + + const uint64_t golden[] = {104199304, 407748384, 206363744, 200989269, + 52144406, 230780884, 174394190, 379684049, + 94840835, 57788823, 531528204, 318265707, + 263149795, 188110467, 501443259, 200747781}; + + // With the fix in FilterBankPrepare, the num_channels/tensor-size mismatch is + // rejected before Eval, so no out-of-bounds access occurs. + EXPECT_EQ( + kTfLiteError, + tflite::testing::TestFilterBank( + input1_shape, input1, input2_shape, input2, input3_shape, input3, + input4_shape, input4, input5_shape, input5, input6_shape, input6, + output_shape, golden, g_gen_data_filter_bank_32_channel, + g_gen_data_size_filter_bank_32_channel, output)); +} + TF_LITE_MICRO_TESTS_MAIN