From 18bd802a974cd434c1d06268f512e9d892e010f6 Mon Sep 17 00:00:00 2001 From: Tamim Abu Ayash Date: Tue, 15 Sep 2026 10:11:30 +0300 Subject: [PATCH] Validate TFLite inputs before copying tensor data --- .../servables/tensorflow/tflite_session.cc | 51 ++++++- .../tensorflow/tflite_session_test.cc | 139 +++++++++++++++++- 2 files changed, 187 insertions(+), 3 deletions(-) diff --git a/tensorflow_serving/servables/tensorflow/tflite_session.cc b/tensorflow_serving/servables/tensorflow/tflite_session.cc index 3bacaaf9737..c790eb6d112 100644 --- a/tensorflow_serving/servables/tensorflow/tflite_session.cc +++ b/tensorflow_serving/servables/tensorflow/tflite_session.cc @@ -28,6 +28,7 @@ limitations under the License. #include "tensorflow/cc/saved_model/signature_constants.h" #include "tensorflow/core/framework/tensor_shape.pb.h" #include "tensorflow/core/framework/tensor_util.h" +#include "tensorflow/core/framework/types.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/notification.h" #include "tensorflow/core/lib/gtl/cleanup.h" @@ -201,6 +202,17 @@ absl::Status SetInputAndInvokeMiniBatch( int tflite_input_idx = tflite_input_indices[i]; auto tflite_input_tensor = interpreter->tensor(tflite_input_idx); const auto& tf_input_tensors = inputs[i]; + DataType expected_tf_type; + TF_RETURN_IF_ERROR( + TfLiteTypeToTfType(tflite_input_tensor->type, &expected_tf_type)); + for (const Tensor* tf_input_tensor : tf_input_tensors) { + if (tf_input_tensor->dtype() != expected_tf_type) { + return absl::InvalidArgumentError(absl::StrCat( + "Expected input '", tflite_input_tensor->name, "' to have type ", + DataTypeString(expected_tf_type), ", but got ", + DataTypeString(tf_input_tensor->dtype()), ".")); + } + } if (tflite_input_tensor->type != kTfLiteString) { const Tensor* tf_input_tensor = tf_input_tensors[0]; // concated.tensor_data() may be accessed later. @@ -231,6 +243,13 @@ absl::Status SetInputAndInvokeMiniBatch( if (interpreter->AllocateTensors() != kTfLiteOk) { return absl::InternalError("Failed to allocate tensors"); } + tflite_input_tensor = interpreter->tensor(tflite_input_idx); + } + if (tensor_bytes.size() != tflite_input_tensor->bytes) { + return absl::InternalError(absl::StrCat( + "Input tensor byte size mismatch for '", tflite_input_tensor->name, + "': source has ", tensor_bytes.size(), " bytes, destination has ", + tflite_input_tensor->bytes, " bytes.")); } std::memcpy(tflite_input_tensor->data.raw, tensor_bytes.data(), tensor_bytes.size()); @@ -246,6 +265,7 @@ absl::Status SetInputAndInvokeMiniBatch( if (interpreter->AllocateTensors() != kTfLiteOk) { return absl::InternalError("Failed to allocate tensors"); } + tflite_input_tensor = interpreter->tensor(tflite_input_idx); } if (fixed_batch_size) { *fixed_batch_size = interpreter_wrapper->GetBatchSize(); @@ -663,7 +683,36 @@ absl::Status MergeInputTensors( return absl::InternalError(absl::StrCat( "Batch size expected to be positive; was ", batch.num_tasks())); } - const int tensors_per_task = batch.task(0).inputs.size(); + const TfLiteBatchTask& reference_task = batch.task(0); + const int tensors_per_task = reference_task.inputs.size(); + if (reference_task.input_indices.size() != tensors_per_task) { + return absl::InvalidArgumentError( + "TFLite batch task input tensors and indices are not aligned."); + } + if (reference_task.output_tensor_names == nullptr) { + return absl::InvalidArgumentError( + "TFLite batch task is missing output tensor names."); + } + for (int i = 1; i < batch.num_tasks(); ++i) { + const TfLiteBatchTask& task = batch.task(i); + if (task.inputs.size() != tensors_per_task) { + return absl::InvalidArgumentError( + "TFLite batch tasks have different input tensor counts."); + } + if (task.input_indices.size() != task.inputs.size()) { + return absl::InvalidArgumentError( + "TFLite batch task input tensors and indices are not aligned."); + } + if (task.input_indices != reference_task.input_indices) { + return absl::InvalidArgumentError( + "TFLite batch tasks have different input tensor indices."); + } + if (task.output_tensor_names == nullptr || + *task.output_tensor_names != *reference_task.output_tensor_names) { + return absl::InvalidArgumentError( + "TFLite batch tasks have different output tensor names."); + } + } *batch_size = 0; // each entry in merged_inputs is a list of task tensors. for (int i = 0; i < tensors_per_task; ++i) { diff --git a/tensorflow_serving/servables/tensorflow/tflite_session_test.cc b/tensorflow_serving/servables/tensorflow/tflite_session_test.cc index 6d2b6fea3bd..a801a4f6db8 100644 --- a/tensorflow_serving/servables/tensorflow/tflite_session_test.cc +++ b/tensorflow_serving/servables/tensorflow/tflite_session_test.cc @@ -63,6 +63,7 @@ namespace serving { namespace { using ::testing::_; +using ::testing::HasSubstr; using ::testing::Pair; using ::testing::SizeIs; using ::testing::UnorderedElementsAre; @@ -141,6 +142,39 @@ TEST(TfLiteSession, BasicTest) { } } +TEST(TfLiteSession, RejectsMismatchedInputTypes) { + std::string model_bytes; + TF_ASSERT_OK(ReadFileToString(tensorflow::Env::Default(), + test_util::TestSrcDirPath(kTestModel), + &model_bytes)); + + ::google::protobuf::Map signatures; + std::unique_ptr session; + tensorflow::SessionOptions options; + TF_ASSERT_OK(TfLiteSession::Create( + std::move(model_bytes), options, absl::GetFlag(FLAGS_num_pools), + absl::GetFlag(FLAGS_num_tflite_interpreters), &session, &signatures)); + + std::vector outputs; + Tensor int64_input(DT_INT64, TensorShape({1})); + absl::Status status = session->Run({{"x", int64_input}}, {"y"}, {}, &outputs); + EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(status.message(), HasSubstr("Expected input 'x'")); + + Tensor string_input = + test::AsTensor({"not a float"}, TensorShape({1})); + status = session->Run({{"x", string_input}}, {"y"}, {}, &outputs); + EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(status.message(), HasSubstr("Expected input 'x'")); + + Tensor valid_input = test::AsTensor({1.0}, TensorShape({1})); + outputs.clear(); + TF_EXPECT_OK(session->Run({{"x", valid_input}}, {"y"}, {}, &outputs)); + ASSERT_EQ(outputs.size(), 1); + test::ExpectTensorEqual( + outputs[0], test::AsTensor({2.5}, TensorShape({1}))); +} + TEST(TfLiteSession, ResizeWithSameNumElementsTest) { std::string model_bytes; TF_ASSERT_OK(ReadFileToString(tensorflow::Env::Default(), @@ -629,10 +663,18 @@ TEST(TfLiteSession, SimpleSignatureDefAndRun) { ASSERT_EQ(sigdef.outputs().at(kSignatureOutput).name(), kTestModelOutput); ASSERT_EQ(sigdef.method_name(), kClassifyMethodName); - Tensor input_list = - test::AsTensor({"a", "b", "c", "d"}, TensorShape({4})); + Tensor invalid_input_list = + test::AsTensor({1, 2, 3, 4}, TensorShape({4})); Tensor input_shape = test::AsTensor({2, 2}, TensorShape({2})); std::vector outputs; + absl::Status status = session->Run({{kTestModelInputList, invalid_input_list}, + {kTestModelInputShape, input_shape}}, + {kTestModelOutput}, {}, &outputs); + EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(status.message(), HasSubstr("to have type string")); + + Tensor input_list = + test::AsTensor({"a", "b", "c", "d"}, TensorShape({4})); TF_EXPECT_OK(session->Run( {{kTestModelInputList, input_list}, {kTestModelInputShape, input_shape}}, {kTestModelOutput}, {}, &outputs)); @@ -642,6 +684,88 @@ TEST(TfLiteSession, SimpleSignatureDefAndRun) { test::AsTensor({"a", "b", "c", "d"}, TensorShape({2, 2}))); } +absl::Status BuildTwoInputBatchedSession( + std::unique_ptr* session) { + auto model_signature_def_map = GetTestSignatureDefMap(); + std::string model_bytes = + BuildTestModel(tflite::TensorType_STRING, /*use_flex_op=*/false, + &model_signature_def_map); + ::google::protobuf::Map signatures; + tensorflow::SessionOptions options; + TF_RETURN_IF_ERROR(TfLiteSession::Create(std::move(model_bytes), options, 1, + 1, session, &signatures)); + + BasicBatchScheduler::Options scheduler_options; + scheduler_options.num_batch_threads = 1; + scheduler_options.max_batch_size = 2; + scheduler_options.batch_timeout_micros = 10 * 1000 * 1000; + return (*session)->SetScheduler( + TfLiteSession::CreateDefaultBasicBatchScheduler, scheduler_options); +} + +TEST(TfLiteSession, BatchedRequestsRejectDifferentInputCounts) { + std::unique_ptr session; + TF_ASSERT_OK(BuildTwoInputBatchedSession(&session)); + + Tensor input_list = test::AsTensor({"a"}, TensorShape({1})); + Tensor input_shape = test::AsTensor({1}, TensorShape({1})); + std::vector first_outputs; + std::vector second_outputs; + absl::Status first_status; + absl::Status second_status; + std::unique_ptr first_request_thread( + Env::Default()->StartThread(ThreadOptions(), "first_request", [&] { + first_status = session->Run({{kTestModelInputList, input_list}, + {kTestModelInputShape, input_shape}}, + {kTestModelOutput}, {}, &first_outputs); + })); + std::unique_ptr second_request_thread( + Env::Default()->StartThread(ThreadOptions(), "second_request", [&] { + second_status = session->Run({{kTestModelInputList, input_list}}, + {kTestModelOutput}, {}, &second_outputs); + })); + first_request_thread.reset(); + second_request_thread.reset(); + + EXPECT_EQ(first_status.code(), absl::StatusCode::kInvalidArgument); + EXPECT_EQ(second_status.code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(first_status.message(), + HasSubstr("different input tensor counts")); + EXPECT_THAT(second_status.message(), + HasSubstr("different input tensor counts")); +} + +TEST(TfLiteSession, BatchedRequestsRejectDifferentInputIndices) { + std::unique_ptr session; + TF_ASSERT_OK(BuildTwoInputBatchedSession(&session)); + + Tensor input_list = test::AsTensor({"a"}, TensorShape({1})); + Tensor input_shape = test::AsTensor({1}, TensorShape({1})); + std::vector first_outputs; + std::vector second_outputs; + absl::Status first_status; + absl::Status second_status; + std::unique_ptr first_request_thread( + Env::Default()->StartThread(ThreadOptions(), "first_request", [&] { + first_status = session->Run({{kTestModelInputList, input_list}}, + {kTestModelOutput}, {}, &first_outputs); + })); + std::unique_ptr second_request_thread( + Env::Default()->StartThread(ThreadOptions(), "second_request", [&] { + second_status = session->Run({{kTestModelInputShape, input_shape}}, + {kTestModelOutput}, {}, &second_outputs); + })); + first_request_thread.reset(); + second_request_thread.reset(); + + EXPECT_EQ(first_status.code(), absl::StatusCode::kInvalidArgument); + EXPECT_EQ(second_status.code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(first_status.message(), + HasSubstr("different input tensor indices")); + EXPECT_THAT(second_status.message(), + HasSubstr("different input tensor indices")); +} + absl::Status BuildSessionInBatch(std::unique_ptr* sess, bool use_model_batch_size, const std::string& model_path) { @@ -729,6 +853,17 @@ TEST_P(TfLiteSessionBatchSizeTest, TestBatchParallelismForFloat) { EXPECT_TRUE(outputs[0].shape().IsSameSize(TensorShape({kBatchSize, 1}))); } +TEST_P(TfLiteSessionBatchSizeTest, RejectsMismatchedInputType) { + std::unique_ptr sess; + TF_ASSERT_OK(BuildSessionInBatch(&sess, GetParam(), kTestModel)); + + Tensor invalid_input(DT_INT64, TensorShape({kBatchSize, 1})); + std::vector outputs; + absl::Status status = sess->Run({{"x", invalid_input}}, {"y"}, {}, &outputs); + EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(status.message(), HasSubstr("Expected input 'x'")); +} + TEST_P(TfLiteSessionBatchSizeTest, TestBatchParallelismForString) { std::unique_ptr sess; TF_ASSERT_OK(BuildSessionInBatch(&sess, GetParam(), kParseExampleModel));