diff --git a/tensorflow/lite/micro/memory_helpers.cc b/tensorflow/lite/micro/memory_helpers.cc index d78e34d4d96..f60dcc25d7b 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" @@ -111,7 +112,15 @@ TfLiteStatus BytesRequiredForTensor(const tflite::Tensor& flatbuffer_tensor, // 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); + int32_t dim = flatbuffer_tensor.shape()->Get(n); + if (dim < 0) { + return kTfLiteError; + } + if (element_count > 0 && + dim > std::numeric_limits::max() / element_count) { + return kTfLiteError; + } + element_count *= dim; } } diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc index ecb1651c71e..629a2a062a4 100644 --- a/tensorflow/lite/micro/micro_allocator.cc +++ b/tensorflow/lite/micro/micro_allocator.cc @@ -17,6 +17,8 @@ limitations under the License. #include #include +#include +#include #include "flatbuffers/flatbuffers.h" // from @flatbuffers #include "tensorflow/lite/c/common.h" @@ -287,6 +289,13 @@ TfLiteStatus InitializeTfLiteTensorFromFlatbuffer( // Populate per-channel quantization params. int channels = src_quantization->scale()->size(); + if (channels <= 0 || + static_cast(channels) > + (std::numeric_limits::max() - sizeof(TfLiteIntArray)) / + sizeof(int)) { + MicroPrintf("Invalid number of quantization channels: %d\n", channels); + return kTfLiteError; + } TfLiteAffineQuantization* quantization = allocate_temp ? reinterpret_cast( diff --git a/tensorflow/lite/micro/micro_allocator_test.cc b/tensorflow/lite/micro/micro_allocator_test.cc index 7a46bd0341c..d21ba95e196 100644 --- a/tensorflow/lite/micro/micro_allocator_test.cc +++ b/tensorflow/lite/micro/micro_allocator_test.cc @@ -1357,4 +1357,133 @@ TEST(MicroAllocatorTest, TestMultiSubgraphNumScratchAllocations) { used_bytes + sizeof(tflite::ScratchBufferHandle) * 2); } +const tflite::Tensor* CreateOverflowFlatbufferTensor( + flatbuffers::FlatBufferBuilder* builder, int dim1, int dim2) { + using flatbuffers::Offset; + const int32_t tensor_shape[] = {dim1, dim2}; + const Offset static_tensor_offset = + tflite::CreateTensor(*builder, builder->CreateVector(tensor_shape, 2), + tflite::TensorType_INT32, 0, + builder->CreateString("test_tensor"), 0, false); + builder->Finish(static_tensor_offset); + void* tensor_pointer = builder->GetBufferPointer(); + const tflite::Tensor* tensor = + flatbuffers::GetRoot(tensor_pointer); + return tensor; +} + +TEST(MicroAllocatorTest, TestTensorShapeIntegerOverflow) { + constexpr size_t arena_size = 1024; + uint8_t arena[arena_size]; + tflite::SingleArenaBufferAllocator* simple_allocator = + tflite::SingleArenaBufferAllocator::Create(arena, arena_size); + + flatbuffers::DefaultAllocator default_allocator; + flatbuffers::FlatBufferBuilder builder(1024, &default_allocator); + + const tflite::Tensor* tensor = + CreateOverflowFlatbufferTensor(&builder, 65536, 65536); + const flatbuffers::Vector>* buffers = + tflite::testing::CreateFlatbufferBuffers(); + + TfLiteTensor allocated_tensor; + EXPECT_EQ(kTfLiteError, + tflite::internal::InitializeTfLiteTensorFromFlatbuffer( + simple_allocator, simple_allocator, /*allocate_temp=*/false, + *tensor, buffers, &allocated_tensor)); + + simple_allocator->~SingleArenaBufferAllocator(); +} + +const tflite::Tensor* CreateMalformedQuantizedTensor( + flatbuffers::FlatBufferBuilder* builder, int size, uint32_t fake_channels) { + using flatbuffers::Offset; + + constexpr size_t quant_params_size = 1; + const float min_array[quant_params_size] = {0.1f}; + const float max_array[quant_params_size] = {0.2f}; + const float scale_array[quant_params_size] = {0.3f}; + const int64_t zero_point_array[quant_params_size] = {100ll}; + + const Offset quant_params = + tflite::CreateQuantizationParameters( + *builder, + /*min=*/builder->CreateVector(min_array, quant_params_size), + /*max=*/builder->CreateVector(max_array, quant_params_size), + /*scale=*/ + builder->CreateVector(scale_array, quant_params_size), + /*zero_point=*/ + builder->CreateVector(zero_point_array, quant_params_size)); + + constexpr size_t tensor_shape_size = 1; + const int32_t tensor_shape[tensor_shape_size] = {size}; + const Offset tensor_offset = tflite::CreateTensor( + *builder, builder->CreateVector(tensor_shape, tensor_shape_size), + tflite::TensorType_INT32, 0, builder->CreateString("test_tensor"), + quant_params, false); + builder->Finish(tensor_offset); + + uint8_t* tensor_pointer = builder->GetBufferPointer(); + size_t buffer_size = builder->GetSize(); + + // Mutate scale vector size + for (size_t i = 0; i < buffer_size - 8; i++) { + if (tensor_pointer[i] == 0x01 && tensor_pointer[i + 1] == 0x00 && + tensor_pointer[i + 2] == 0x00 && tensor_pointer[i + 3] == 0x00 && + tensor_pointer[i + 4] == 0x9a && tensor_pointer[i + 5] == 0x99 && + tensor_pointer[i + 6] == 0x99 && tensor_pointer[i + 7] == 0x3e) { + uint32_t* size_ptr = reinterpret_cast(&tensor_pointer[i]); + *size_ptr = fake_channels; + break; + } + } + + // Mutate zero_point vector size + for (size_t i = 0; i < buffer_size - 16; i++) { + if (tensor_pointer[i] == 0x01 && tensor_pointer[i + 1] == 0x00 && + tensor_pointer[i + 2] == 0x00 && tensor_pointer[i + 3] == 0x00 && + tensor_pointer[i + 8] == 0x64 && tensor_pointer[i + 9] == 0x00 && + tensor_pointer[i + 10] == 0x00 && tensor_pointer[i + 11] == 0x00 && + tensor_pointer[i + 12] == 0x00 && tensor_pointer[i + 13] == 0x00 && + tensor_pointer[i + 14] == 0x00 && tensor_pointer[i + 15] == 0x00) { + uint32_t* size_ptr = reinterpret_cast(&tensor_pointer[i]); + *size_ptr = fake_channels; + break; + } + } + + const tflite::Tensor* tensor = + flatbuffers::GetRoot(tensor_pointer); + return tensor; +} + +TEST(MicroAllocatorTest, TestQuantizationChannelsIntegerOverflow) { + // On a 32-bit target (sizeof(size_t) == 4), fake_channels = 0x3fffffff + // triggers an integer overflow in TfLiteIntArrayGetSizeInBytes (allocates 0 + // bytes), followed by a heap out-of-bounds write of zero points. We include + // this test to document the vulnerability. + if (sizeof(size_t) == 4) { + constexpr size_t arena_size = 1024; + uint8_t arena[arena_size]; + tflite::SingleArenaBufferAllocator* simple_allocator = + tflite::SingleArenaBufferAllocator::Create(arena, arena_size); + + flatbuffers::DefaultAllocator default_allocator; + flatbuffers::FlatBufferBuilder builder(1024, &default_allocator); + + const tflite::Tensor* tensor = + CreateMalformedQuantizedTensor(&builder, 100, 0x3fffffff); + const flatbuffers::Vector>* buffers = + tflite::testing::CreateFlatbufferBuffers(); + + TfLiteTensor allocated_tensor; + EXPECT_EQ(kTfLiteError, + tflite::internal::InitializeTfLiteTensorFromFlatbuffer( + simple_allocator, simple_allocator, /*allocate_temp=*/false, + *tensor, buffers, &allocated_tensor)); + + simple_allocator->~SingleArenaBufferAllocator(); + } +} + TF_LITE_MICRO_TESTS_MAIN