[runtime] Fix null operator overload logging and named data key dereferences

extension/flat_tensor/flat_tensor_data_map.cpp

@@ -63,14 +63,18 @@ Result<const flat_tensor_flatbuffer::NamedData*> get_named_data(
     return Error::InvalidExternalData;
   }
   for (flatbuffers::uoffset_t i = 0; i < named_data->size(); ++i) {
-    if (key.size() == named_data->Get(i)->key()->size() &&
-        std::strncmp(
-            named_data->Get(i)->key()->c_str(),
-            key.data(),
-            named_data->Get(i)->key()->size()) == 0) {
-      const auto* found = named_data->Get(i);
+    const auto* entry = named_data->Get(i);
+    ET_CHECK_OR_RETURN_ERROR(
+        entry != nullptr && entry->key() != nullptr,
+        InvalidExternalData,
+        "Named data entry or key at index %" PRIu32
+        " is null. Malformed PTD file.",
+        i);
+    if (key.size() == entry->key()->size() &&
+        std::strncmp(entry->key()->c_str(), key.data(), entry->key()->size()) ==
+            0) {
       // Validate the named_data.
-      const flatbuffers::uoffset_t segment_index = found->segment_index();
+      const flatbuffers::uoffset_t segment_index = entry->segment_index();
       ET_CHECK_OR_RETURN_ERROR(
           segment_index < segments->size(),
           InvalidExternalData,
@@ -94,7 +98,7 @@ Result<const flat_tensor_flatbuffer::NamedData*> get_named_data(
           segments->Get(segment_index)->offset(),
           segments->Get(segment_index)->size(),
           segment_data_size);
-      return found;
+      return entry;
     }
   }
   return Error::NotFound;
@@ -254,14 +258,21 @@ ET_NODISCARD Result<uint32_t> FlatTensorDataMap::get_num_keys() const {
 
 ET_NODISCARD Result<const char*> FlatTensorDataMap::get_key(
     uint32_t index) const {
-  uint32_t num_keys = get_num_keys().get();
+  const auto* named_data = flat_tensor_->named_data();
   ET_CHECK_OR_RETURN_ERROR(
-      index < num_keys,
+      index < named_data->size(),
       InvalidArgument,
       "Index %" PRIu32 " out of range of size %" PRIu32,
       index,
-      num_keys);
-  return flat_tensor_->named_data()->Get(index)->key()->c_str();
+      named_data->size());
+  const auto* entry = named_data->Get(index);
+  ET_CHECK_OR_RETURN_ERROR(
+      entry != nullptr && entry->key() != nullptr,
+      InvalidExternalData,
+      "Named data entry or key at index %" PRIu32
+      " is null. Malformed PTD file.",
+      index);
+  return entry->key()->c_str();
 }
 
 /* static */ Result<FlatTensorDataMap> FlatTensorDataMap::load(

extension/flat_tensor/test/flat_tensor_data_map_test.cpp

@@ -208,7 +208,8 @@ struct SegmentSpec {
 // test loader, so large source offsets do not require a large allocation.
 std::vector<uint8_t> CreateDataWithVersion(
     uint32_t version,
-    const SegmentSpec* segment = nullptr) {
+    const SegmentSpec* segment = nullptr,
+    bool with_null_key = false) {
   flatbuffers::FlatBufferBuilder builder;
 
   std::vector<flatbuffers::Offset<flat_tensor_flatbuffer::DataSegment>>
@@ -223,13 +224,16 @@ std::vector<uint8_t> CreateDataWithVersion(
         executorch_flatbuffer::ScalarType::FLOAT,
         builder.CreateVector(sizes),
         builder.CreateVector(dim_order));
-    segments.push_back(flat_tensor_flatbuffer::CreateDataSegment(
-        builder, segment->offset, segment->size));
-    named_data.push_back(flat_tensor_flatbuffer::CreateNamedData(
-        builder,
-        builder.CreateString(kWideTensorKey),
-        /*segment_index=*/0,
-        tensor_layout));
+    segments.push_back(
+        flat_tensor_flatbuffer::CreateDataSegment(
+            builder, segment->offset, segment->size));
+    named_data.push_back(
+        flat_tensor_flatbuffer::CreateNamedData(
+            builder,
+            with_null_key ? flatbuffers::Offset<flatbuffers::String>()
+                          : builder.CreateString(kWideTensorKey),
+            /*segment_index=*/0,
+            tensor_layout));
   }
 
   auto flat_tensor = flat_tensor_flatbuffer::CreateFlatTensor(
@@ -384,6 +388,43 @@ TEST_F(FlatTensorDataMapTest, NewerSchemaVersionFailsToLoad) {
   EXPECT_EQ(data_map.error(), Error::InvalidExternalData);
 }
 
+TEST_F(FlatTensorDataMapTest, KeylessNamedDataReturnsInvalidExternalData) {
+  constexpr SegmentSpec kSegment{0, sizeof(float), sizeof(float)};
+  std::vector<uint8_t> data = CreateDataWithVersion(
+      FlatTensorDataMap::kMaxSupportedSchemaVersion,
+      &kSegment,
+      /*with_null_key=*/true);
+
+  alignas(std::max_align_t) std::array<uint8_t, 1024> aligned_buffer{};
+  ASSERT_LE(data.size(), aligned_buffer.size());
+  std::memcpy(aligned_buffer.data(), data.data(), data.size());
+
+  Result<FlatTensorHeader> header =
+      FlatTensorHeader::Parse(aligned_buffer.data(), data.size());
+  ASSERT_TRUE(header.ok());
+  const uint64_t absolute_offset = header->segment_base_offset;
+  const float segment_data = 1.0f;
+  WideOffsetDataLoader loader(
+      aligned_buffer.data(),
+      data.size(),
+      reinterpret_cast<const uint8_t*>(&segment_data),
+      sizeof(segment_data),
+      absolute_offset + sizeof(segment_data));
+
+  Result<FlatTensorDataMap> data_map = FlatTensorDataMap::load(&loader);
+  ASSERT_TRUE(data_map.ok());
+
+  // get_key for the keyless entry must fail with InvalidExternalData.
+  EXPECT_EQ(data_map->get_key(0).error(), Error::InvalidExternalData);
+
+  // Linear search in get_data and get_tensor_layout must encounter the keyless
+  // entry and fail with InvalidExternalData instead of silently skipping it.
+  EXPECT_EQ(data_map->get_data("any_key").error(), Error::InvalidExternalData);
+  EXPECT_EQ(
+      data_map->get_tensor_layout("any_key").error(),
+      Error::InvalidExternalData);
+}
+
 TEST_F(FlatTensorDataMapTest, RejectsOutOfBoundsRootOffset) {
   // A valid file loads.
   Result<FreeableBuffer> valid = data_map_loader_->load(

runtime/executor/method.cpp

@@ -1735,11 +1735,15 @@ Error Method::execute_instruction() {
         ET_LOG(
             Error,
             "KernelCall failed at instruction %" ET_PRIsize_t ":%" ET_PRIsize_t
-            " in operator %s.%s: 0x%x",
+            " in operator %s%s%s: 0x%x",
             step_state_.chain_idx,
             step_state_.instr_idx,
             op->name()->c_str(),
-            op->overload()->c_str(),
+            (op->overload() != nullptr && op->overload()->size() > 0) ? "."
+                                                                      : "",
+            (op->overload() != nullptr && op->overload()->size() > 0)
+                ? op->overload()->c_str()
+                : "",
             (unsigned int)err);
         for (size_t i = 0; i < args.size(); ++i) {
           ET_LOG(

runtime/executor/test/CMakeLists.txt

@@ -215,6 +215,19 @@ target_include_directories(
                        "${EXECUTORCH_ROOT}/third-party/flatbuffers/include"
 )
 
+et_cxx_test(
+  program_validation_test SOURCES program_validation_test.cpp EXTRA_LIBS
+  extension_data_loader program_schema
+)
+add_dependencies(program_validation_test generated_pte_files)
+set_property(TEST program_validation_test PROPERTY ENVIRONMENT ${test_env})
+
+target_include_directories(
+  program_validation_test
+  PRIVATE "${CMAKE_INSTALL_PREFIX}/schema/include"
+          "${EXECUTORCH_ROOT}/third-party/flatbuffers/include"
+)
+
 et_cxx_test(
   kernel_resolution_test SOURCES kernel_resolution_test.cpp EXTRA_LIBS
   extension_data_loader

runtime/executor/test/program_validation_test.cpp

@@ -15,6 +15,7 @@
 #include <executorch/runtime/executor/method.h>
 #include <executorch/runtime/executor/program.h>
 #include <executorch/runtime/executor/test/managed_memory_manager.h>
+#include <executorch/runtime/kernel/kernel_runtime_context.h>
 #include <executorch/runtime/kernel/operator_registry.h>
 #include <executorch/runtime/platform/runtime.h>
 #include <executorch/schema/program_generated.h>
@@ -78,7 +79,8 @@ std::vector<uint8_t> CreateTestProgram(
     const std::vector<int32_t>& input_indices = {},
     // If true, emit a chain of two KernelCalls: the first's op_index is out of
     // range so resolve_operator() fails on a non-final instruction.
-    bool with_out_of_range_kernel_call = false) {
+    bool with_out_of_range_kernel_call = false,
+    bool with_failing_kernel_no_overload = false) {
   flatbuffers::FlatBufferBuilder builder(1024);
 
   std::vector<flatbuffers::Offset<executorch_flatbuffer::EValue>> evalues;
@@ -161,12 +163,33 @@ std::vector<uint8_t> CreateTestProgram(
         /*outputs=*/builder.CreateVector(std::vector<int32_t>{}),
         builder.CreateVector(instructions),
         /*stacktrace=*/0));
+  } else if (with_failing_kernel_no_overload) {
+    auto kernel_call = executorch_flatbuffer::CreateKernelCall(
+        builder, /*op_index=*/0, builder.CreateVector(std::vector<int32_t>{}));
+    std::vector<flatbuffers::Offset<executorch_flatbuffer::Instruction>>
+        instructions{executorch_flatbuffer::CreateInstruction(
+            builder,
+            executorch_flatbuffer::InstructionArguments::KernelCall,
+            kernel_call.Union())};
+    chains.push_back(
+        executorch_flatbuffer::CreateChain(
+            builder,
+            /*inputs=*/builder.CreateVector(std::vector<int32_t>{}),
+            /*outputs=*/builder.CreateVector(std::vector<int32_t>{}),
+            builder.CreateVector(instructions),
+            /*stacktrace=*/0));
   }
   auto chain_vec = builder.CreateVector(chains);
   std::vector<flatbuffers::Offset<executorch_flatbuffer::Operator>> operators;
   if (with_out_of_range_kernel_call) {
     operators.push_back(executorch_flatbuffer::CreateOperator(
         builder, builder.CreateString("test::noop"), builder.CreateString("")));
+  } else if (with_failing_kernel_no_overload) {
+    operators.push_back(
+        executorch_flatbuffer::CreateOperator(
+            builder,
+            builder.CreateString("test::failing_no_overload"),
+            /*overload=*/0));
   }
   auto operators_vec = builder.CreateVector(operators);
   auto empty_delegates_vec = builder.CreateVector(
@@ -189,14 +212,19 @@ std::vector<uint8_t> CreateTestProgram(
   plans.push_back(execution_plan);
   auto plans_vec = builder.CreateVector(plans);
 
+  auto constant_segment = executorch_flatbuffer::CreateSubsegmentOffsets(
+      builder,
+      /*segment_index=*/0,
+      builder.CreateVector(std::vector<uint64_t>{0}));
+
   auto program = executorch_flatbuffer::CreateProgram(
       builder,
       /*version=*/0,
       plans_vec,
       /*constant_buffer=*/0,
       /*backend_delegate_data=*/0,
       /*segments=*/0,
-      /*constant_segment=*/0,
+      constant_segment,
       /*mutable_data_segments=*/0,
       /*named_data=*/0);
 
@@ -404,3 +432,36 @@ TEST_F(ProgramValidationTest, ResolveOperatorFailureSurfacesFromInit) {
       program->load_method("forward", &mmm.get());
   EXPECT_EQ(method.error(), Error::InvalidProgram);
 }
+
+TEST_F(ProgramValidationTest, KernelCallFailureWithNullOverloadDoesNotCrash) {
+  static const auto fail_kernel = executorch::runtime::Kernel(
+      "test::failing_no_overload",
+      [](executorch::runtime::KernelRuntimeContext& context,
+         executorch::runtime::Span<executorch::runtime::EValue*>) {
+        context.fail(Error::Internal);
+      });
+  ASSERT_EQ(executorch::runtime::register_kernel(fail_kernel), Error::Ok);
+
+  AlignedBuffer buf(CreateTestProgram(
+      /*configs=*/{},
+      /*input_indices=*/{},
+      /*with_out_of_range_kernel_call=*/false,
+      /*with_failing_kernel_no_overload=*/true));
+  auto loader = buf.loader();
+
+  Result<Program> program =
+      Program::load(&loader, Program::Verification::InternalConsistency);
+  ASSERT_EQ(program.error(), Error::Ok);
+
+  executorch::runtime::testing::ManagedMemoryManager mmm(
+      /*planned_memory_bytes=*/32 * 1024U,
+      /*method_allocator_bytes=*/32 * 1024U);
+  Result<executorch::runtime::Method> method =
+      program->load_method("forward", &mmm.get());
+  ASSERT_EQ(method.error(), Error::Ok);
+
+  // KernelCall failure logs operator overload. With a null overload in the
+  // flatbuffer operator table, this previously segfaulted (#23955).
+  Error err = method->execute();
+  EXPECT_EQ(err, Error::Internal);
+}