Fix segfault in rd_grav_add_survey_FIP

equinor · Aug 6, 2024 · b223ff5 · b223ff5
1 parent b1121a9
commit b223ff5
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 3 deletions.
diff --git a/lib/resdata/rd_grav.cpp b/lib/resdata/rd_grav.cpp
@@ -203,6 +203,10 @@ rd_grav_phase_alloc(rd_grav_type *rd_grav, rd_grav_survey_type *survey,
                 rd_file_iget_named_kw(init_file, PVTNUM_KW, 0);
             const std::vector<double> std_density =
                 rd_grav->std_density[std::string(rd_get_phase_name(phase))];
+
+            if (std_density.size() < size)
+                return NULL;
+
             rd_kw_type *fip_kw;
 
             if (phase == RD_OIL_PHASE)
@@ -440,6 +444,10 @@ rd_grav_survey_alloc_RPORV(rd_grav_type *rd_grav,
                            const char *name) {
     rd_grav_survey_type *survey =
         rd_grav_survey_alloc_empty(rd_grav, name, GRAV_CALC_RPORV);
+
+    if (survey == NULL)
+        return NULL;
+
     if (rd_file_view_has_kw(restart_file, RPORV_KW)) {
         rd_kw_type *rporv_kw =
             rd_file_view_iget_named_kw(restart_file, RPORV_KW, 0);
@@ -466,6 +474,10 @@ rd_grav_survey_alloc_PORMOD(rd_grav_type *rd_grav,
     rd::rd_grid_cache &grid_cache = *(rd_grav->grid_cache);
     rd_grav_survey_type *survey =
         rd_grav_survey_alloc_empty(rd_grav, name, GRAV_CALC_PORMOD);
+
+    if (survey == NULL)
+        return NULL;
+
     rd_kw_type *init_porv_kw = rd_file_iget_named_kw(
         rd_grav->init_file, PORV_KW, 0); /* Global indexing */
     rd_kw_type *pormod_kw = rd_file_view_iget_named_kw(restart_file, PORMOD_KW,
@@ -497,7 +509,9 @@ rd_grav_survey_alloc_FIP(rd_grav_type *rd_grav,
 
     rd_grav_survey_type *survey =
         rd_grav_survey_alloc_empty(rd_grav, name, GRAV_CALC_FIP);
-    rd_grav_survey_add_phases(rd_grav, survey, restart_file, GRAV_CALC_FIP);
+
+    if (survey != NULL)
+        rd_grav_survey_add_phases(rd_grav, survey, restart_file, GRAV_CALC_FIP);
 
     return survey;
 }
@@ -509,7 +523,10 @@ rd_grav_survey_alloc_RFIP(rd_grav_type *rd_grav,
 
     rd_grav_survey_type *survey =
         rd_grav_survey_alloc_empty(rd_grav, name, GRAV_CALC_RFIP);
-    rd_grav_survey_add_phases(rd_grav, survey, restart_file, GRAV_CALC_RFIP);
+
+    if (survey != NULL)
+        rd_grav_survey_add_phases(rd_grav, survey, restart_file,
+                                  GRAV_CALC_RFIP);
 
     return survey;
 }
@@ -574,6 +591,8 @@ rd_grav_add_survey_RPORV(rd_grav_type *grav, const char *name,
                          const rd_file_view_type *restart_file) {
     rd_grav_survey_type *survey =
         rd_grav_survey_alloc_RPORV(grav, restart_file, name);
+    if (survey == NULL)
+        return NULL;
     rd_grav_add_survey__(grav, name, survey);
     return survey;
 }
@@ -583,6 +602,8 @@ rd_grav_add_survey_FIP(rd_grav_type *grav, const char *name,
                        const rd_file_view_type *restart_file) {
     rd_grav_survey_type *survey =
         rd_grav_survey_alloc_FIP(grav, restart_file, name);
+    if (survey == NULL)
+        return NULL;
     rd_grav_add_survey__(grav, name, survey);
     return survey;
 }
@@ -592,6 +613,8 @@ rd_grav_add_survey_RFIP(rd_grav_type *grav, const char *name,
                         const rd_file_view_type *restart_file) {
     rd_grav_survey_type *survey =
         rd_grav_survey_alloc_RFIP(grav, restart_file, name);
+    if (survey == NULL)
+        return NULL;
     rd_grav_add_survey__(grav, name, survey);
     return survey;
 }
@@ -601,6 +624,8 @@ rd_grav_add_survey_PORMOD(rd_grav_type *grav, const char *name,
                           const rd_file_view_type *restart_file) {
     rd_grav_survey_type *survey =
         rd_grav_survey_alloc_PORMOD(grav, restart_file, name);
+    if (survey == NULL)
+        return NULL;
     rd_grav_add_survey__(grav, name, survey);
     return survey;
 }

diff --git a/python/resdata/gravimetry/rd_grav.py b/python/resdata/gravimetry/rd_grav.py
@@ -127,7 +127,9 @@ def add_survey_FIP(self, survey_name: str, restart_view: ResdataFileView):
         the new_std_density() (and possibly also add_std_density())
         method before calling the add_survey_FIP() method.
         """
-        self._add_survey_FIP(survey_name, restart_view)
+        void_ptr = self._add_survey_FIP(survey_name, restart_view)
+        if void_ptr is None:
+            raise ValueError("Could not add FIP survey due to missing std_density")
 
     def add_survey_RFIP(self, survey_name: str, restart_view: ResdataFileView):
         """

diff --git a/python/tests/rd_tests/test_grav.py b/python/tests/rd_tests/test_grav.py
@@ -98,3 +98,8 @@ def test_create(self):
             grav.add_survey_RFIP("fip", restart_view)
 
             grav.eval("rporv", "pormod", (0, 0, 0), phase_mask=1)
+
+            # Test that missing std_density raises
+            grav = ResdataGrav(self.grid, self.init)
+            with self.assertRaises(ValueError):
+                grav.add_survey_FIP("fip", restart_view)