diff --git a/signal/micro/kernels/energy.cc b/signal/micro/kernels/energy.cc index 6a86366d42f..4d75f4b19a7 100644 --- a/signal/micro/kernels/energy.cc +++ b/signal/micro/kernels/energy.cc @@ -78,6 +78,15 @@ TfLiteStatus EnergyPrepare(TfLiteContext* context, TfLiteNode* node) { TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteInt16); TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteUInt32); + // Validate start_index/end_index against the input tensor size to prevent + // out-of-bounds access in SpectrumToEnergy (energy.cc uses Complex, + // so the element count in complex pairs is dims[0] / 2). + auto* params = reinterpret_cast(node->user_data); + const int input_elems = input->dims->data[0] / 2; // complex pairs + TF_LITE_ENSURE(context, params->start_index >= 0); + TF_LITE_ENSURE(context, params->end_index > params->start_index); + TF_LITE_ENSURE(context, params->end_index <= input_elems); + micro_context->DeallocateTempTfLiteTensor(input); micro_context->DeallocateTempTfLiteTensor(output); return kTfLiteOk; diff --git a/signal/micro/kernels/filter_bank_spectral_subtraction.cc b/signal/micro/kernels/filter_bank_spectral_subtraction.cc index e0693235dbd..ad0bcdef594 100644 --- a/signal/micro/kernels/filter_bank_spectral_subtraction.cc +++ b/signal/micro/kernels/filter_bank_spectral_subtraction.cc @@ -125,6 +125,17 @@ TfLiteStatus FilterBankSpectralSubtractionPrepare(TfLiteContext* context, auto* params = reinterpret_cast(node->user_data); + + // Validate that num_channels (from flexbuffer) matches all tensor sizes + // to prevent out-of-bounds access in FilterbankSpectralSubtraction. + const int input_elems = ElementCount(*input->dims); + const int output_elems = ElementCount(*output->dims); + const int noise_elems = ElementCount(*noise_estimate->dims); + TF_LITE_ENSURE(context, params->config.num_channels > 0); + TF_LITE_ENSURE_EQ(context, params->config.num_channels, input_elems); + TF_LITE_ENSURE_EQ(context, params->config.num_channels, output_elems); + TF_LITE_ENSURE_EQ(context, params->config.num_channels, noise_elems); + TfLiteTypeSizeOf(output->type, ¶ms->noise_estimate_size); params->noise_estimate_size *= ElementCount(*noise_estimate->dims); diff --git a/signal/micro/kernels/overlap_add.cc b/signal/micro/kernels/overlap_add.cc index 885ad3e4a2f..0e473801dfb 100644 --- a/signal/micro/kernels/overlap_add.cc +++ b/signal/micro/kernels/overlap_add.cc @@ -102,6 +102,11 @@ TfLiteStatus OverlapAddPrepare(TfLiteContext* context, TfLiteNode* node) { params->frame_size = input_shape.Dims(input_shape.DimensionsCount() - 1); params->n_frames = input_shape.Dims(input_shape.DimensionsCount() - 2); + // Validate dimension values and op param to prevent OOB/divide-by-zero. + TF_LITE_ENSURE(context, params->frame_size > 0); + TF_LITE_ENSURE(context, params->n_frames > 0); + TF_LITE_ENSURE(context, params->frame_step > 0); + TF_LITE_ENSURE(context, params->frame_step <= params->frame_size); params->outer_dims = input_shape.FlatSize() / (params->frame_size * params->n_frames); params->state_buffers = static_cast(context->AllocatePersistentBuffer( diff --git a/signal/micro/kernels/pcan.cc b/signal/micro/kernels/pcan.cc index 9473e1b4fab..323f410d6ba 100644 --- a/signal/micro/kernels/pcan.cc +++ b/signal/micro/kernels/pcan.cc @@ -81,6 +81,15 @@ TfLiteStatus PcanPrepare(TfLiteContext* context, TfLiteNode* node) { TF_LITE_ENSURE_TYPES_EQ(context, gain_lut->type, kTfLiteInt16); TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteUInt32); + // Validate that input/noise_estimate/output dimensions match to prevent + // out-of-bounds access in ApplyPcanAutoGainControlFixed. + const int input_elems = ElementCount(*input->dims); + const int noise_elems = ElementCount(*noise_estimate->dims); + const int output_elems = ElementCount(*output->dims); + TF_LITE_ENSURE(context, input_elems > 0); + TF_LITE_ENSURE_EQ(context, input_elems, noise_elems); + TF_LITE_ENSURE_EQ(context, input_elems, output_elems); + micro_context->DeallocateTempTfLiteTensor(input); micro_context->DeallocateTempTfLiteTensor(output); micro_context->DeallocateTempTfLiteTensor(noise_estimate); diff --git a/signal/micro/kernels/rfft.cc b/signal/micro/kernels/rfft.cc index c9472b05657..a7fb9730b46 100644 --- a/signal/micro/kernels/rfft.cc +++ b/signal/micro/kernels/rfft.cc @@ -97,6 +97,11 @@ TfLiteStatus RfftPrepare(TfLiteContext* context, TfLiteNode* node) { RuntimeShape output_shape = GetTensorShape(output); params->input_length = input_shape.Dims(input_shape.DimensionsCount() - 1); params->input_size = input_shape.FlatSize(); + // Validate that the input fits within the FFT scratch buffer to prevent + // heap-buffer-overflow in the Eval zero-padding memcpy/memset. + TF_LITE_ENSURE(context, params->fft_length > 0); + TF_LITE_ENSURE(context, params->input_length > 0); + TF_LITE_ENSURE(context, params->input_length <= params->fft_length); // Divide by 2 because output is complex. params->output_length = output_shape.Dims(output_shape.DimensionsCount() - 1) / 2; diff --git a/signal/micro/kernels/window.cc b/signal/micro/kernels/window.cc index cd9c4623768..028fefd0387 100644 --- a/signal/micro/kernels/window.cc +++ b/signal/micro/kernels/window.cc @@ -81,6 +81,15 @@ TfLiteStatus WindowPrepare(TfLiteContext* context, TfLiteNode* node) { RuntimeShape input_shape = GetTensorShape(input); params->input_size = input_shape.FlatSize(); + // Validate that input size is a multiple of the window (weights) size + // and that all buffers match, preventing OOB in ApplyWindow. + const int weight_size = weights->dims->data[0]; + TF_LITE_ENSURE(context, weight_size > 0); + TF_LITE_ENSURE(context, params->input_size > 0); + TF_LITE_ENSURE_EQ(context, params->input_size % weight_size, 0); + TF_LITE_ENSURE_EQ(context, params->input_size, + GetTensorShape(output).FlatSize()); + micro_context->DeallocateTempTfLiteTensor(input); micro_context->DeallocateTempTfLiteTensor(weights); micro_context->DeallocateTempTfLiteTensor(output);