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& 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); + }); }); });