Fix edge case for MicroRep allocation.

Because of the interaction between Ceil and the rounding down we could end up
with a `delete` that used less bytes than the `new`.

Changed the test to not use an arena to exercise the `delete` path.

PiperOrigin-RevId: 761087000

Author: sbenzaquen

src/google/protobuf/micro_string.cc

@@ -113,10 +113,13 @@ MicroString::MicroRep* MicroString::AllocateMicroRep(size_t size,
   MicroRep* h;
   size_t capacity = size;
   if (arena == nullptr) {
-    const internal::SizedPtr alloc = internal::AllocateAtLeast(
-        ArenaAlignDefault::Ceil(MicroRepSize(capacity)));
+    size_t requested_size = ArenaAlignDefault::Ceil(MicroRepSize(capacity));
+    const internal::SizedPtr alloc = internal::AllocateAtLeast(requested_size);
     // Maybe we rounded up too much.
     capacity = std::min(kMaxMicroRepCapacity, alloc.n - sizeof(MicroRep));
+    // Verify that the size we are going to free later is at least what we asked
+    // for.
+    ABSL_DCHECK_LE(requested_size, MicroRepSize(capacity));
     h = reinterpret_cast<MicroRep*>(alloc.p);
   } else {
     capacity =

src/google/protobuf/micro_string.h

@@ -8,6 +8,7 @@
 #ifndef GOOGLE_PROTOBUF_MICRO_STRING_H__
 #define GOOGLE_PROTOBUF_MICRO_STRING_H__
 
+#include <cstddef>
 #include <cstdint>
 
 #include "absl/base/config.h"
@@ -97,12 +98,34 @@ class PROTOBUF_EXPORT MicroString {
     }
   };
 
+  struct MicroRep {
+    uint8_t size;
+    uint8_t capacity;
+
+    char* data() { return reinterpret_cast<char*>(this + 1); }
+    const char* data() const { return reinterpret_cast<const char*>(this + 1); }
+    absl::string_view view() const { return {data(), size}; }
+
+    void SetInitialSize(uint8_t size) {
+      PoisonMemoryRegion(data() + size, capacity - size);
+      this->size = size;
+    }
+
+    void Unpoison() { UnpoisonMemoryRegion(data(), capacity); }
+
+    void ChangeSize(uint8_t new_size) {
+      PoisonMemoryRegion(data() + new_size, capacity - new_size);
+      UnpoisonMemoryRegion(data(), new_size);
+      size = new_size;
+    }
+  };
+
  public:
   // We don't allow extra capacity in big-endian because it is harder to manage
   // the pointer to the MicroString "base".
   static constexpr bool kAllowExtraCapacity = IsLittleEndian();
   static constexpr size_t kInlineCapacity = sizeof(uintptr_t) - 1;
-  static constexpr size_t kMaxMicroRepCapacity = 255;
+  static constexpr size_t kMaxMicroRepCapacity = 256 - sizeof(MicroRep);
 
   // Empty string.
   constexpr MicroString() : rep_() {}
@@ -311,27 +334,6 @@ class PROTOBUF_EXPORT MicroString {
     return cap >= kOwned ? kOwned : static_cast<LargeRepKind>(cap);
   }
 
-  struct MicroRep {
-    uint8_t size;
-    uint8_t capacity;
-
-    char* data() { return reinterpret_cast<char*>(this + 1); }
-    const char* data() const { return reinterpret_cast<const char*>(this + 1); }
-    absl::string_view view() const { return {data(), size}; }
-
-    void SetInitialSize(uint8_t size) {
-      PoisonMemoryRegion(data() + size, capacity - size);
-      this->size = size;
-    }
-
-    void Unpoison() { UnpoisonMemoryRegion(data(), capacity); }
-
-    void ChangeSize(uint8_t new_size) {
-      PoisonMemoryRegion(data() + new_size, capacity - new_size);
-      UnpoisonMemoryRegion(data(), new_size);
-      size = new_size;
-    }
-  };
   // Micro-optimization: by using kIsMicroRepTag as 2, the MicroRep `rep_`
   // pointer (with the tag) is already pointing into the data buffer.
   static_assert(sizeof(MicroRep) == kIsMicroRepTag);

src/google/protobuf/micro_string_test.cc

@@ -319,19 +319,25 @@ TEST(MicroStringTest, CapacityIsRoundedUpOnHeap) {
 }
 
 TEST(MicroStringTest, CapacityRoundingUpStaysWithinBoundsForMicroRep) {
-  Arena arena;
-  MicroString str;
+  const auto get_capacity_for_size = [&](size_t size) {
+    MicroString str;
+    std::string input = std::string(size, 'x');
+    str.Set(input, nullptr);
+    EXPECT_EQ(str.Get(), input);
+    size_t cap = str.Capacity();
+    str.Destroy();
+    return cap;
+  };
 
-  std::string input = std::string(200, 'x');
-  str.Set(input, &arena);
-  EXPECT_EQ(str.Capacity(), 208 - kMicroRepSize);
-  EXPECT_EQ(str.Get(), input);
+  EXPECT_EQ(get_capacity_for_size(200), 208 - kMicroRepSize);
 
-  input = std::string(255, 'x');
-  str.Set(input, &arena);
-  // It caps at 255 even though the allocated block is larger.
-  EXPECT_EQ(str.Capacity(), 255);
-  EXPECT_EQ(str.Get(), input);
+  // These are in the boundary
+  EXPECT_EQ(get_capacity_for_size(253), 256 - kMicroRepSize);
+  // This is the maximum capacity for MicroRep
+  EXPECT_EQ(get_capacity_for_size(254), 256 - kMicroRepSize);
+
+  // This one jumps to LargeRep
+  EXPECT_GE(get_capacity_for_size(255), 256);
 }
 
 TEST(MicroStringTest, PoisonsTheUnusedCapacity) {