From fc580c84a097b9ac79702580fa93dce0439dd13e Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Fri, 24 Apr 2026 15:01:44 -0700 Subject: [PATCH] Numel and Nbytes Validation (#19121) Summary: Attempt 3 to check numel and nbytes overflow. This time we defer checking dynamic sized inputs until their size is realized. Reviewed By: lucylq Differential Revision: D98148157 --- runtime/executor/method.cpp | 28 ++++ runtime/executor/program_validation.cpp | 90 ++++++++----- runtime/executor/program_validation.h | 9 +- runtime/executor/test/method_test.cpp | 25 ++++ .../executor/test/program_validation_test.cpp | 126 ++++++++++++++---- 5 files changed, 214 insertions(+), 64 deletions(-) diff --git a/runtime/executor/method.cpp b/runtime/executor/method.cpp index 577691dc44b..1610804586d 100644 --- a/runtime/executor/method.cpp +++ b/runtime/executor/method.cpp @@ -9,6 +9,7 @@ #include #include +#include #include #include // @donotremove #include @@ -1194,6 +1195,33 @@ Method::set_input(const EValue& input_evalue, size_t input_idx) { input_idx, executorch::runtime::toString(t_dst.scalar_type()), executorch::runtime::toString(t_src.scalar_type())); + + ssize_t numel = 1; + for (ssize_t i = 0; i < t_src.dim(); i++) { + bool overflow = c10::mul_overflows( + numel, static_cast(t_src.size(i)), &numel); + ET_CHECK_OR_RETURN_ERROR( + !overflow, + InvalidArgument, + "Input %" ET_PRIsize_t + ": numel overflowed at dimension %zd with size %zd", + input_idx, + (size_t)i, + (size_t)t_src.size(i)); + } + size_t nbytes; + bool nbytes_overflow = c10::mul_overflows( + static_cast(numel), + executorch::runtime::elementSize(t_src.scalar_type()), + &nbytes); + ET_CHECK_OR_RETURN_ERROR( + !nbytes_overflow, + InvalidArgument, + "Input %" ET_PRIsize_t + ": nbytes overflowed: numel %zd with element size %zu", + input_idx, + numel, + executorch::runtime::elementSize(t_src.scalar_type())); // Reset the shape for the Method's input as the size of forwarded input // tensor for shape dynamism. Also is a safety check if need memcpy. ET_CHECK_OK_OR_RETURN_ERROR( diff --git a/runtime/executor/program_validation.cpp b/runtime/executor/program_validation.cpp index 92243fea289..448edc61ce5 100644 --- a/runtime/executor/program_validation.cpp +++ b/runtime/executor/program_validation.cpp @@ -14,7 +14,7 @@ #include #include -// #include +#include namespace executorch { namespace runtime { @@ -32,7 +32,8 @@ validate_tensor(const executorch_flatbuffer::Tensor* tensor) { return Error::InvalidProgram; } - // ssize_t numel = 1; + ssize_t numel = 1; + bool numel_overflowed = false; for (flatbuffers::uoffset_t i = 0; i < sizes->size(); i++) { int32_t size = sizes->Get(i); @@ -45,16 +46,10 @@ validate_tensor(const executorch_flatbuffer::Tensor* tensor) { return Error::InvalidProgram; } - // bool overflow = - // c10::mul_overflows(numel, static_cast(size), &numel); - // if (overflow) { - // ET_LOG( - // Error, - // "numel overflowed at dimension %u with size %d", - // static_cast(i), - // size); - // return Error::InvalidProgram; - // } + if (!numel_overflowed) { + numel_overflowed = + c10::mul_overflows(numel, static_cast(size), &numel); + } } auto scalar_type = @@ -64,19 +59,18 @@ validate_tensor(const executorch_flatbuffer::Tensor* tensor) { return Error::InvalidProgram; } - // size_t nbytes; - // bool nbytes_overflow = c10::mul_overflows( - // static_cast(numel), - // executorch::runtime::elementSize(scalar_type), - // &nbytes); - // if (nbytes_overflow) { - // ET_LOG( - // Error, - // "nbytes overflowed: numel %zd with element size %zu", - // numel, - // executorch::runtime::elementSize(scalar_type)); - // return Error::InvalidProgram; - // } + if (numel_overflowed) { + return Error::InvalidProgram; + } + + size_t nbytes; + bool nbytes_overflow = c10::mul_overflows( + static_cast(numel), + executorch::runtime::elementSize(scalar_type), + &nbytes); + if (nbytes_overflow) { + return Error::InvalidProgram; + } return Error::Ok; } @@ -114,6 +108,27 @@ validate_program(const executorch_flatbuffer::Program* program) { return Error::InvalidProgram; } + const auto* inputs = plan->inputs(); + auto is_dynamic_input = [&](flatbuffers::uoffset_t idx) -> bool { + if (inputs == nullptr) { + return false; + } + for (flatbuffers::uoffset_t i = 0; i < inputs->size(); i++) { + if (inputs->Get(i) == static_cast(idx)) { + const auto* value = values->Get(idx); + if (value == nullptr) { + return false; + } + const auto* tensor = + static_cast(value->val()); + return tensor != nullptr && + tensor->shape_dynamism() != + executorch_flatbuffer::TensorShapeDynamism::STATIC; + } + } + return false; + }; + for (flatbuffers::uoffset_t value_idx = 0; value_idx < values->size(); value_idx++) { const auto* value = values->Get(value_idx); @@ -128,12 +143,25 @@ validate_program(const executorch_flatbuffer::Program* program) { Error err = validate_tensor(tensor); if (err != Error::Ok) { - ET_LOG( - Error, - "Tensor validation failed for value %u in execution plan %u", - static_cast(value_idx), - static_cast(plan_idx)); - return err; + // Dynamic input tensors may have upper-bound sizes serialized for + // 64-bit machines that would overflow on 32-bit. Since their actual + // sizes are provided at set_input time, we defer overflow checks + // for those to Method::set_input. + if (is_dynamic_input(value_idx)) { + ET_LOG( + Info, + "Skipping validation failure for dynamic input tensor " + "at value %u in execution plan %u", + static_cast(value_idx), + static_cast(plan_idx)); + } else { + ET_LOG( + Error, + "Tensor validation failed for value %u in execution plan %u", + static_cast(value_idx), + static_cast(plan_idx)); + return err; + } } } diff --git a/runtime/executor/program_validation.h b/runtime/executor/program_validation.h index bb42a29423c..68e4ff7eb81 100644 --- a/runtime/executor/program_validation.h +++ b/runtime/executor/program_validation.h @@ -22,13 +22,12 @@ namespace executorch { namespace runtime { /** - * Validates that computing numel (number of elements) from the tensor's sizes - * will not overflow. This check should be performed before creating TensorImpl - * objects to prevent undefined behavior from integer overflow. + * Validates that a tensor's metadata is semantically valid: sizes are + * non-negative, scalar type is valid, and computing numel/nbytes will not + * overflow. * * @param[in] tensor The flatbuffer Tensor to validate. - * @return Error::Ok if the numel calculation is safe, Error::InvalidProgram - * if computing numel would overflow. + * @return Error::Ok if validation passes, Error::InvalidProgram otherwise. */ ET_NODISCARD Error validate_tensor(const executorch_flatbuffer::Tensor* tensor); diff --git a/runtime/executor/test/method_test.cpp b/runtime/executor/test/method_test.cpp index dc926184049..36a0c6f169b 100644 --- a/runtime/executor/test/method_test.cpp +++ b/runtime/executor/test/method_test.cpp @@ -310,6 +310,31 @@ TEST_F(MethodTest, AliasedIOTest) { } } +TEST_F(MethodTest, SetInputRejectsOverflowingSizes) { + // The "cat" model (ModuleDynamicCatUnallocatedIO) has a 2D input. + // set_input validates numel/nbytes overflow before resize_tensor. + ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes); + Result method = programs_["cat"]->load_method("forward", &mmm.get()); + ASSERT_EQ(method.error(), Error::Ok); + + // Create a 2D tensor with enormous sizes. On 32-bit platforms the numel + // multiplication overflows; on 64-bit the resize bounds check rejects it. + int32_t sizes[2] = {2000000000, 2000000000}; + uint8_t dim_order[2] = {0, 1}; + int32_t strides[2] = {1, 1}; + executorch::aten::TensorImpl impl( + executorch::aten::ScalarType::Float, + 2, + sizes, + nullptr, + dim_order, + strides); + + auto input_err = + method->set_input(EValue(executorch::aten::Tensor(&impl)), 0); + EXPECT_NE(input_err, Error::Ok); +} + TEST_F(MethodTest, ConstantSegmentTest) { // Execute model with constants stored in segment. ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes); diff --git a/runtime/executor/test/program_validation_test.cpp b/runtime/executor/test/program_validation_test.cpp index 73d3e5b9cb6..a133ff074ee 100644 --- a/runtime/executor/test/program_validation_test.cpp +++ b/runtime/executor/test/program_validation_test.cpp @@ -64,12 +64,15 @@ struct EValueConfig { EValueType type; std::vector tensor_sizes; // For Tensor type. std::vector tensor_list_items; // For TensorList type (indices). + bool is_dynamic = false; // For Tensor type: if true, uses DYNAMIC_BOUND. }; // Unified helper to create a minimal valid PTE flatbuffer with configurable -// evalues. Returns a buffer containing the flatbuffer data. +// evalues. Returns a buffer containing the flatbuffer data. input_indices +// specifies which value indices appear in the execution plan's inputs list. std::vector CreateTestProgram( - const std::vector& configs) { + const std::vector& configs, + const std::vector& input_indices = {}) { flatbuffers::FlatBufferBuilder builder(1024); std::vector> evalues; @@ -93,7 +96,9 @@ std::vector CreateTestProgram( /*data_buffer_idx=*/0, /*allocation_info=*/0, /*layout=*/0, - executorch_flatbuffer::TensorShapeDynamism::STATIC, + config.is_dynamic + ? executorch_flatbuffer::TensorShapeDynamism::DYNAMIC_BOUND + : executorch_flatbuffer::TensorShapeDynamism::STATIC, /*extra_tensor_info=*/0); evalues.push_back(executorch_flatbuffer::CreateEValue( builder, @@ -124,6 +129,7 @@ std::vector CreateTestProgram( auto values_vec = builder.CreateVector(evalues); auto plan_name = builder.CreateString("forward"); + auto inputs_vec = builder.CreateVector(input_indices); auto empty_int_vec = builder.CreateVector(std::vector{}); auto empty_int64_vec = builder.CreateVector(std::vector{0}); auto empty_chain_vec = builder.CreateVector( @@ -139,8 +145,8 @@ std::vector CreateTestProgram( plan_name, /*container_meta_type=*/0, values_vec, - empty_int_vec, - empty_int_vec, + /*inputs=*/inputs_vec, + /*outputs=*/empty_int_vec, empty_chain_vec, empty_operators_vec, empty_delegates_vec, @@ -206,29 +212,93 @@ TEST_F(ProgramValidationTest, InternalConsistencyDetectsTruncatedData) { ASSERT_EQ(program.error(), Error::InvalidProgram); } -// TEST_F(ProgramValidationTest, TensorNumelOverflowDetected) { -// std::vector configs = { -// {EValueType::Tensor, {2000000000, 2000000000, 2000000000}, {}}}; -// -// AlignedBuffer buf(CreateTestProgram(configs)); -// auto loader = buf.loader(); -// -// Result program = -// Program::load(&loader, Program::Verification::InternalConsistency); -// EXPECT_EQ(program.error(), Error::InvalidProgram); -// } - -// TEST_F(ProgramValidationTest, TensorNumelOverflowNotDetectedWithMinimal) { -// std::vector configs = { -// {EValueType::Tensor, {2000000000, 2000000000, 2000000000}, {}}}; -// -// AlignedBuffer buf(CreateTestProgram(configs)); -// auto loader = buf.loader(); -// -// // Minimal verification doesn't run program validation. -// Result program = -// Program::load(&loader, Program::Verification::Minimal); -// } +TEST_F(ProgramValidationTest, TensorNumelOverflowDetectedForStaticTensor) { + // Static tensors always have their overflow checked at validation time. + std::vector configs = { + {EValueType::Tensor, + {2000000000, 2000000000, 2000000000}, + {}, + /*is_dynamic=*/false}}; + + AlignedBuffer buf(CreateTestProgram(configs)); + auto loader = buf.loader(); + + Result program = + Program::load(&loader, Program::Verification::InternalConsistency); + EXPECT_EQ(program.error(), Error::InvalidProgram); +} + +TEST_F(ProgramValidationTest, TensorNumelOverflowNotDetectedWithMinimal) { + std::vector configs = { + {EValueType::Tensor, + {2000000000, 2000000000, 2000000000}, + {}, + /*is_dynamic=*/false}}; + + AlignedBuffer buf(CreateTestProgram(configs)); + auto loader = buf.loader(); + + // Minimal verification doesn't run program validation. + Result program = + Program::load(&loader, Program::Verification::Minimal); + EXPECT_EQ(program.error(), Error::Ok); +} + +TEST_F(ProgramValidationTest, TensorNumelOverflowSkippedForDynamicInput) { + // Dynamic input tensors skip overflow checks at validation time; the check + // is deferred to set_input where actual sizes are known. + std::vector configs = { + {EValueType::Tensor, + {2000000000, 2000000000, 2000000000}, + {}, + /*is_dynamic=*/true}}; + + // Mark value index 0 as a plan input. + AlignedBuffer buf(CreateTestProgram(configs, /*input_indices=*/{0})); + auto loader = buf.loader(); + + Result program = + Program::load(&loader, Program::Verification::InternalConsistency); + EXPECT_EQ(program.error(), Error::Ok); +} + +TEST_F( + ProgramValidationTest, + TensorNumelOverflowDetectedForDynamicNonInputTensor) { + // A dynamic tensor that is NOT in the inputs list should still have its + // overflow checked at validation time. + std::vector configs = { + {EValueType::Tensor, + {2000000000, 2000000000, 2000000000}, + {}, + /*is_dynamic=*/true}}; + + // No input indices — the tensor is not a plan input. + AlignedBuffer buf(CreateTestProgram(configs)); + auto loader = buf.loader(); + + Result program = + Program::load(&loader, Program::Verification::InternalConsistency); + EXPECT_EQ(program.error(), Error::InvalidProgram); +} + +TEST_F(ProgramValidationTest, TensorNumelOverflowDetectedForStaticInputTensor) { + // A static input tensor should still have its overflow checked at + // validation time since its sizes cannot change. + std::vector configs = { + {EValueType::Tensor, + {2000000000, 2000000000, 2000000000}, + {}, + /*is_dynamic=*/false}}; + + // Mark value index 0 as a plan input. + AlignedBuffer buf(CreateTestProgram(configs, /*input_indices=*/{0})); + auto loader = buf.loader(); + + Result program = + Program::load(&loader, Program::Verification::InternalConsistency); + EXPECT_EQ(program.error(), Error::InvalidProgram); +} TEST_F(ProgramValidationTest, NegativeSizeDetected) { std::vector configs = {{EValueType::Tensor, {10, -5, 10}, {}}};