From 1ff8aeeb6b9f1c37ae75ed151813c71b48d86a51 Mon Sep 17 00:00:00 2001
From: nonoge <2359216537tjf@gmail.com>
Date: Sun, 5 Apr 2026 01:03:37 +0800
Subject: [PATCH] fix: reject negative EXP flat_size before size_t cast

The EXP kernel computed flat_size as a signed int and then cast it to size_t without validating negative values. A malformed tensor shape such as [-1] turns flat_size into a negative number, and the cast expands it into a huge unsigned length that drives reference_ops::Exp into out-of-bounds memory access.

This change adds a TF_LITE_ENSURE check in ExpEval so malformed negative flat sizes are rejected before the cast. It also adds a regression test that demonstrates the kernel now returns kTfLiteError for a negative dimension instead of crashing.
---
 tensorflow/lite/micro/kernels/exp.cc      |  1 +
 tensorflow/lite/micro/kernels/exp_test.cc | 28 +++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/tensorflow/lite/micro/kernels/exp.cc b/tensorflow/lite/micro/kernels/exp.cc
index 8d1da8faf00..5c41cdc3c88 100644
--- a/tensorflow/lite/micro/kernels/exp.cc
+++ b/tensorflow/lite/micro/kernels/exp.cc
@@ -58,6 +58,7 @@ TfLiteStatus ExpEval(TfLiteContext* context, TfLiteNode* node) {
       tflite::micro::GetEvalOutput(context, node, kOutputTensor);
   int flat_size = MatchingFlatSize(tflite::micro::GetTensorShape(input),
                                    tflite::micro::GetTensorShape(output));
+  TF_LITE_ENSURE(context, flat_size >= 0);
 
   if (input->type == kTfLiteFloat32) {
     reference_ops::Exp(tflite::micro::GetTensorData<float>(input),
diff --git a/tensorflow/lite/micro/kernels/exp_test.cc b/tensorflow/lite/micro/kernels/exp_test.cc
index 30e0e73ed15..39c0b0ab1f6 100644
--- a/tensorflow/lite/micro/kernels/exp_test.cc
+++ b/tensorflow/lite/micro/kernels/exp_test.cc
@@ -71,4 +71,32 @@ TEST(ExpTest, SingleDim) {
   tflite::testing::TestExp(input_dims, input_values, golden, output_data);
 }
 
+TEST(ExpTest, NegativeDimensionShallFail) {
+  int dims_data[] = {1, -1};
+  TfLiteIntArray* dims = tflite::testing::IntArrayFromInts(dims_data);
+  float input_data[1] = {0.0f};
+  float output_data[1] = {0.0f};
+
+  TfLiteTensor tensors[2] = {
+      tflite::testing::CreateTensor(input_data, dims),
+      tflite::testing::CreateTensor(output_data, dims),
+  };
+
+  int inputs_array_data[] = {1, 0};
+  TfLiteIntArray* inputs_array =
+      tflite::testing::IntArrayFromInts(inputs_array_data);
+  int outputs_array_data[] = {1, 1};
+  TfLiteIntArray* outputs_array =
+      tflite::testing::IntArrayFromInts(outputs_array_data);
+
+  const TFLMRegistration registration = tflite::Register_EXP();
+  tflite::micro::KernelRunner runner(registration, tensors,
+                                     /*tensors_size=*/2, inputs_array,
+                                     outputs_array,
+                                     /*builtin_data=*/nullptr);
+
+  EXPECT_EQ(kTfLiteOk, runner.InitAndPrepare());
+  EXPECT_EQ(kTfLiteError, runner.Invoke());
+}
+
 TF_LITE_MICRO_TESTS_MAIN