Validate STACKER init parameters against overflow and tensor sizes

signal/micro/kernels/stacker.cc

@@ -112,9 +112,36 @@ TfLiteStatus StackerPrepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteInt16);
   TF_LITE_ENSURE_TYPES_EQ(context, output_valid->type, kTfLiteBool);
 
+  const int input_size = ElementCount(*input->dims);
+  const int output_size = ElementCount(*output->dims);
+
   micro_context->DeallocateTempTfLiteTensor(input);
   micro_context->DeallocateTempTfLiteTensor(output);
   micro_context->DeallocateTempTfLiteTensor(output_valid);
+
+  // Validate the init-flexbuffer parameters before they are used to size the
+  // circular buffer and to copy data. Without this, negative or overflowing
+  // values produce a zero/undersized buffer that StackerEval then writes and
+  // reads out of bounds (CircularBufferWrite copies num_channels values per
+  // frame and the output receives buffer_size values).
+  auto* params = reinterpret_cast<TFLMSignalStackerParams*>(node->user_data);
+  TF_LITE_ENSURE(context, params != nullptr);
+  TF_LITE_ENSURE(context, params->num_channels > 0);
+  TF_LITE_ENSURE(context, params->stacker_left_context >= 0);
+  TF_LITE_ENSURE(context, params->stacker_right_context >= 0);
+  TF_LITE_ENSURE(context, params->stacker_step > 0);
+  TF_LITE_ENSURE(context, params->num_channels <= input_size);
+
+  // Recompute buffer_size in 64-bit to confirm the products did not overflow,
+  // then require the output to be large enough to receive it.
+  const int64_t frames = static_cast<int64_t>(params->stacker_left_context) +
+                         params->stacker_right_context + 1;
+  const int64_t buffer_size =
+      static_cast<int64_t>(params->num_channels) * frames;
+  TF_LITE_ENSURE(context,
+                 buffer_size == static_cast<int64_t>(params->buffer_size));
+  TF_LITE_ENSURE(context, buffer_size <= output_size);
+
   return kTfLiteOk;
 }
 

signal/micro/kernels/stacker_test.cc

@@ -238,4 +238,25 @@ TEST(StackerTest, StackerTestReset10ChannelStep2_2ndTest) {
                            g_gen_data_size_stacker_10_channels_step_2);
 }
 
+
+// Regression test: a crafted init flexbuffer with stacker_right_context = -1 makes
+// buffer_size = num_channels * (left + right + 1) = 0. Before the fix this produced a
+// zero-capacity circular buffer that StackerEval wrote/read out of bounds; now
+// StackerPrepare rejects the invalid parameters.
+TEST(StackerTest, StackerInvalidParamsRejected) {
+  int input_shape[] = {1, 8};
+  int output_shape[] = {1, 8};
+  int output_ready_shape[] = {0};
+  const int16_t input[8] = {1, 2, 3, 4, 5, 6, 7, 8};
+  int16_t output[8];
+  bool output_ready = false;
+  const unsigned char evil_flex[] = {110,117,109,95,99,104,97,110,110,101,108,115,0,115,116,97,99,107,101,114,95,108,101,102,116,95,99,111,110,116,101,120,116,0,115,116,97,99,107,101,114,95,114,105,103,104,116,95,99,111,110,116,101,120,116,0,115,116,97,99,107,101,114,95,115,116,101,112,0,4,70,58,38,17,4,1,4,8,0,255,1,4,4,4,4,8,36,1};
+  tflite::StackerKernelRunner stacker_runner(input_shape, input, output_shape,
+                                             output, output_ready_shape,
+                                             &output_ready);
+  EXPECT_EQ(kTfLiteError,
+            stacker_runner.kernel_runner()->InitAndPrepare(
+                reinterpret_cast<const char*>(evil_flex), sizeof(evil_flex)));
+}
+
 TF_LITE_MICRO_TESTS_MAIN