From b1dadabfedfeb33b2c6c202314003085f53a7c6c Mon Sep 17 00:00:00 2001 From: Github Executorch Date: Thu, 22 Jan 2026 13:18:26 -0800 Subject: [PATCH] Check size before memcpy --- extension/runner_util/inputs.cpp | 48 +++++- extension/runner_util/test/CMakeLists.txt | 9 +- extension/runner_util/test/inputs_test.cpp | 192 ++++++++++++++++++--- extension/runner_util/test/targets.bzl | 1 + test/models/export_program.py | 12 ++ test/models/targets.bzl | 1 + 6 files changed, 224 insertions(+), 39 deletions(-) diff --git a/extension/runner_util/inputs.cpp b/extension/runner_util/inputs.cpp index c1112489afb..646950e6cf8 100644 --- a/extension/runner_util/inputs.cpp +++ b/extension/runner_util/inputs.cpp @@ -86,17 +86,47 @@ Result prepare_input_tensors( Debug, "Verifying and setting input for non-tensor input %zu", i); if (tag.get() == Tag::Int) { - int64_t int_input; - std::memcpy(&int_input, buffer, buffer_size); - err = method.set_input(runtime::EValue(int_input), i); + if (buffer_size != sizeof(int64_t)) { + ET_LOG( + Error, + "Int input at index %zu has size %zu, expected sizeof(int64_t) %zu", + i, + buffer_size, + sizeof(int64_t)); + err = Error::InvalidArgument; + } else { + int64_t int_input; + std::memcpy(&int_input, buffer, buffer_size); + err = method.set_input(runtime::EValue(int_input), i); + } } else if (tag.get() == Tag::Double) { - double double_input; - std::memcpy(&double_input, buffer, buffer_size); - err = method.set_input(runtime::EValue(double_input), i); + if (buffer_size != sizeof(double)) { + ET_LOG( + Error, + "Double input at index %zu has size %zu, expected sizeof(double) %zu", + i, + buffer_size, + sizeof(double)); + err = Error::InvalidArgument; + } else { + double double_input; + std::memcpy(&double_input, buffer, buffer_size); + err = method.set_input(runtime::EValue(double_input), i); + } } else if (tag.get() == Tag::Bool) { - bool bool_input; - std::memcpy(&bool_input, buffer, buffer_size); - err = method.set_input(runtime::EValue(bool_input), i); + if (buffer_size != sizeof(bool)) { + ET_LOG( + Error, + "Bool input at index %zu has size %zu, expected sizeof(bool) %zu", + i, + buffer_size, + sizeof(bool)); + err = Error::InvalidArgument; + } else { + bool bool_input; + std::memcpy(&bool_input, buffer, buffer_size); + err = method.set_input(runtime::EValue(bool_input), i); + } } else { ET_LOG( Error, diff --git a/extension/runner_util/test/CMakeLists.txt b/extension/runner_util/test/CMakeLists.txt index 44b85a7fced..4f5c39aca65 100644 --- a/extension/runner_util/test/CMakeLists.txt +++ b/extension/runner_util/test/CMakeLists.txt @@ -19,17 +19,22 @@ include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) add_custom_command( OUTPUT "${CMAKE_CURRENT_BINARY_DIR}/ModuleAdd.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModuleIntBool.pte" COMMAND ${PYTHON_EXECUTABLE} -m test.models.export_program --modules - "ModuleAdd" --outdir "${CMAKE_CURRENT_BINARY_DIR}" + "ModuleAdd,ModuleIntBool" --outdir "${CMAKE_CURRENT_BINARY_DIR}" WORKING_DIRECTORY ${EXECUTORCH_ROOT} ) add_custom_target( executorch_runner_util_test_resources DEPENDS "${CMAKE_CURRENT_BINARY_DIR}/ModuleAdd.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModuleIntBool.pte" ) -set(test_env "ET_MODULE_ADD_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleAdd.pte") +set(test_env + "ET_MODULE_ADD_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleAdd.pte" + "ET_MODULE_INTBOOL_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleIntBool.pte" +) set(_test_srcs inputs_test.cpp) diff --git a/extension/runner_util/test/inputs_test.cpp b/extension/runner_util/test/inputs_test.cpp index aa3af2e145b..9720a0b9234 100644 --- a/extension/runner_util/test/inputs_test.cpp +++ b/extension/runner_util/test/inputs_test.cpp @@ -8,6 +8,9 @@ #include +#include +#include + #include #include #include @@ -40,52 +43,81 @@ class InputsTest : public ::testing::Test { void SetUp() override { torch::executor::runtime_init(); - // Create a loader for the serialized ModuleAdd program. - const char* path = std::getenv("ET_MODULE_ADD_PATH"); - Result loader = FileDataLoader::from(path); - ASSERT_EQ(loader.error(), Error::Ok); - loader_ = std::make_unique(std::move(loader.get())); + // Load ModuleAdd + const char* add_path = std::getenv("ET_MODULE_ADD_PATH"); + ASSERT_NE(add_path, nullptr) + << "ET_MODULE_ADD_PATH environment variable must be set"; + Result add_loader = FileDataLoader::from(add_path); + ASSERT_EQ(add_loader.error(), Error::Ok); + add_loader_ = std::make_unique(std::move(add_loader.get())); + + Result add_program = Program::load( + add_loader_.get(), Program::Verification::InternalConsistency); + ASSERT_EQ(add_program.error(), Error::Ok); + add_program_ = std::make_unique(std::move(add_program.get())); + + add_mmm_ = std::make_unique( + /*planned_memory_bytes=*/32 * 1024U, + /*method_allocator_bytes=*/32 * 1024U); + + Result add_method = + add_program_->load_method("forward", &add_mmm_->get()); + ASSERT_EQ(add_method.error(), Error::Ok); + add_method_ = std::make_unique(std::move(add_method.get())); + + // Load ModuleIntBool + const char* intbool_path = std::getenv("ET_MODULE_INTBOOL_PATH"); + ASSERT_NE(intbool_path, nullptr) + << "ET_MODULE_INTBOOL_PATH environment variable must be set"; + Result intbool_loader = FileDataLoader::from(intbool_path); + ASSERT_EQ(intbool_loader.error(), Error::Ok); + intbool_loader_ = + std::make_unique(std::move(intbool_loader.get())); - // Use it to load the program. - Result program = Program::load( - loader_.get(), Program::Verification::InternalConsistency); - ASSERT_EQ(program.error(), Error::Ok); - program_ = std::make_unique(std::move(program.get())); + Result intbool_program = Program::load( + intbool_loader_.get(), Program::Verification::InternalConsistency); + ASSERT_EQ(intbool_program.error(), Error::Ok); + intbool_program_ = + std::make_unique(std::move(intbool_program.get())); - mmm_ = std::make_unique( + intbool_mmm_ = std::make_unique( /*planned_memory_bytes=*/32 * 1024U, /*method_allocator_bytes=*/32 * 1024U); - // Load the forward method. - Result method = program_->load_method("forward", &mmm_->get()); - ASSERT_EQ(method.error(), Error::Ok); - method_ = std::make_unique(std::move(method.get())); + Result intbool_method = + intbool_program_->load_method("forward", &intbool_mmm_->get()); + ASSERT_EQ(intbool_method.error(), Error::Ok); + intbool_method_ = std::make_unique(std::move(intbool_method.get())); } private: - // Must outlive method_, but tests shouldn't need to touch them. - std::unique_ptr loader_; - std::unique_ptr mmm_; - std::unique_ptr program_; + std::unique_ptr add_loader_; + std::unique_ptr add_program_; + std::unique_ptr add_mmm_; + + std::unique_ptr intbool_loader_; + std::unique_ptr intbool_program_; + std::unique_ptr intbool_mmm_; protected: - std::unique_ptr method_; + std::unique_ptr add_method_; + std::unique_ptr intbool_method_; }; TEST_F(InputsTest, Smoke) { - Result input_buffers = prepare_input_tensors(*method_); + Result input_buffers = prepare_input_tensors(*add_method_); ASSERT_EQ(input_buffers.error(), Error::Ok); - auto input_err = method_->set_input(executorch::runtime::EValue(1.0), 2); + auto input_err = add_method_->set_input(executorch::runtime::EValue(1.0), 2); ASSERT_EQ(input_err, Error::Ok); // We can't look at the input tensors, but we can check that the outputs make // sense after executing the method. - Error status = method_->execute(); + Error status = add_method_->execute(); ASSERT_EQ(status, Error::Ok); // Get the single output, which should be a floating-point Tensor. - ASSERT_EQ(method_->outputs_size(), 1); - const EValue& output_value = method_->get_output(0); + ASSERT_EQ(add_method_->outputs_size(), 1); + const EValue& output_value = add_method_->get_output(0); ASSERT_EQ(output_value.tag, Tag::Tensor); Tensor output = output_value.toTensor(); ASSERT_EQ(output.scalar_type(), ScalarType::Float); @@ -107,14 +139,14 @@ TEST_F(InputsTest, ExceedingInputCountLimitFails) { // The smoke test above demonstrated that we can prepare inputs with the // default limits. It should fail if we lower the max below the number of // actual inputs. - MethodMeta method_meta = method_->method_meta(); + MethodMeta method_meta = add_method_->method_meta(); size_t num_inputs = method_meta.num_inputs(); ASSERT_GE(num_inputs, 1); executorch::extension::PrepareInputTensorsOptions options; options.max_inputs = num_inputs - 1; Result input_buffers = - prepare_input_tensors(*method_, options); + prepare_input_tensors(*add_method_, options); ASSERT_NE(input_buffers.error(), Error::Ok); } @@ -128,7 +160,7 @@ TEST_F(InputsTest, ExceedingInputAllocationLimitFails) { options.max_total_allocation_size = 1; Result input_buffers = - prepare_input_tensors(*method_, options); + prepare_input_tensors(*add_method_, options); ASSERT_NE(input_buffers.error(), Error::Ok); } @@ -186,3 +218,107 @@ TEST(BufferCleanupTest, Smoke) { // complaint. bc2.reset(); } + +TEST_F(InputsTest, DoubleInputWrongSizeFails) { + MethodMeta method_meta = add_method_->method_meta(); + + // ModuleAdd has 3 inputs: tensor, tensor, double (alpha) + ASSERT_EQ(method_meta.num_inputs(), 3); + + // Verify input 2 is a Double + auto tag = method_meta.input_tag(2); + ASSERT_TRUE(tag.ok()); + ASSERT_EQ(tag.get(), Tag::Double); + + // Create input_buffers with wrong size for the Double input + std::vector> input_buffers; + + // Allocate correct buffers for tensors (inputs 0 and 1) + auto tensor0_meta = method_meta.input_tensor_meta(0); + auto tensor1_meta = method_meta.input_tensor_meta(1); + ASSERT_TRUE(tensor0_meta.ok()); + ASSERT_TRUE(tensor1_meta.ok()); + + std::vector buf0(tensor0_meta->nbytes(), 0); + std::vector buf1(tensor1_meta->nbytes(), 0); + + // ModuleAdd expects alpha=1.0. Need to set this correctly, otherwise + // set_input fails validation before the buffer overflow happens. + double alpha = 1.0; + // Double is size 8; use a larger buffer to invoke overflow. + char large_buffer[16]; + memcpy(large_buffer, &alpha, sizeof(double)); + + input_buffers.push_back({buf0.data(), buf0.size()}); + input_buffers.push_back({buf1.data(), buf1.size()}); + input_buffers.push_back({large_buffer, sizeof(large_buffer)}); + + Result result = + prepare_input_tensors(*add_method_, {}, input_buffers); + EXPECT_EQ(result.error(), Error::InvalidArgument); +} + +TEST_F(InputsTest, IntBoolInputWrongSizeFails) { + MethodMeta method_meta = intbool_method_->method_meta(); + + // ModuleIntBool has 3 inputs: tensor, int, bool + ASSERT_EQ(method_meta.num_inputs(), 3); + + // Verify input types + auto int_tag = method_meta.input_tag(1); + ASSERT_TRUE(int_tag.ok()); + ASSERT_EQ(int_tag.get(), Tag::Int); + + auto bool_tag = method_meta.input_tag(2); + ASSERT_TRUE(bool_tag.ok()); + ASSERT_EQ(bool_tag.get(), Tag::Bool); + + // Allocate correct buffer for tensor (input 0) + auto tensor0_meta = method_meta.input_tensor_meta(0); + ASSERT_TRUE(tensor0_meta.ok()); + std::vector buf0(tensor0_meta->nbytes(), 0); + + // Prepare scalar values + int64_t y = 1; + bool z = true; + + // Test 1: Int input with wrong size + { + std::vector> input_buffers; + + // Int is size 8; use a larger buffer to invoke overflow. + char large_int_buffer[16]; + memcpy(large_int_buffer, &y, sizeof(int64_t)); + + char bool_buffer[sizeof(bool)]; + memcpy(bool_buffer, &z, sizeof(bool)); + + input_buffers.push_back({buf0.data(), buf0.size()}); + input_buffers.push_back({large_int_buffer, sizeof(large_int_buffer)}); + input_buffers.push_back({bool_buffer, sizeof(bool_buffer)}); + + Result result = + prepare_input_tensors(*intbool_method_, {}, input_buffers); + EXPECT_EQ(result.error(), Error::InvalidArgument); + } + + // Test 2: Bool input with wrong size + { + std::vector> input_buffers; + + char int_buffer[sizeof(int64_t)]; + memcpy(int_buffer, &y, sizeof(int64_t)); + + // Bool is size 1; use a larger buffer to invoke overflow. + char large_bool_buffer[8]; + memcpy(large_bool_buffer, &z, sizeof(bool)); + + input_buffers.push_back({buf0.data(), buf0.size()}); + input_buffers.push_back({int_buffer, sizeof(int_buffer)}); + input_buffers.push_back({large_bool_buffer, sizeof(large_bool_buffer)}); + + Result result = + prepare_input_tensors(*intbool_method_, {}, input_buffers); + EXPECT_EQ(result.error(), Error::InvalidArgument); + } +} diff --git a/extension/runner_util/test/targets.bzl b/extension/runner_util/test/targets.bzl index 95d5804ecdf..92746e3c9e3 100644 --- a/extension/runner_util/test/targets.bzl +++ b/extension/runner_util/test/targets.bzl @@ -28,5 +28,6 @@ def define_common_targets(is_fbcode = False): ], env = { "ET_MODULE_ADD_PATH": "$(location fbcode//executorch/test/models:exported_programs[ModuleAdd.pte])", + "ET_MODULE_INTBOOL_PATH": "$(location fbcode//executorch/test/models:exported_programs[ModuleIntBool.pte])", }, ) diff --git a/test/models/export_program.py b/test/models/export_program.py index ff5708f6685..2ee7d9b5e38 100644 --- a/test/models/export_program.py +++ b/test/models/export_program.py @@ -67,6 +67,18 @@ def get_random_inputs(self): return (torch.randn(10, 10, 10),) +# Used for testing int and bool inputs. +class ModuleIntBool(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x: torch.Tensor, y: int, z: bool): + return x + y + int(z) + + def get_random_inputs(self): + return (torch.ones(1), 1, True) + + class ModuleNoOp(nn.Module): def __init__(self): super(ModuleNoOp, self).__init__() diff --git a/test/models/targets.bzl b/test/models/targets.bzl index 2ae26d8217d..506d0a801a5 100644 --- a/test/models/targets.bzl +++ b/test/models/targets.bzl @@ -68,6 +68,7 @@ def define_common_targets(): "ModuleMultipleEntry", "ModuleNoKVCache", "ModuleIndex", + "ModuleIntBool", "ModuleDynamicCatUnallocatedIO", "ModuleSimpleTrain", "ModuleStateful",