diff --git a/toolbelt/payload_buffer.cc b/toolbelt/payload_buffer.cc index 0e249d7..b2a75b1 100644 --- a/toolbelt/payload_buffer.cc +++ b/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(old_str); + const char *old_end = + old_begin + sizeof(uint32_t) + *static_cast(old_str); + if (std::less_equal()(old_begin, s) && + std::less()(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(str); p[0] = uint32_t(len); - memcpy(p + 1, s, len); + memcpy(p + 1, source == 0 ? s : reinterpret_cast(*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(free_block->next); - FreeBlockHeader *newb = reinterpret_cast( - 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(free_block); *len_ptr = new_length; @@ -615,6 +625,18 @@ uint64_t *PayloadBuffer::MergeWithFreeBlockBelow( memset(reinterpret_cast(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( + 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; } diff --git a/toolbelt/payload_buffer.h b/toolbelt/payload_buffer.h index 3feca78..d94ee35 100644 --- a/toolbelt/payload_buffer.h +++ b/toolbelt/payload_buffer.h @@ -471,6 +471,13 @@ struct PayloadBuffer { reinterpret_cast(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(p); + const uintptr_t start = reinterpret_cast(this); + return at > start && at - start < full_size ? BufferOffset(at - start) : 0; + } + void InsertNewFreeBlockAtEnd(FreeBlockHeader *free_block, FreeBlockHeader *prev, uint32_t length); diff --git a/toolbelt/payload_buffer_test.cc b/toolbelt/payload_buffer_test.cc index 357a771..305e4b6 100644 --- a/toolbelt/payload_buffer_test.cc +++ b/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(PayloadBuffer::Allocate(&pb, 256, true, false)); + auto *upper = + static_cast(PayloadBuffer::Allocate(&pb, 256, true, false)); + for (uint32_t i = 0; i < 256; ++i) { + block[i] = static_cast(i); + } + memset(upper, 0xc5, 256); + const uint64_t upper_size = *(reinterpret_cast(upper) - 1); + pb->Free(lower); + + auto *grown = static_cast( + 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(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(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(realloc(*p, new_size)); +#if defined(__clang__) +#pragma clang diagnostic pop +#elif defined(__GNUC__) +#pragma GCC diagnostic pop +#endif + buffer = reinterpret_cast(*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('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);