From f281cd060d28652ffa8f36b97b8fd9cc79d6c9d8 Mon Sep 17 00:00:00 2001
From: Joe Zuntz <joezuntz@googlemail.com>
Date: Mon, 7 Nov 2022 13:36:08 +0000
Subject: [PATCH] Fix hard crash when using wrong array types with put/get

---
 cosmosis/datablock/c_datablock.cc | 53 ++++++++++++++++++++-----------
 cosmosis/test/test_block.py       | 31 +++++++++++++++++-
 2 files changed, 65 insertions(+), 19 deletions(-)

diff --git a/cosmosis/datablock/c_datablock.cc b/cosmosis/datablock/c_datablock.cc
index 49f7b1db..abc6d046 100644
--- a/cosmosis/datablock/c_datablock.cc
+++ b/cosmosis/datablock/c_datablock.cc
@@ -564,10 +564,16 @@ extern "C"
     if (sz == nullptr) return DBS_SIZE_NULL;
 
     auto p = static_cast<DataBlock *>(s);
-    vector<int> const& r = p->view<vector<int>>(section, name);
-    *sz = r.size();
-    if (r.size() > static_cast<size_t>(maxsize)) return DBS_SIZE_INSUFFICIENT;
-    std::copy(r.cbegin(), r.cend(), val);
+    try {
+      vector<int> const& r = p->view<vector<int>>(section, name);
+      *sz = r.size();
+      if (r.size() > static_cast<size_t>(maxsize)) return DBS_SIZE_INSUFFICIENT;
+      std::copy(r.cbegin(), r.cend(), val);
+    }
+    catch (DataBlock::BadDataBlockAccess const&) { return DBS_SECTION_NOT_FOUND; }
+    catch (Section::BadSectionAccess const&) { return DBS_NAME_NOT_FOUND; }
+    catch (Entry::BadEntry const&) { return DBS_WRONG_VALUE_TYPE; }
+    catch (...) { return DBS_LOGIC_ERROR; }
     // If we are asked to clear out the remainder of the input buffer,
     // the following line should be used.
     //    std::fill(val + *sz, val+maxsize, 0);
@@ -590,10 +596,16 @@ extern "C"
     if (sz == nullptr) return DBS_SIZE_NULL;
 
     auto p = static_cast<DataBlock *>(s);
-    vector<double> const& r = p->view<vector<double>>(section, name);
-    *sz = r.size();
-    if (r.size() > static_cast<size_t>(maxsize)) return DBS_SIZE_INSUFFICIENT;
-    std::copy(r.cbegin(), r.cend(), val);
+    try {
+      vector<double> const& r = p->view<vector<double>>(section, name);
+      *sz = r.size();
+      if (r.size() > static_cast<size_t>(maxsize)) return DBS_SIZE_INSUFFICIENT;
+      std::copy(r.cbegin(), r.cend(), val);
+    }
+    catch (DataBlock::BadDataBlockAccess const&) { return DBS_SECTION_NOT_FOUND; }
+    catch (Section::BadSectionAccess const&) { return DBS_NAME_NOT_FOUND; }
+    catch (Entry::BadEntry const&) { return DBS_WRONG_VALUE_TYPE; }
+    catch (...) { return DBS_LOGIC_ERROR; }
     // If we are asked to clear out the remainder of the input buffer,
     // the following line should be used.
     //    std::fill(val + *sz, val+maxsize, 0);
@@ -616,14 +628,21 @@ extern "C"
     if (sz == nullptr) return DBS_SIZE_NULL;
 
     auto p = static_cast<DataBlock *>(s);
-    vector<complex_t> const& r = p->view<vector<complex_t>>(section, name);
-    *sz = r.size();
-    if (r.size() > static_cast<size_t>(maxsize)) return DBS_SIZE_INSUFFICIENT;
-    //std::copy(r.cbegin(), r.cend(), val);
-    for (size_t i = 0, n = r.size(); i != n; ++i)
-      {
-        val[i] = from_complex(r[i]);
-      }
+    try{
+      vector<complex_t> const& r = p->view<vector<complex_t>>(section, name);
+      *sz = r.size();
+      if (r.size() > static_cast<size_t>(maxsize)) return DBS_SIZE_INSUFFICIENT;
+      //std::copy(r.cbegin(), r.cend(), val);
+      for (size_t i = 0, n = r.size(); i != n; ++i)
+        {
+          val[i] = from_complex(r[i]);
+        }
+    }
+    catch (DataBlock::BadDataBlockAccess const&) { return DBS_SECTION_NOT_FOUND; }
+    catch (Section::BadSectionAccess const&) { return DBS_NAME_NOT_FOUND; }
+    catch (Entry::BadEntry const&) { return DBS_WRONG_VALUE_TYPE; }
+    catch (...) { return DBS_LOGIC_ERROR; }
+
 
     // If we are asked to clear out the remainder of the input buffer,
     // the following line should be used.
@@ -652,8 +671,6 @@ extern "C"
       for (int i=0; i<*sz; i++){
         val[i] = strdup(r[i].c_str());
       }
-
-      *sz = r.size();
     }
     catch (DataBlock::BadDataBlockAccess const&) { return DBS_SECTION_NOT_FOUND; }
     catch (Section::BadSectionAccess const&) { return DBS_NAME_NOT_FOUND; }
diff --git a/cosmosis/test/test_block.py b/cosmosis/test/test_block.py
index 8ddd6a23..fca4d66f 100644
--- a/cosmosis/test/test_block.py
+++ b/cosmosis/test/test_block.py
@@ -149,8 +149,37 @@ def test_keys():
         assert k in b
 
 
+def test_wrong_array_type():
+    puts = {
+        int:   "put_int_array_1d",
+        float: "put_double_array_1d",
+        str:   "put_string_array_1d",
+    }
+    gets = {
+        int:   "get_int_array_1d",
+        float: "get_double_array_1d",
+        str:   "get_string_array_1d",
+    }
+    dtypes = list(puts.keys())
+
+    for d1 in dtypes[:]:
+        for d2 in dtypes[:]:
+            if d1 is d2:
+                continue
+
+            b = DataBlock()
+            section = 'section'
+            key = 'key'
+            put = getattr(b, puts[d1])
+            get = getattr(b, gets[d2])
+
+            value = np.array([1, 2, 3], dtype=d1)
+            put(section, key, value)
+            with pytest.raises(errors.BlockWrongValueType):
+                get(section, key)
 
 
 if __name__ == '__main__':
     # test_string_array()
-    test_string_array_save()    
\ No newline at end of file
+    # test_string_array_save()
+    test_wrong_array_type()