Make only the region of the tag bitmap for the requested slots accessible.

Bug: 1092288
Change-Id: I0b0f32d16a01128087fecf2c6b1b3b32ca07d4ba
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2299941
Commit-Queue: Takashi Sakamoto <tasak@google.com>
Reviewed-by: Kentaro Hara <haraken@chromium.org>
Reviewed-by: Bartek Nowierski <bartekn@chromium.org>
Cr-Commit-Position: refs/heads/master@{#789882}

base/allocator/partition_allocator/partition_alloc.h

@@ -419,6 +419,9 @@ struct BASE_EXPORT PartitionRoot {
 #if ENABLE_TAG_FOR_CHECKED_PTR2 || ENABLE_TAG_FOR_MTE_CHECKED_PTR
   internal::PartitionTag current_partition_tag = 0;
 #endif
+#if ENABLE_TAG_FOR_MTE_CHECKED_PTR
+  char* next_tag_bitmap_page = nullptr;
+#endif
 
   // Some pre-computed constants.
   size_t order_index_shifts[kBitsPerSizeT + 1] = {};
@@ -566,15 +569,17 @@ ALWAYS_INLINE void* PartitionRoot<thread_safe>::AllocFromBucket(Bucket* bucket,
     PA_DCHECK(raw_size == size);
     new_slot_size = raw_size;
   }
-  // Layout inside the slot: |tag|cookie|object|[empty]|cookie|
-  //                                    <--a--->
-  //                                    <------b------->
-  //                         <---------------c---------------->
+  // Layout inside the slot: |[tag]|cookie|object|[empty]|cookie|
+  //                                      <--a--->
+  //                                      <------b------->
+  //                         <-----------------c---------------->
   //   a: size
   //   b: size_with_no_extras
   //   c: new_slot_size
   // Note, empty space occurs if the slot size is larger than needed to
   // accommodate the request.
+  // The tag may or may not exist in the slot, depending on CheckedPtr
+  // implementation.
   size_t size_with_no_extras =
       internal::PartitionSizeAdjustSubtract(allow_extras, new_slot_size);
   // The value given to the application is just after the tag and cookie.
@@ -600,7 +605,9 @@ ALWAYS_INLINE void* PartitionRoot<thread_safe>::AllocFromBucket(Bucket* bucket,
   }
 
   if (allow_extras && !bucket->is_direct_mapped()) {
-    internal::PartitionTagSetValue(ret, page->bucket->slot_size,
+    size_t slot_size_with_no_extras = internal::PartitionSizeAdjustSubtract(
+        allow_extras, page->bucket->slot_size);
+    internal::PartitionTagSetValue(ret, slot_size_with_no_extras,
                                    GetNewPartitionTag());
   }
 
@@ -627,8 +634,10 @@ ALWAYS_INLINE void PartitionRoot<thread_safe>::Free(void* ptr) {
   PA_DCHECK(IsValidPage(page));
   auto* root = PartitionRoot<thread_safe>::FromPage(page);
   if (root->allow_extras && !page->bucket->is_direct_mapped()) {
+    size_t size_with_no_extras = internal::PartitionSizeAdjustSubtract(
+        root->allow_extras, page->bucket->slot_size);
     // TODO(tasak): clear partition tag. Temporarily set the tag to be 0.
-    internal::PartitionTagClearValue(ptr, page->bucket->slot_size);
+    internal::PartitionTagClearValue(ptr, size_with_no_extras);
   }
   ptr = internal::PartitionPointerAdjustSubtract(root->allow_extras, ptr);
   internal::DeferredUnmap deferred_unmap;

base/allocator/partition_allocator/partition_alloc_unittest.cc

@@ -2528,21 +2528,19 @@ TEST_F(PartitionAllocTest, TagBasic) {
   PartitionTagSetValue(ptr2, page->bucket->slot_size, kTag2);
   PartitionTagSetValue(ptr3, page->bucket->slot_size, kTag3);
 
-  memset(ptr1, 0, kTestAllocSize);
-  memset(ptr2, 0, kTestAllocSize);
-  memset(ptr3, 0, kTestAllocSize);
+  memset(ptr1, 0, alloc_size);
+  memset(ptr2, 0, alloc_size);
+  memset(ptr3, 0, alloc_size);
 
   EXPECT_EQ(kTag1, PartitionTagGetValue(ptr1));
   EXPECT_EQ(kTag2, PartitionTagGetValue(ptr2));
   EXPECT_EQ(kTag3, PartitionTagGetValue(ptr3));
 
-  EXPECT_TRUE(!memchr(ptr1, static_cast<uint8_t>(kTag1), kTestAllocSize));
-  EXPECT_TRUE(!memchr(ptr2, static_cast<uint8_t>(kTag2), kTestAllocSize));
+  EXPECT_TRUE(!memchr(ptr1, static_cast<uint8_t>(kTag1), alloc_size));
+  EXPECT_TRUE(!memchr(ptr2, static_cast<uint8_t>(kTag2), alloc_size));
   if (sizeof(PartitionTag) > 1) {
-    EXPECT_TRUE(
-        !memchr(ptr1, static_cast<uint8_t>(kTag1 >> 8), kTestAllocSize));
-    EXPECT_TRUE(
-        !memchr(ptr2, static_cast<uint8_t>(kTag2 >> 8), kTestAllocSize));
+    EXPECT_TRUE(!memchr(ptr1, static_cast<uint8_t>(kTag1 >> 8), alloc_size));
+    EXPECT_TRUE(!memchr(ptr2, static_cast<uint8_t>(kTag2 >> 8), alloc_size));
   }
 
   allocator.root()->Free(ptr1);
@@ -2583,8 +2581,9 @@ TEST_F(PartitionAllocTest, TagForDirectMap) {
 
   constexpr PartitionTag kTag1 = 0xBADA;
   constexpr PartitionTag kTag2 = 0xDB8A;
-  PartitionTagSetValue(ptr1, kTag1);
-  PartitionTagSetValue(ptr2, kTag2);
+  constexpr size_t unused_size = 0;
+  PartitionTagSetValue(ptr1, unused_size, kTag1);
+  PartitionTagSetValue(ptr2, unused_size, kTag2);
 
   memset(ptr1, 0, request_size);
   memset(ptr2, 0, request_size);

base/allocator/partition_allocator/partition_bucket.cc

@@ -239,6 +239,27 @@ ALWAYS_INLINE void* PartitionBucket<thread_safe>::AllocNewSlotSpan(
 
     root->next_partition_page += total_size;
     root->IncreaseCommittedPages(total_size);
+
+#if ENABLE_TAG_FOR_MTE_CHECKED_PTR
+    PA_DCHECK(root->next_tag_bitmap_page);
+    char* next_tag_bitmap_page = reinterpret_cast<char*>(
+        bits::Align(reinterpret_cast<uintptr_t>(
+                        PartitionTagPointer(root->next_partition_page)),
+                    kSystemPageSize));
+    if (root->next_tag_bitmap_page < next_tag_bitmap_page) {
+#if DCHECK_IS_ON()
+      char* super_page = reinterpret_cast<char*>(
+          reinterpret_cast<uintptr_t>(ret) & kSuperPageBaseMask);
+      char* tag_bitmap = super_page + kPartitionPageSize;
+      PA_DCHECK(next_tag_bitmap_page <= tag_bitmap + kActualTagBitmapSize);
+      PA_DCHECK(next_tag_bitmap_page > tag_bitmap);
+#endif
+      SetSystemPagesAccess(root->next_tag_bitmap_page,
+                           next_tag_bitmap_page - root->next_tag_bitmap_page,
+                           PageReadWrite);
+      root->next_tag_bitmap_page = next_tag_bitmap_page;
+    }
+#endif
     return ret;
   }
 
@@ -289,12 +310,22 @@ ALWAYS_INLINE void* PartitionBucket<thread_safe>::AllocNewSlotSpan(
   SetSystemPagesAccess(super_page + (kSystemPageSize * 2),
                        kPartitionPageSize - (kSystemPageSize * 2),
                        PageInaccessible);
-  if (kActualTagBitmapSize < kReservedTagBitmapSize) {
-    // Make guard pages between tag bitmap and the first slotspan if possible.
-    SetSystemPagesAccess(tag_bitmap + kActualTagBitmapSize,
-                         kReservedTagBitmapSize - kActualTagBitmapSize,
-                         PageInaccessible);
-  }
+#if ENABLE_TAG_FOR_MTE_CHECKED_PTR
+  // Make the first |total_size| region of the tag bitmap accessible.
+  // The rest of the region is set to inaccessible.
+  char* next_tag_bitmap_page = reinterpret_cast<char*>(
+      bits::Align(reinterpret_cast<uintptr_t>(
+                      PartitionTagPointer(root->next_partition_page)),
+                  kSystemPageSize));
+  PA_DCHECK(next_tag_bitmap_page <= tag_bitmap + kActualTagBitmapSize);
+  PA_DCHECK(next_tag_bitmap_page > tag_bitmap);
+  // |ret| points at the end of the tag bitmap.
+  PA_DCHECK(next_tag_bitmap_page <= ret);
+  SetSystemPagesAccess(next_tag_bitmap_page, ret - next_tag_bitmap_page,
+                       PageInaccessible);
+  root->next_tag_bitmap_page = next_tag_bitmap_page;
+#endif
+
   //  SetSystemPagesAccess(super_page + (kSuperPageSize -
   //  kPartitionPageSize),
   //                             kPartitionPageSize, PageInaccessible);