From d9a15ab395a6548cee1b3f4d9fba6841e043fbe5 Mon Sep 17 00:00:00 2001
From: Bao Nguyen <waycaosi1994@gmail.com>
Date: Thu, 13 Aug 2026 16:49:23 +0700
Subject: [PATCH] fix: restore input buffer size validation lost in the Nitro
 migration

TfLiteTensorCopyFromBuffer returns kTfLiteError and leaves the tensor
untouched when the supplied byte size does not match the tensor's
expected size. copyInputBuffers discarded that status, so a wrong-sized
input ArrayBuffer was silently dropped and inference ran on whatever the
tensor held from the previous call, returning plausible-looking but stale
results instead of throwing.

v2 threw "Input Buffer size (N) does not match the Input Tensor's expected
size (M)!" from cpp/TensorHelpers.cpp; that check was removed when
cpp/TensorHelpers.cpp was deleted in 6f86153 (#172, the Nitro migration).

Check the status and throw with the actual and expected byte sizes, and
add harness coverage asserting a wrong-sized buffer throws instead of
reusing the previous input.
---
 cpp/HybridTfliteModel.cpp           | 12 ++++++-
 example/__tests__/tflite.harness.ts | 55 +++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/cpp/HybridTfliteModel.cpp b/cpp/HybridTfliteModel.cpp
index 1a7a3c41..78f0d42c 100644
--- a/cpp/HybridTfliteModel.cpp
+++ b/cpp/HybridTfliteModel.cpp
@@ -95,7 +95,17 @@ void HybridTfliteModel::copyInputBuffers(const std::vector<std::shared_ptr<Array
   for (int32_t i = 0; i < inputCount; i++) {
     TfLiteTensor* tensor = TfLiteInterpreterGetInputTensor(_interpreter, i);
     const std::shared_ptr<ArrayBuffer>& buffer = input[i];
-    TfLiteTensorCopyFromBuffer(tensor, buffer->data(), buffer->size());
+    TfLiteStatus status = TfLiteTensorCopyFromBuffer(tensor, buffer->data(), buffer->size());
+    if (status != kTfLiteOk) [[unlikely]] {
+      // TfLiteTensorCopyFromBuffer requires input_data_size == TfLiteTensorByteSize(tensor).
+      // On mismatch it leaves the tensor untouched, so without this check inference would
+      // silently run on the previous (or zero-initialized) contents of the tensor.
+      throw std::runtime_error("TFLite: Input buffer " + std::to_string(i) + " size (" +
+                               std::to_string(buffer->size()) + ") does not match input tensor \"" +
+                               std::string(TfLiteTensorName(tensor)) + "\" expected size (" +
+                               std::to_string(TfLiteTensorByteSize(tensor)) +
+                               ")! Status: " + tfLiteStatusToString(status));
+    }
   }
 }
 
diff --git a/example/__tests__/tflite.harness.ts b/example/__tests__/tflite.harness.ts
index d0c6854e..6ecf8ec7 100644
--- a/example/__tests__/tflite.harness.ts
+++ b/example/__tests__/tflite.harness.ts
@@ -69,6 +69,11 @@ function filledInputBuffer(input: Tensor, byte: number): ArrayBuffer {
   return buf;
 }
 
+/** Output buffers are pre-allocated per tensor and reused, so snapshot before re-running. */
+function copyOf(buffer: ArrayBuffer): ArrayBuffer {
+  return new Uint8Array(new Uint8Array(buffer)).buffer as ArrayBuffer;
+}
+
 function buffersEqual(a: ArrayBuffer, b: ArrayBuffer): boolean {
   if (a.byteLength !== b.byteLength) {
     return false;
@@ -405,5 +410,55 @@ describe('react-native-fast-tflite (harness)', () => {
       expect(top5Indices).toContain(GIANT_PANDA_INDEX);
     });
 
+    it('runSync rejects an input buffer that is too small', () => {
+      const input = firstInputTensor(model);
+      const tooSmall = new ArrayBuffer(tensorByteLength(input) - 1);
+      expect(() => model.runSync([tooSmall])).toThrow(
+        /does not match input tensor/i,
+      );
+    });
+
+    it('runSync rejects an input buffer that is too large', () => {
+      const input = firstInputTensor(model);
+      const tooLarge = new ArrayBuffer(tensorByteLength(input) + 1);
+      expect(() => model.runSync([tooLarge])).toThrow(
+        /does not match input tensor/i,
+      );
+    });
+
+    // Regression guard: a wrong-sized buffer used to be dropped silently, so
+    // inference ran on whatever was left in the tensor from the previous call
+    // and returned plausible-looking (but stale) scores instead of throwing.
+    it('does not silently run inference on the previous input when the buffer size is wrong', () => {
+      const input = firstInputTensor(model);
+      const byteLength = tensorByteLength(input);
+
+      // Prime the input tensor with a known image and keep a copy of its scores.
+      const primed = copyOf(model.runSync([filledInputBuffer(input, 0x10)])[0]!);
+
+      // A visibly different image, but one byte short.
+      const wrongSized = new Uint8Array(byteLength - 1).fill(0xf0)
+        .buffer as ArrayBuffer;
+
+      let threw = false;
+      let reusedStaleInput = false;
+      try {
+        const outputs = model.runSync([wrongSized]);
+        reusedStaleInput = buffersEqual(copyOf(outputs[0]!), primed);
+      } catch {
+        threw = true;
+      }
+
+      expect(reusedStaleInput).toBe(false);
+      expect(threw).toBe(true);
+
+      // The rejection must be specific to the size mismatch: a correctly sized
+      // buffer with the same content still runs and reaches the tensor, so the
+      // scores differ from the primed ones.
+      const correctlySized = filledInputBuffer(input, 0xf0);
+      const after = copyOf(model.runSync([correctlySized])[0]!);
+      expect(after.byteLength).toBe(tensorByteLength(model.outputs[0]!));
+      expect(buffersEqual(after, primed)).toBe(false);
+    });
   });
 });