From f1bc52b14e41cbe426afee639299d7c6afbd6b3d Mon Sep 17 00:00:00 2001
From: Esun Kim <veblush@google.com>
Date: Tue, 1 Sep 2026 11:03:46 -0700
Subject: [PATCH] Fortified rfft

---
 signal/micro/kernels/fft_test.cc              | 16 ++++++++++++++++
 signal/micro/kernels/rfft.cc                  |  2 ++
 tensorflow/lite/micro/testing/micro_test_v2.h |  8 +++++---
 3 files changed, 23 insertions(+), 3 deletions(-)

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<float>(
+                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),
                                        &params->scratch_buffer_index);
   micro_context->DeallocateTempTfLiteTensor(input);
diff --git a/tensorflow/lite/micro/testing/micro_test_v2.h b/tensorflow/lite/micro/testing/micro_test_v2.h
index e6fef1dc95c..0db2b4bc71b 100644
--- a/tensorflow/lite/micro/testing/micro_test_v2.h
+++ b/tensorflow/lite/micro/testing/micro_test_v2.h
@@ -147,12 +147,12 @@ inline bool AreStringsEqual(const char* s1, const char* s2) {
 // Singleton class to manage test registration and execution.
 class TestRunner {
  public:
-  static TestRunner& Get() {
+  static TestRunner& Get() noexcept {
     static TestRunner instance;
     return instance;
   }
 
-  void RegisterTest(TestInfo* test) {
+  void RegisterTest(TestInfo* test) noexcept {
     test->next = tests_;
     tests_ = test;
   }
@@ -229,7 +229,9 @@ class TestRunner {
 // Helper class to register tests at startup time.
 class TestRegistrar {
  public:
-  TestRegistrar(TestInfo* info) { TestRunner::Get().RegisterTest(info); }
+  explicit TestRegistrar(TestInfo* info) noexcept {
+    TestRunner::Get().RegisterTest(info);
+  }
 };
 
 }  // namespace internal