Guard against null flatbuffer containers when getting constant data

Differential Revision: D69562793

Pull Request resolved: #8438

Author: dbort

runtime/executor/program.cpp

@@ -294,7 +294,12 @@ Result<const void*> Program::get_constant_buffer_data(
   // loaded during Program::load, or stored inside the flatbuffer data
   // (constant_buffer).
   if (constant_segment_data_.data() != nullptr) {
-    size_t num_elems = internal_program->constant_segment()->offsets()->size();
+    const auto* constant_segment = internal_program->constant_segment();
+    size_t num_elems = constant_segment == nullptr
+        ? 0
+        : (constant_segment->offsets() == nullptr
+               ? 0
+               : constant_segment->offsets()->size());
     ET_CHECK_OR_RETURN_ERROR(
         buffer_index < num_elems,
         InvalidArgument,
@@ -326,25 +331,27 @@ Result<const void*> Program::get_constant_buffer_data(
         offset);
   } else {
     // Otherwise, the constant data is stored inside Program.constant_buffer.
-    size_t num_elems = internal_program->constant_buffer()->size();
+    const auto* constant_buffer_ptr = internal_program->constant_buffer();
+    size_t num_elems =
+        constant_buffer_ptr == nullptr ? 0 : constant_buffer_ptr->size();
     ET_CHECK_OR_RETURN_ERROR(
         buffer_index < num_elems,
         InvalidArgument,
         "Constant buffer index %zu invalid for program constant buffer range %zu",
         buffer_index,
         num_elems);
 
-    const auto& constant_buffer = *internal_program->constant_buffer();
-
+    const auto& constant_buffer = *constant_buffer_ptr;
+    const auto* storage = constant_buffer[buffer_index]->storage();
+    auto storage_size = storage == nullptr ? 0 : storage->size();
     ET_CHECK_OR_RETURN_ERROR(
-        constant_buffer[buffer_index]->storage()->size() <= nbytes,
+        storage_size <= nbytes,
         InvalidArgument,
         "Constant buffer size %u larger than allocated nbytes %zu",
-        constant_buffer[buffer_index]->storage()->size(),
+        storage_size,
         nbytes);
 
-    return static_cast<const void*>(
-        constant_buffer[buffer_index]->storage()->data());
+    return storage->data();
   }
 }