diff --git a/runtime/executor/tensor_parser_exec_aten.cpp b/runtime/executor/tensor_parser_exec_aten.cpp index 31ec2377f16..1f2ee0e5565 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( "Scalar type mismatch. Expected %hhd, got %hhd.", static_cast(s_tensor->scalar_type()), static_cast(expected_layout.scalar_type())); + ET_CHECK_OR_RETURN_ERROR( + s_tensor->sizes() != nullptr, InvalidExternalData, "Missing sizes field"); + ET_CHECK_OR_RETURN_ERROR( + s_tensor->dim_order() != nullptr, + InvalidExternalData, + "Missing dim_order field"); int dim = s_tensor->sizes()->size(); ET_CHECK_OR_RETURN_ERROR( dim >= 0, InvalidExternalData, "Dim is negative: %d", dim) diff --git a/runtime/executor/test/tensor_parser_test.cpp b/runtime/executor/test/tensor_parser_test.cpp index bf102d7d1f6..1214d0ce731 100644 --- a/runtime/executor/test/tensor_parser_test.cpp +++ b/runtime/executor/test/tensor_parser_test.cpp @@ -229,6 +229,72 @@ TEST(ValidateTensorLayoutTest, DimOrderSizeMismatchIsRejected) { validateTensorLayout(s_tensor, layout.get()), Error::InvalidExternalData); } +// Tests that validateTensorLayout rejects tensors with a null sizes field +// instead of dereferencing it, which would SIGSEGV under the default +// Verification::Minimal load mode. +TEST(ValidateTensorLayoutTest, NullSizesIsRejected) { + flatbuffers::FlatBufferBuilder builder; + + std::vector dim_order = {0, 1, 2}; + + // Pass 0 for the sizes offset to serialize a null sizes field. + auto tensor_offset = executorch_flatbuffer::CreateTensor( + builder, + executorch_flatbuffer::ScalarType::FLOAT, + /*storage_offset=*/0, + /*sizes=*/0, + builder.CreateVector(dim_order)); + builder.Finish(tensor_offset); + + const auto* s_tensor = flatbuffers::GetRoot( + builder.GetBufferPointer()); + ASSERT_EQ(s_tensor->sizes(), nullptr); + + 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); +} + +// Tests that validateTensorLayout rejects tensors with a null dim_order field +// instead of dereferencing it, which would SIGSEGV under the default +// Verification::Minimal load mode. +TEST(ValidateTensorLayoutTest, NullDimOrderIsRejected) { + flatbuffers::FlatBufferBuilder builder; + + std::vector sizes = {2, 3, 4}; + + // Pass 0 for the dim_order offset to serialize a null dim_order field. + auto tensor_offset = executorch_flatbuffer::CreateTensor( + builder, + executorch_flatbuffer::ScalarType::FLOAT, + /*storage_offset=*/0, + builder.CreateVector(sizes), + /*dim_order=*/0); + builder.Finish(tensor_offset); + + const auto* s_tensor = flatbuffers::GetRoot( + builder.GetBufferPointer()); + ASSERT_EQ(s_tensor->dim_order(), nullptr); + + 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); +} + // Helper to construct a flatbuffers::Vector from raw data. // FlatBuffer vectors are stored as [uint32_t length][T elements...]. namespace {