Merge pull request #2779 from MRtrix3/tckgen_tensorprob_fix

Fix buffer overflow in DWI bootstrapping

src/dwi/bootstrap.h

@@ -45,8 +45,10 @@ class Bootstrap : public Adapter::Base<Bootstrap<ImageType, Functor, NUM_VOX_PER
   };
 
   Bootstrap(const ImageType &Image, const Functor &functor)
-      : base_type(Image), func(functor), next_voxel(nullptr), last_voxel(nullptr) {
+      : base_type(Image), func(functor), next_voxel(nullptr), last_voxel(nullptr), current_chunk(0) {
     assert(ndim() == 4);
+    voxel_buffer.push_back(std::vector<value_type>(NUM_VOX_PER_CHUNK * size(3)));
+    clear();
   }
 
   value_type value() { return get_voxel()[index(3)]; }
@@ -64,8 +66,6 @@ class Bootstrap : public Adapter::Base<Bootstrap<ImageType, Functor, NUM_VOX_PER
 
   void clear() {
     voxels.clear();
-    if (voxel_buffer.empty())
-      voxel_buffer.push_back(std::vector<value_type>(NUM_VOX_PER_CHUNK * size(3)));
     next_voxel = &voxel_buffer[0][0];
     last_voxel = next_voxel + NUM_VOX_PER_CHUNK * size(3);
     current_chunk = 0;
@@ -81,11 +81,10 @@ class Bootstrap : public Adapter::Base<Bootstrap<ImageType, Functor, NUM_VOX_PER
 
   value_type *allocate_voxel() {
     if (next_voxel == last_voxel) {
-      ++current_chunk;
-      if (current_chunk >= voxel_buffer.size())
+      if (++current_chunk >= voxel_buffer.size())
         voxel_buffer.push_back(std::vector<value_type>(NUM_VOX_PER_CHUNK * size(3)));
       assert(current_chunk < voxel_buffer.size());
-      next_voxel = &voxel_buffer.back()[0];
+      next_voxel = &voxel_buffer[current_chunk][0];
       last_voxel = next_voxel + NUM_VOX_PER_CHUNK * size(3);
     }
     value_type *retval = next_voxel;
@@ -94,16 +93,17 @@ class Bootstrap : public Adapter::Base<Bootstrap<ImageType, Functor, NUM_VOX_PER
   }
 
   value_type *get_voxel() {
-    value_type *&data(
-        voxels.insert(std::make_pair(Eigen::Vector3i(index(0), index(1), index(2)), nullptr)).first->second);
-    if (!data) {
-      data = allocate_voxel();
-      ssize_t pos = index(3);
-      for (auto l = Loop(3)(*this); l; ++l)
-        data[index(3)] = base_type::value();
-      index(3) = pos;
-      func(data);
-    }
+    const Eigen::Vector3i voxel(index(0), index(1), index(2));
+    const typename std::map<Eigen::Vector3i, value_type *, IndexCompare>::const_iterator existing = voxels.find(voxel);
+    if (existing != voxels.end())
+      return existing->second;
+    value_type *const data = allocate_voxel();
+    ssize_t pos = index(3);
+    for (auto l = Loop(3)(*this); l; ++l)
+      data[index(3)] = base_type::value();
+    index(3) = pos;
+    func(data);
+    voxels.insert({voxel, data});
     return data;
   }
 };