diff --git a/runtime/executor/tensor_parser_exec_aten.cpp b/runtime/executor/tensor_parser_exec_aten.cpp index 74fb602cdf4..c61be43503e 100644 --- a/runtime/executor/tensor_parser_exec_aten.cpp +++ b/runtime/executor/tensor_parser_exec_aten.cpp @@ -129,6 +129,12 @@ ET_NODISCARD Error validateTensorLayout( "Dim mismatch. Expected %d, got %zu.", dim, expected_layout.sizes().size()); + ET_CHECK_OR_RETURN_ERROR( + s_tensor->dim_order()->size() == static_cast(dim), + InvalidExternalData, + "Dim order size mismatch. Expected %d, got %zu.", + dim, + s_tensor->dim_order()->size()); for (int i = 0; i < dim; i++) { ET_CHECK_OR_RETURN_ERROR( s_tensor->sizes()->Get(i) == expected_layout.sizes()[i], diff --git a/runtime/executor/test/tensor_parser_test.cpp b/runtime/executor/test/tensor_parser_test.cpp index b193ab229b0..df82c2f4e72 100644 --- a/runtime/executor/test/tensor_parser_test.cpp +++ b/runtime/executor/test/tensor_parser_test.cpp @@ -10,6 +10,7 @@ #include #include +#include #include #include @@ -23,7 +24,10 @@ using executorch::runtime::EValue; using executorch::runtime::FreeableBuffer; using executorch::runtime::Program; using executorch::runtime::Result; +using executorch::runtime::Span; +using executorch::runtime::TensorLayout; using executorch::runtime::deserialization::parseTensor; +using executorch::runtime::deserialization::validateTensorLayout; using executorch::runtime::testing::ManagedMemoryManager; using torch::executor::util::FileDataLoader; @@ -188,3 +192,34 @@ TEST_F(TensorParserTest, TestMutableState) { } ASSERT_EQ(num_mutable_tensors, 2); } + +// Tests that validateTensorLayout rejects tensors where dim_order is shorter +// than sizes, preventing out-of-bounds reads. +TEST(ValidateTensorLayoutTest, DimOrderSizeMismatchIsRejected) { + flatbuffers::FlatBufferBuilder builder; + + std::vector sizes = {2, 3, 4}; + std::vector dim_order_short = {0}; + + auto tensor_offset = executorch_flatbuffer::CreateTensor( + builder, + executorch_flatbuffer::ScalarType::FLOAT, + 0, + builder.CreateVector(sizes), + builder.CreateVector(dim_order_short)); + builder.Finish(tensor_offset); + + const auto* s_tensor = flatbuffers::GetRoot( + builder.GetBufferPointer()); + + std::vector expected_sizes = {2, 3, 4}; + std::vector expected_dim_order = {0, 1, 2}; + auto layout = TensorLayout::create( + Span(expected_sizes.data(), expected_sizes.size()), + Span(expected_dim_order.data(), expected_dim_order.size()), + ScalarType::Float); + ASSERT_TRUE(layout.ok()); + + EXPECT_EQ( + validateTensorLayout(s_tensor, layout.get()), Error::InvalidExternalData); +}