Check for missing arrays in Program

runtime/executor/method.cpp

@@ -388,12 +388,11 @@ Error Method::parse_values() {
         // subtract one to keep the output in 0 based indexing for a
         // disgruntled debugger seeing this error message and checking
         // schema.fbs
-        ET_CHECK_MSG(
-            false,
-            "Enum KernelTypes type: %" PRIu32
-            " not supported. Please look in executorch/schema/program.fbs "
-            "to see which type this is.",
+        ET_LOG(
+            Error,
+            "Unknown KernelTypes value %" PRIu32,
             static_cast<uint32_t>(serialization_value->val_type()) - 1);
+        return Error::InvalidProgram;
     }
 
     // ~Method() will try to clean up n_value_ entries in the values_ array.

runtime/executor/program.cpp

@@ -146,30 +146,33 @@ Result<executorch_flatbuffer::ExecutionPlan*> get_execution_plan(
 
   // Constant data may live inside the flatbuffer data (constant_buffer) or in a
   // separate segment (constant_segment). It should not be in both.
-  const auto& constant_buffer = flatbuffer_program->constant_buffer();
-  const auto& constant_segment = flatbuffer_program->constant_segment();
-
-  // Check if the constant data is inside a separate segment.
-  if (constant_segment != nullptr && constant_segment->offsets()->size() > 0) {
+  const auto* constant_segment = flatbuffer_program->constant_segment();
+  if (constant_segment != nullptr && constant_segment->offsets() != nullptr &&
+      constant_segment->offsets()->size() > 0) {
+    // The constant data is inside a separate segment.
+    const auto* constant_buffer = flatbuffer_program->constant_buffer();
     ET_CHECK_OR_RETURN_ERROR(
-        constant_buffer->size() == 0,
-        InvalidState,
-        "constant_buffer contains %u items, constant_segment.offsets contains %u items. Only one should be used.",
+        constant_buffer == nullptr || constant_buffer->size() == 0,
+        InvalidProgram,
+        "constant_buffer contains %u items, "
+        "constant_segment.offsets contains %u items. Only one should be used.",
         constant_buffer->size(),
         constant_segment->offsets()->size());
+    const auto* segments = flatbuffer_program->segments();
+    ET_CHECK_OR_RETURN_ERROR(
+        segments != nullptr, InvalidProgram, "No segments in program");
 
     // Load constant segment.
     // TODO(T171839323): Add test for segment_index > num available segments.
     ET_CHECK_OR_RETURN_ERROR(
-        constant_segment->segment_index() <
-            flatbuffer_program->segments()->size(),
-        InvalidArgument,
+        constant_segment->segment_index() < segments->size(),
+        InvalidProgram,
         "Constant segment index %d invalid for program segments range %d",
         constant_segment->segment_index(),
-        flatbuffer_program->segments()->size());
+        segments->size());
 
     const executorch_flatbuffer::DataSegment* data_segment =
-        flatbuffer_program->segments()->Get(constant_segment->segment_index());
+        segments->Get(constant_segment->segment_index());
     Result<FreeableBuffer> constant_segment_data = loader->Load(
         segment_base_offset + data_segment->offset(), data_segment->size());
     if (!constant_segment_data.ok()) {
@@ -199,7 +202,12 @@ Result<executorch_flatbuffer::ExecutionPlan*> get_execution_plan(
 size_t Program::num_methods() const {
   auto internal_program =
       static_cast<const executorch_flatbuffer::Program*>(internal_program_);
-  return internal_program->execution_plan()->size();
+  const auto execution_plan = internal_program->execution_plan();
+  if (execution_plan != nullptr) {
+    return execution_plan->size();
+  } else {
+    return 0;
+  }
 }
 
 Result<const char*> Program::get_method_name(size_t plan_index) const {