Fix VectorHeader/pointer invalidation after buffer relocation

PayloadBuffer's VectorPush, VectorReserve, VectorResize, Realloc,
PrimeBitmapAllocator, and BitMapRun::Allocate all dereferenced pointers
into the buffer (hdr, p) after Allocate/Realloc/AllocateBitMapRun calls
that may have relocated the underlying buffer, leaving those pointers
dangling. Save the offset before any allocation and re-derive the
pointer afterward across all six call sites.

Also drop a leftover write through the dangling hdr in VectorPush that
slipped through a previous partial fix.

Add VectorPushWithResize, VectorReserveWithResize, and
VectorResizeWithResize regression tests that drive the resizer with a
small (256-byte) resizable buffer and verify data integrity afterward.
The existing tests all used fixed-size buffers and could not trigger
the bug.

Original fix by Damian Stulich <damjan.stulic@getcruise.com>.

Author: Dave Allison

toolbelt/payload_buffer.cc

@@ -584,10 +584,13 @@ void *PayloadBuffer::Realloc(PayloadBuffer **buffer, void *p, uint32_t n,
       }
       // Need to free the old block and allocate a new one as the small block
       // index is different.
+      BufferOffset p_offset = (*buffer)->ToOffset(p);
       void *newp = Allocate(buffer, n, false, enable_small_block);
       if (newp == NULL) {
         return NULL;
       }
