From 1aef6504775d3e07e2221c58b5b0ca189a93e8ad Mon Sep 17 00:00:00 2001
From: Marcus Holland-Moritz <github@mhxnet.de>
Date: Sun, 30 Mar 2025 14:14:07 +0200
Subject: [PATCH] refactor(cached_block): safeguard against moving block
 reallocation

---
 include/dwarfs/byte_buffer.h         |  6 ++++
 src/reader/internal/cached_block.cpp | 16 ++++++++--
 src/vector_byte_buffer.cpp           | 45 +++++++++++++++++++++++++---
 3 files changed, 60 insertions(+), 7 deletions(-)

diff --git a/include/dwarfs/byte_buffer.h b/include/dwarfs/byte_buffer.h
index 4ec9fe7b..24153ea6 100644
--- a/include/dwarfs/byte_buffer.h
+++ b/include/dwarfs/byte_buffer.h
@@ -61,6 +61,10 @@ class mutable_byte_buffer_interface : public byte_buffer_interface {
   virtual void resize(size_t size) = 0;
   virtual void shrink_to_fit() = 0;
 
+  // Freezes the location of the buffer in memory, i.e. all further calls
+  // that would reallocate the buffer will throw.
+  virtual void freeze_location() = 0;
+
   // TODO: See if we can do without this. This will *only* be implemented
   //       in the vector_byte_buffer, other implementations will throw.
   virtual std::vector<uint8_t>& raw_vector() = 0;
@@ -147,6 +151,8 @@ class mutable_byte_buffer {
 
   void shrink_to_fit() { bb_->shrink_to_fit(); }
 
+  void freeze_location() { bb_->freeze_location(); }
+
   std::vector<uint8_t>& raw_vector() { return bb_->raw_vector(); }
 
   void swap(mutable_byte_buffer& other) noexcept { std::swap(bb_, other.bb_); }
diff --git a/src/reader/internal/cached_block.cpp b/src/reader/internal/cached_block.cpp
index 0f037b45..bdc875b0 100644
--- a/src/reader/internal/cached_block.cpp
+++ b/src/reader/internal/cached_block.cpp
@@ -76,7 +76,9 @@ class cached_block_ final : public cached_block {
   // once the block is fully decompressed, we can reset the decompressor_
 
   // This can be called from any thread
-  size_t range_end() const override { return range_end_.load(); }
+  size_t range_end() const override {
+    return range_end_.load(std::memory_order_acquire);
+  }
 
   // TODO: The code relies on the fact that the data_ buffer is never
   // reallocated once block decompression has started. I would like to
@@ -84,7 +86,9 @@ class cached_block_ final : public cached_block {
   uint8_t const* data() const override { return data_.data(); }
 
   void decompress_until(size_t end) override {
-    while (data_.size() < end) {
+    auto pos = data_.size();
+
+    while (pos < end) {
       if (!decompressor_) {
         DWARFS_THROW(runtime_error, "no decompressor for block");
       }
@@ -97,7 +101,13 @@ class cached_block_ final : public cached_block {
         try_release();
       }
 
-      range_end_ = data_.size();
+      if (pos == 0) {
+        // Freeze the location of the data buffer
+        data_.freeze_location();
+      }
+
+      pos = data_.size();
+      range_end_.store(pos, std::memory_order_release);
     }
   }
 
diff --git a/src/vector_byte_buffer.cpp b/src/vector_byte_buffer.cpp
index da853e38..f5cbdc6e 100644
--- a/src/vector_byte_buffer.cpp
+++ b/src/vector_byte_buffer.cpp
@@ -19,6 +19,10 @@
  * along with dwarfs.  If not, see <https://www.gnu.org/licenses/>.
  */
 
+#include <atomic>
+#include <stdexcept>
+#include <string>
+#include <string_view>
 #include <vector>
 
 #include <dwarfs/vector_byte_buffer.h>
@@ -55,18 +59,51 @@ class vector_byte_buffer_impl : public mutable_byte_buffer_interface {
     return {data_.data(), data_.size()};
   }
 
-  void clear() override { data_.clear(); }
+  void clear() override {
+    assert_not_frozen("clear");
+    data_.clear();
+  }
 
-  void reserve(size_t size) override { data_.reserve(size); }
+  void reserve(size_t size) override {
+    assert_not_frozen("reserve");
+    data_.reserve(size);
+  }
 
-  void resize(size_t size) override { data_.resize(size); }
+  void resize(size_t size) override {
+    if (frozen() && size > data_.capacity()) {
+      frozen_error("resize beyond capacity");
+    }
+    data_.resize(size);
+  }
 
-  void shrink_to_fit() override { data_.shrink_to_fit(); }
+  void shrink_to_fit() override {
+    assert_not_frozen("shrink_to_fit");
+    data_.shrink_to_fit();
+  }
+
+  void freeze_location() override {
+    frozen_.store(true, std::memory_order_release);
+  }
 
   std::vector<uint8_t>& raw_vector() override { return data_; }
 
  private:
+  void assert_not_frozen(std::string_view what) const {
+    if (frozen()) {
+      frozen_error(what);
+    }
+  }
+
+  void frozen_error(std::string_view what) const {
+    throw std::runtime_error("operation not allowed on frozen buffer: " +
+                             std::string{what});
+  }
+
+  bool frozen() const { return frozen_.load(std::memory_order_acquire); }
+
   std::vector<uint8_t> data_;
+  std::atomic<bool> frozen_{false};
+  static_assert(std::atomic<bool>::is_always_lock_free);
 };
 
 } // namespace