diff --git a/tensorflow/lite/micro/memory_helpers.cc b/tensorflow/lite/micro/memory_helpers.cc index d78e34d4d96..82689697bd9 100644 --- a/tensorflow/lite/micro/memory_helpers.cc +++ b/tensorflow/lite/micro/memory_helpers.cc @@ -17,6 +17,7 @@ limitations under the License. #include #include +#include #include "flatbuffers/flatbuffers.h" // from @flatbuffers #include "tensorflow/lite/c/common.h" @@ -106,12 +107,29 @@ TfLiteStatus TfLiteTypeSizeOf(TfLiteType type, size_t* size) { TfLiteStatus BytesRequiredForTensor(const tflite::Tensor& flatbuffer_tensor, size_t* bytes, size_t* type_size) { - int element_count = 1; + // Shape dimensions and the element type size come from the partially-trusted + // flatbuffer model, so an oversized shape (e.g. [65536, 65536]) is an invalid + // model topology that -- per the Error Handling & Defensive Programming Guide + // -- must be rejected during the Setup phase (this runs under MicroAllocator). + // The running product is done in size_t (matching the result type) and every + // multiplication is guarded against wrap-around. This helper is not given a + // TfLiteContext, so the guide's low-cost `return kTfLiteError;` branch is used + // rather than TF_LITE_ENSURE. + size_t element_count = 1; // If flatbuffer_tensor.shape == nullptr, then flatbuffer_tensor is a scalar // so has 1 element. if (flatbuffer_tensor.shape() != nullptr) { for (size_t n = 0; n < flatbuffer_tensor.shape()->size(); ++n) { - element_count *= flatbuffer_tensor.shape()->Get(n); + const int32_t dim = flatbuffer_tensor.shape()->Get(n); + if (dim < 0) { + return kTfLiteError; + } + const size_t udim = static_cast(dim); + if (udim != 0 && + element_count > std::numeric_limits::max() / udim) { + return kTfLiteError; + } + element_count *= udim; } } @@ -119,6 +137,10 @@ TfLiteStatus BytesRequiredForTensor(const tflite::Tensor& flatbuffer_tensor, TF_LITE_ENSURE_STATUS( ConvertTensorType(flatbuffer_tensor.type(), &tf_lite_type)); TF_LITE_ENSURE_STATUS(TfLiteTypeSizeOf(tf_lite_type, type_size)); + if (*type_size != 0 && + element_count > std::numeric_limits::max() / *type_size) { + return kTfLiteError; + } *bytes = element_count * (*type_size); return kTfLiteOk; } @@ -127,15 +149,33 @@ TfLiteStatus TfLiteEvalTensorByteLength(const TfLiteEvalTensor* eval_tensor, size_t* out_bytes) { TFLITE_DCHECK(out_bytes != nullptr); - int element_count = 1; + // Called from both the Setup planning path and from Eval helpers (e.g. + // squeeze, kernel_util). The running product is computed in size_t and each + // multiplication is guarded against wrap-around. Because there is no + // TfLiteContext here and the guard protects against runtime memory + // corruption, the guide's low-cost `return kTfLiteError;` branch is used. + size_t element_count = 1; // If eval_tensor->dims == nullptr, then tensor is a scalar so has 1 element. if (eval_tensor->dims != nullptr) { for (int n = 0; n < eval_tensor->dims->size; ++n) { - element_count *= eval_tensor->dims->data[n]; + const int dim = eval_tensor->dims->data[n]; + if (dim < 0) { + return kTfLiteError; + } + const size_t udim = static_cast(dim); + if (udim != 0 && + element_count > std::numeric_limits::max() / udim) { + return kTfLiteError; + } + element_count *= udim; } } size_t type_size; TF_LITE_ENSURE_STATUS(TfLiteTypeSizeOf(eval_tensor->type, &type_size)); + if (type_size != 0 && + element_count > std::numeric_limits::max() / type_size) { + return kTfLiteError; + } *out_bytes = element_count * type_size; return kTfLiteOk; } @@ -151,17 +191,35 @@ TfLiteStatus AllocateOutputDimensionsFromInput(TfLiteContext* context, input = input1->dims->size > input2->dims->size ? input1 : input2; TF_LITE_ENSURE(context, output->type == input->type); - size_t size = 0; - TfLiteTypeSizeOf(input->type, &size); + // `bytes` is the tensor storage size (element size * shape product) that is + // stored in output->bytes. Keep it distinct from `dimensions_count`: the + // output->dims array is sized from the number of dimension entries, not from + // the byte count. + size_t bytes = 0; + TF_LITE_ENSURE_STATUS(TfLiteTypeSizeOf(input->type, &bytes)); const int dimensions_count = tflite::GetTensorShape(input).DimensionsCount(); for (int i = 0; i < dimensions_count; i++) { - size *= input->dims->data[i]; + const int dim = input->dims->data[i]; + const size_t udim = static_cast(dim); + // This helper runs in the Setup/Prepare phase, so per the Error Handling & + // Defensive Programming Guide a bad (partially-trusted model) shape is + // rejected with TF_LITE_ENSURE. The two preconditions -- a non-negative + // dimension and a running product that does not overflow size_t -- are + // combined into a single TF_LITE_ENSURE so only one error string is emitted + // into .rodata per call site (see "The Hidden Cost of TF_LITE_ENSURE"). + // The `dim >= 0` term is evaluated first, so the division guard is never + // reached with a wrapped-around `udim`. + TF_LITE_ENSURE( + context, + dim >= 0 && + (udim == 0 || bytes <= std::numeric_limits::max() / udim)); + bytes *= udim; } - output->bytes = size; + output->bytes = bytes; output->dims = reinterpret_cast(context->AllocatePersistentBuffer( - context, TfLiteIntArrayGetSizeInBytes(size))); + context, TfLiteIntArrayGetSizeInBytes(dimensions_count))); output->dims->size = input->dims->size; for (int i = 0; i < dimensions_count; i++) { output->dims->data[i] = input->dims->data[i]; diff --git a/tensorflow/lite/micro/memory_helpers_test.cc b/tensorflow/lite/micro/memory_helpers_test.cc index a004bb61bcd..c3f6b7213aa 100644 --- a/tensorflow/lite/micro/memory_helpers_test.cc +++ b/tensorflow/lite/micro/memory_helpers_test.cc @@ -34,6 +34,15 @@ void* FakeAllocatePersistentBuffer(TfLiteContext* context, size_t bytes) { return reinterpret_cast(global_persistent_buffer); } +// Records the size argument of the most recent AllocatePersistentBuffer call so +// a test can assert that output->dims is sized from the dimension count rather +// than from the (much larger) tensor byte count. +size_t g_last_requested_bytes = 0; +void* RecordingAllocatePersistentBuffer(TfLiteContext* context, size_t bytes) { + g_last_requested_bytes = bytes; + return reinterpret_cast(global_persistent_buffer); +} + } // namespace TEST(MemoryHelpersTest, TestAlignPointerUp) { @@ -187,6 +196,58 @@ TEST(MemoryHelpersTest, TestBytesRequiredForTensor) { EXPECT_EQ(static_cast(4), type_size); } +TEST(MemoryHelpersTest, + TfLiteEvalTensorByteLengthDoesNotTruncateAcrossInt32Boundary) { + // Shape [65536, 65536] with float32: 65536 * 65536 * 4 = 17179869184 bytes. + // The original implementation used a signed 32-bit running product, so the + // element count wrapped to 0 and *out_bytes was reported as 0 even though + // any subsequent allocation would address ~17 GiB. The hardened + // implementation performs the arithmetic in size_t and must report the + // mathematically correct value (where size_t is wide enough), or refuse the + // request via kTfLiteError when it would otherwise overflow size_t. + int dims[] = {2, 65536, 65536}; + TfLiteEvalTensor eval_tensor = {}; + eval_tensor.dims = tflite::testing::IntArrayFromInts(dims); + eval_tensor.type = kTfLiteFloat32; + + size_t out_bytes = 0; + const TfLiteStatus status = + tflite::TfLiteEvalTensorByteLength(&eval_tensor, &out_bytes); + + if (sizeof(size_t) >= 8) { + EXPECT_EQ(kTfLiteOk, status); + EXPECT_EQ(static_cast(17179869184ULL), out_bytes); + } else { + // 32-bit size_t cannot represent the result; the hardened code must + // refuse rather than silently truncate. + EXPECT_EQ(kTfLiteError, status); + } +} + +TEST(MemoryHelpersTest, TfLiteEvalTensorByteLengthRejectsSizeTOverflow) { + // A shape whose product overflows size_t even on 64-bit platforms must be + // rejected. INT32_MAX^4 * 4 vastly exceeds 2^64. + int dims[] = {4, 0x7fffffff, 0x7fffffff, 0x7fffffff, 0x7fffffff}; + TfLiteEvalTensor eval_tensor = {}; + eval_tensor.dims = tflite::testing::IntArrayFromInts(dims); + eval_tensor.type = kTfLiteFloat32; + + size_t out_bytes = 0; + EXPECT_EQ(kTfLiteError, + tflite::TfLiteEvalTensorByteLength(&eval_tensor, &out_bytes)); +} + +TEST(MemoryHelpersTest, TfLiteEvalTensorByteLengthRejectsNegativeDimension) { + int dims[] = {2, -1, 4}; + TfLiteEvalTensor eval_tensor = {}; + eval_tensor.dims = tflite::testing::IntArrayFromInts(dims); + eval_tensor.type = kTfLiteFloat32; + + size_t out_bytes = 0; + EXPECT_EQ(kTfLiteError, + tflite::TfLiteEvalTensorByteLength(&eval_tensor, &out_bytes)); +} + TEST(MemoryHelpersTest, TestAllocateOutputDimensionsFromInput) { constexpr int kDimsLen = 4; int input1_dims[] = {1, 1}; @@ -223,4 +284,34 @@ TEST(MemoryHelpersTest, TestAllocateOutputDimensionsFromInput) { } EXPECT_EQ(output_tensor.bytes, input_tensor2.bytes); } + +TEST(MemoryHelpersTest, AllocateOutputDimensionsSizesDimsByDimensionCount) { + // Regression test for the coupling between the byte count and the dimension + // count: output->dims must be allocated from the number of dimension entries, + // not from the tensor byte count. int32 [5, 5, 5, 5] is 2500 bytes but only 4 + // dimensions, so a correct implementation requests + // TfLiteIntArrayGetSizeInBytes(4), never TfLiteIntArrayGetSizeInBytes(2500). + constexpr int kDimsLen = 4; + int input1_dims[] = {1, 1}; + int input2_dims[] = {kDimsLen, 5, 5, 5, 5}; + int output_dims[] = {0, 0, 0, 0, 0}; + TfLiteTensor input_tensor1 = tflite::testing::CreateTensor( + nullptr, tflite::testing::IntArrayFromInts(input1_dims)); + TfLiteTensor input_tensor2 = tflite::testing::CreateTensor( + nullptr, tflite::testing::IntArrayFromInts(input2_dims)); + TfLiteTensor output_tensor = tflite::testing::CreateTensor( + nullptr, tflite::testing::IntArrayFromInts(output_dims)); + TfLiteContext context; + context.AllocatePersistentBuffer = RecordingAllocatePersistentBuffer; + + g_last_requested_bytes = 0; + EXPECT_EQ(kTfLiteOk, + tflite::AllocateOutputDimensionsFromInput( + &context, &input_tensor1, &input_tensor2, &output_tensor)); + + EXPECT_EQ(static_cast(TfLiteIntArrayGetSizeInBytes(kDimsLen)), + g_last_requested_bytes); + EXPECT_EQ(kDimsLen, output_tensor.dims->size); + EXPECT_EQ(input_tensor2.bytes, output_tensor.bytes); +} TF_LITE_MICRO_TESTS_MAIN