diff --git a/signal/micro/kernels/fft_test.cc b/signal/micro/kernels/fft_test.cc index deda153ea53..4e6b84ad164 100644 --- a/signal/micro/kernels/fft_test.cc +++ b/signal/micro/kernels/fft_test.cc @@ -500,4 +500,20 @@ TEST(FftTest, FftAutoScaleTestLarge) { nullptr, 0, output, &scale_bit)); } +TEST(FftTest, RfftInputLengthGreaterThanFftLengthFails) { + constexpr int kOutputLen = 66; + int input_shape[] = {1, 128}; + float input[128] = {0}; + int output_shape[] = {1, kOutputLen}; + float golden[kOutputLen] = {0}; + float output[kOutputLen]; + const TFLMRegistration* registration = + tflite::tflm_signal::Register_RFFT_FLOAT(); + EXPECT_EQ(kTfLiteError, + tflite::testing::TestFFT( + input_shape, input, output_shape, golden, *registration, + g_gen_data_fft_length_64_float, + g_gen_data_size_fft_length_64_float, output, 1e-7)); +} + TF_LITE_MICRO_TESTS_MAIN diff --git a/signal/micro/kernels/rfft.cc b/signal/micro/kernels/rfft.cc index c9472b05657..24d00ba0ed2 100644 --- a/signal/micro/kernels/rfft.cc +++ b/signal/micro/kernels/rfft.cc @@ -101,6 +101,8 @@ TfLiteStatus RfftPrepare(TfLiteContext* context, TfLiteNode* node) { params->output_length = output_shape.Dims(output_shape.DimensionsCount() - 1) / 2; + TF_LITE_ENSURE(context, params->input_length <= params->fft_length); + context->RequestScratchBufferInArena(context, params->fft_length * sizeof(T), ¶ms->scratch_buffer_index); micro_context->DeallocateTempTfLiteTensor(input);