+      // Re-derive p since Allocate may have triggered a buffer resize.
+      p = (*buffer)->ToAddress(p_offset);
       memcpy(newp, p, decoded_length);
       if (clear && n > static_cast<uint32_t>(decoded_length)) {
         memset(reinterpret_cast<char *>(newp) + decoded_length, 0,
@@ -662,10 +665,13 @@ void *PayloadBuffer::Realloc(PayloadBuffer **buffer, void *p, uint32_t n,
   // one, copy the memory and free the old block.  We are guaranteed that
   // the new block is larger than the original one since if it was smaller
   // we can always reuse the block.
+  BufferOffset p_offset = (*buffer)->ToOffset(p);
   void *newp = Allocate(buffer, n, false, enable_small_block);
   if (newp == NULL) {
     return NULL;
   }
+  // Re-derive p since Allocate may have triggered a buffer resize.
+  p = (*buffer)->ToAddress(p_offset);
   memcpy(newp, p, orig_length);
   if (clear) {
     memset(reinterpret_cast<char *>(newp) + orig_length, 0, n - orig_length);
@@ -686,15 +692,14 @@ bool PayloadBuffer::PrimeBitmapAllocator(PayloadBuffer **self, size_t size) {
     return false;
   }
   (*self)->bitmaps[index] = offset;
-  VectorHeader *hdr = (*self)->ToAddress<VectorHeader>((*self)->bitmaps[index]);
 
   BitMapRun *run = PayloadBuffer::AllocateBitMapRun(
       self, bitmp_run_infos[index].size, bitmp_run_infos[index].num);
   if (run == nullptr) {
     return false;
   }
-  // Add to the vector, this might move the vector contents but the header
-  // will stay where it is.
+  // Re-derive hdr since AllocateBitMapRun may have triggered a buffer resize.
+  VectorHeader *hdr = (*self)->ToAddress<VectorHeader>((*self)->bitmaps[index]);
   (*self)->VectorPush<BufferOffset>(self, hdr, (*self)->ToOffset(run), false);
   return true;
 }
@@ -742,8 +747,10 @@ void *BitMapRun::Allocate(PayloadBuffer **pb, int index, uint32_t, int size,
     }
     (*pb)->bitmaps[index] = offset;
   }
-  VectorHeader *hdr = (*pb)->ToAddress<VectorHeader>((*pb)->bitmaps[index]);
   for (;;) {
+    // Re-derive hdr each iteration since allocations below may trigger a
+    // buffer resize (realloc), invalidating any previous pointer.
+    VectorHeader *hdr = (*pb)->ToAddress<VectorHeader>((*pb)->bitmaps[index]);
     // Go backwards through the elements as that is most likely to find a free
     // bit.
     for (int i = hdr->num_elements - 1; i >= 0; i--) {
@@ -781,8 +788,9 @@ void *BitMapRun::Allocate(PayloadBuffer **pb, int index, uint32_t, int size,
     if (run == nullptr) {
       return nullptr;
     }
-    // Add to the vector, this might move the vector contents but the header
-    // will stay where it is.
+    // Re-derive hdr since AllocateBitMapRun may have triggered a buffer
+    // resize, invalidating the previous pointer.
+    hdr = (*pb)->ToAddress<VectorHeader>((*pb)->bitmaps[index]);
     (*pb)->VectorPush<BufferOffset>(pb, hdr, (*pb)->ToOffset(run), false);
   }
 }

toolbelt/payload_buffer.h

@@ -487,7 +487,6 @@ inline void PayloadBuffer::VectorPush(PayloadBuffer **self, VectorHeader *hdr,
   if (hdr->data == 0) {
     // The vector is empty, allocate it with a default size of 2.
     void *vecp = Allocate(self, 2 * sizeof(T), true, enable_small_block);
-    hdr->data = (*self)->ToOffset(vecp);
     VectorHeader *new_hdr = (*self)->ToAddress<VectorHeader>(hdr_offset);
     new_hdr->data = (*self)->ToOffset(vecp);
     hdr = new_hdr;

toolbelt/payload_buffer_test.cc

@@ -503,6 +503,121 @@ TEST(BufferTest, VectorExpandMore) {
   free(buffer);
 }
 
+TEST(BufferTest, VectorPushWithResize) {
+  char *buffer = (char *)calloc(256, 1);
+  bool resized = false;
+  PayloadBuffer *pb = new (buffer) PayloadBuffer(
+      256, [&resized, &buffer](PayloadBuffer **p, size_t, size_t new_size) {
+#if defined(__clang__)
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wclass-memaccess"
+#elif defined(__GNUC__)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wclass-memaccess"
+#endif
+        *p = reinterpret_cast<PayloadBuffer *>(realloc(*p, new_size));
+#if defined(__clang__)
+#pragma clang diagnostic pop
+#elif defined(__GNUC__)
+#pragma GCC diagnostic pop
+#endif
+        buffer = reinterpret_cast<char *>(*p);
+        resized = true;
+      });
+
+  PayloadBuffer::AllocateMainMessage(&pb, sizeof(VectorHeader));
+  BufferOffset msg_offset = pb->message;
+
+  constexpr int kCount = 200;
+  for (int i = 0; i < kCount; i++) {
+    VectorHeader *hdr = pb->ToAddress<VectorHeader>(msg_offset);
+    PayloadBuffer::VectorPush<uint32_t>(&pb, hdr, i + 1);
+  }
+  ASSERT_TRUE(resized);
+
+  VectorHeader *hdr = pb->ToAddress<VectorHeader>(msg_offset);
+  ASSERT_EQ(kCount, hdr->num_elements);
+  for (int i = 0; i < kCount; i++) {
+    uint32_t v = pb->VectorGet<uint32_t>(hdr, i);
+    ASSERT_EQ(i + 1, v);
+  }
+
+  pb->~PayloadBuffer();
+  free(buffer);
+}
+
+TEST(BufferTest, VectorReserveWithResize) {
+  char *buffer = (char *)calloc(256, 1);
+  bool resized = false;
+  PayloadBuffer *pb = new (buffer) PayloadBuffer(
+      256, [&resized, &buffer](PayloadBuffer **p, size_t, size_t new_size) {
+#if defined(__clang__)
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wclass-memaccess"
+#elif defined(__GNUC__)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wclass-memaccess"
+#endif
+        *p = reinterpret_cast<PayloadBuffer *>(realloc(*p, new_size));
+#if defined(__clang__)
+#pragma clang diagnostic pop
+#elif defined(__GNUC__)
+#pragma GCC diagnostic pop
+#endif
+        buffer = reinterpret_cast<char *>(*p);
+        resized = true;
+      });
+
+  PayloadBuffer::AllocateMainMessage(&pb, sizeof(VectorHeader));
+  BufferOffset msg_offset = pb->message;
+
+  VectorHeader *hdr = pb->ToAddress<VectorHeader>(msg_offset);
+  PayloadBuffer::VectorReserve<uint32_t>(&pb, hdr, 500);
+  ASSERT_TRUE(resized);
+
+  hdr = pb->ToAddress<VectorHeader>(msg_offset);
+  ASSERT_NE(0u, hdr->data);
+
+  pb->~PayloadBuffer();
+  free(buffer);
+}
+
+TEST(BufferTest, VectorResizeWithResize) {
+  char *buffer = (char *)calloc(256, 1);
+  bool resized = false;
+  PayloadBuffer *pb = new (buffer) PayloadBuffer(
+      256, [&resized, &buffer](PayloadBuffer **p, size_t, size_t new_size) {
+#if defined(__clang__)
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wclass-memaccess"
+#elif defined(__GNUC__)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wclass-memaccess"
+#endif
+        *p = reinterpret_cast<PayloadBuffer *>(realloc(*p, new_size));
+#if defined(__clang__)
+#pragma clang diagnostic pop
+#elif defined(__GNUC__)
+#pragma GCC diagnostic pop
+#endif
+        buffer = reinterpret_cast<char *>(*p);
+        resized = true;
+      });
+
+  PayloadBuffer::AllocateMainMessage(&pb, sizeof(VectorHeader));
+  BufferOffset msg_offset = pb->message;
+
+  VectorHeader *hdr = pb->ToAddress<VectorHeader>(msg_offset);
+  PayloadBuffer::VectorResize<uint32_t>(&pb, hdr, 500);
+  ASSERT_TRUE(resized);
+
+  hdr = pb->ToAddress<VectorHeader>(msg_offset);
+  ASSERT_EQ(500u, hdr->num_elements);
+
+  pb->~PayloadBuffer();
+  free(buffer);
+}
+
 TEST(BufferTest, Resizeable) {
   char *buffer = (char *)calloc(1, 512);
   bool resized = false;