Check for missing arrays in Program (#1519)

Summary:

Flatbuffer array fields can be missing, so we need to check for `nullptr` before calling `size()` on them. Discovered by lionhead fuzzing.

Differential Revision: D52493423

Author: dbort

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 {