Fix data loss when Realloc moves a block into the free block below

toolbelt/payload_buffer.cc

@@ -59,6 +59,21 @@ char *PayloadBuffer::SetString(PayloadBuffer **self, const char *s, size_t len,
   BufferOffset str_ptr = *hdr;
   void *old_str = (*self)->ToAddress(str_ptr);
 
+  // 's' can be in this buffer, which allocating can move, so it is found
+  // again there by its offset.
+  const BufferOffset source = (*self)->ContainedOffset(s);
+  if (old_str != nullptr) {
+    const char *old_begin = static_cast<const char *>(old_str);
+    const char *old_end =
+        old_begin + sizeof(uint32_t) + *static_cast<const uint32_t *>(old_str);
+    if (std::less_equal<const char *>()(old_begin, s) &&
+        std::less<const char *>()(s, old_end)) {
+      // Resizing the old string can overwrite 's', which is in it.
+      const std::string copy(s, len);
+      return SetString(self, copy.data(), len, header_offset);
+    }
+  }
+
   // If this contains a valid (non-zero) offset, reallocate the
   // data it points to, otherwise allocate new data.
   if (old_str != nullptr) {
@@ -68,7 +83,8 @@ char *PayloadBuffer::SetString(PayloadBuffer **self, const char *s, size_t len,
   }
   uint32_t *p = reinterpret_cast<uint32_t *>(str);
   p[0] = uint32_t(len);
-  memcpy(p + 1, s, len);
+  memcpy(p + 1, source == 0 ? s : reinterpret_cast<char *>(*self) + source,
+         len);
 
   // The buffer may have moved.  Reassign the address of the string
   // back into the header.
@@ -599,14 +615,8 @@ uint64_t *PayloadBuffer::MergeWithFreeBlockBelow(
   } else {
     next_ptr = &prev->next;
   }
-  // Move FreeBlockHeader to end of allocated block.  This is inside
-  // the combined free block and block being reallocated.
-  FreeBlockHeader *next = ToAddress<FreeBlockHeader>(free_block->next);
-  FreeBlockHeader *newb = reinterpret_cast<FreeBlockHeader *>(
-      free_addr + new_length + sizeof(uint64_t));
-  newb->length = free_block->length + orig_length - new_length;
-  newb->next = ToOffset(next);
-  *next_ptr = ToOffset(newb);
+  const BufferOffset next = free_block->next;
+  const uint32_t remaining = free_block->length + orig_length - new_length;
 
   uint64_t *len_ptr = reinterpret_cast<uint64_t *>(free_block);
   *len_ptr = new_length;
@@ -615,6 +625,18 @@ uint64_t *PayloadBuffer::MergeWithFreeBlockBelow(
     memset(reinterpret_cast<char *>(len_ptr) + sizeof(uint64_t) + orig_length, 0,
            new_length - orig_length);
   }
+  // What is left of the combined blocks follows the moved block, and can
+  // start inside the block's old place, so it is made after the move.  With
+  // nothing left, the next block starts there.
+  if (remaining >= sizeof(FreeBlockHeader)) {
+    FreeBlockHeader *newb = reinterpret_cast<FreeBlockHeader *>(
+        free_addr + new_length + sizeof(uint64_t));
+    newb->length = remaining;
+    newb->next = next;
+    *next_ptr = ToOffset(newb);
+  } else {
+    *next_ptr = next;
+  }
   return len_ptr + 1;
 }
 

toolbelt/payload_buffer.h

@@ -471,6 +471,13 @@ struct PayloadBuffer {
            reinterpret_cast<const char *>(this);
   }
 
+  // The offset of 'p' if it points into this buffer, otherwise 0.
+  BufferOffset ContainedOffset(const void *p) const {
+    const uintptr_t at = reinterpret_cast<uintptr_t>(p);
+    const uintptr_t start = reinterpret_cast<uintptr_t>(this);
+    return at > start && at - start < full_size ? BufferOffset(at - start) : 0;
+  }
+
   void InsertNewFreeBlockAtEnd(FreeBlockHeader *free_block,
                                FreeBlockHeader *prev, uint32_t length);
 

toolbelt/payload_buffer_test.cc

@@ -94,6 +94,46 @@ TEST(BufferTest, FreeThenAlloc) {
   free(buffer);
 }
 
+// A block that grows into the free block below it moves down.  The free
+// blocks below leave 8 bytes, the smallest free block, and nothing.
+TEST(BufferTest, ReallocIntoFreeBlockBelow) {
+  for (const uint32_t below : {256u, 248u}) {
+    char *buffer = (char *)calloc(1, 4096);
+    PayloadBuffer *pb = new (buffer) PayloadBuffer(4096);
+    void *lower = PayloadBuffer::Allocate(&pb, below, true, false);
+    auto *block =
+        static_cast<uint8_t *>(PayloadBuffer::Allocate(&pb, 256, true, false));
+    auto *upper =
+        static_cast<uint8_t *>(PayloadBuffer::Allocate(&pb, 256, true, false));
+    for (uint32_t i = 0; i < 256; ++i) {
+      block[i] = static_cast<uint8_t>(i);
+    }
+    memset(upper, 0xc5, 256);
+    const uint64_t upper_size = *(reinterpret_cast<uint64_t *>(upper) - 1);
+    pb->Free(lower);
+
+    auto *grown = static_cast<uint8_t *>(
+        PayloadBuffer::Realloc(&pb, block, 512, true, false));
+    ASSERT_EQ(grown, lower) << below;
+    for (uint32_t i = 0; i < 256; ++i) {
+      ASSERT_EQ(grown[i], i) << below;
+    }
+    for (uint32_t i = 256; i < 512; ++i) {
+      ASSERT_EQ(grown[i], 0) << below;
+    }
+    EXPECT_EQ(*(reinterpret_cast<uint64_t *>(upper) - 1), upper_size) << below;
+    for (uint32_t i = 0; i < 256; ++i) {
+      ASSERT_EQ(upper[i], 0xc5) << below;
+    }
+    for (const auto *free_block = pb->FreeList(); free_block != nullptr;
+         free_block =
+             pb->ToAddress<toolbelt::FreeBlockHeader>(free_block->next)) {
+      EXPECT_GE(free_block->length, sizeof(toolbelt::FreeBlockHeader)) << below;
+    }
+    free(buffer);
+  }
+}
+
 TEST(BufferTest, SmallBlockAllocSimple) {
   char *buffer = (char *)calloc(1, 4096);
   PayloadBuffer *pb = new (buffer) PayloadBuffer(4096);
@@ -425,6 +465,59 @@ TEST(BufferTest, String) {
   free(buffer);
 }
 
+TEST(BufferTest, StringCopiedFromTheSameBuffer) {
+  char *buffer = (char *)calloc(256, 1);
+  int resizes = 0;
+  PayloadBuffer *pb = new (buffer) PayloadBuffer(
+      256, [&resizes, &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);
+        ++resizes;
+      });
+  constexpr int kCopies = 7;
+  PayloadBuffer::AllocateMainMessage(&pb, (kCopies + 1) * sizeof(BufferOffset));
+  const BufferOffset first = pb->message;
+  auto copy = [first](int i) {
+    return first + BufferOffset((i + 1) * sizeof(BufferOffset));
+  };
+
+  std::string value(1500, '\0');
+  for (size_t i = 0; i < value.size(); ++i) {
+    value[i] = static_cast<char>('a' + i % 26);
+  }
+  PayloadBuffer::SetString(&pb, value, first);
+
+  const int resizes_before_copies = resizes;
+  for (int i = 0; i < kCopies; ++i) {
+    PayloadBuffer::SetString(&pb, pb->GetStringView(first), copy(i));
+  }
+  ASSERT_GT(resizes, resizes_before_copies);
+  for (int i = 0; i < kCopies; ++i) {
+    ASSERT_EQ(value, pb->GetStringView(copy(i))) << i;
+  }
+
+  // A part of itself, ending past the end of the shorter string.
+  PayloadBuffer::SetString(&pb, pb->GetStringView(first).substr(10, 600),
+                           first);
+  ASSERT_EQ(value.substr(10, 600), pb->GetStringView(first));
+  ASSERT_EQ(value, pb->GetStringView(copy(0)));
+
+  pb->~PayloadBuffer();
+  free(buffer);
+}
+
 TEST(BufferTest, HostileStringLengthIsClamped) {
   char *buffer = (char *)calloc(1, 4096);
   PayloadBuffer *pb = new (buffer) PayloadBuffer(4096);