Initialize partition tags for the entire slot (for MTECheckedPtr).

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

base/BUILD.gn

@@ -1747,6 +1747,7 @@ jumbo_component("base") {
         "allocator/partition_allocator/address_pool_manager.h",
         "allocator/partition_allocator/address_space_randomization.cc",
         "allocator/partition_allocator/address_space_randomization.h",
+        "allocator/partition_allocator/checked_ptr_support.h",
         "allocator/partition_allocator/memory_reclaimer.cc",
         "allocator/partition_allocator/memory_reclaimer.h",
         "allocator/partition_allocator/oom.h",

base/allocator/partition_allocator/checked_ptr_support.h

+// Copyright (c) 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef BASE_ALLOCATOR_PARTITION_ALLOCATOR_CHECKED_PTR_SUPPORT_H_
+#define BASE_ALLOCATOR_PARTITION_ALLOCATOR_CHECKED_PTR_SUPPORT_H_
+
+#define ENABLE_TAG_FOR_CHECKED_PTR2 0
+#define ENABLE_TAG_FOR_MTE_CHECKED_PTR 0
+
+#endif  // BASE_ALLOCATOR_PARTITION_ALLOCATOR_CHECKED_PTR_SUPPORT_H_

base/allocator/partition_allocator/partition_alloc.cc

@@ -71,8 +71,13 @@ static_assert(kGenericMaxDirectMapped <=
                   (1UL << 31) + kPageAllocationGranularity,
               "maximum direct mapped allocation");
 // Check that some of our zanier calculations worked out as expected.
+#if !ENABLE_TAG_FOR_MTE_CHECKED_PTR
 static_assert(kGenericSmallestBucket == alignof(std::max_align_t),
               "generic smallest bucket");
+#else
+static_assert(kGenericSmallestBucket >= alignof(std::max_align_t),
+              "generic smallest bucket");
+#endif
 static_assert(kGenericMaxBucketed == 983040, "generic max bucketed");
 static_assert(kMaxSystemPagesPerSlotSpan < (1 << 8),
               "System pages per slot span must be less than 128.");

base/allocator/partition_allocator/partition_alloc.h

@@ -598,7 +598,8 @@ ALWAYS_INLINE void* PartitionRoot<thread_safe>::AllocFromBucket(Bucket* bucket,
   }
 
   if (allow_extras && !bucket->is_direct_mapped()) {
-    internal::PartitionTagSetValue(ret, GetNewPartitionTag());
+    internal::PartitionTagSetValue(ret, page->bucket->slot_size,
+                                   GetNewPartitionTag());
   }
 
   return ret;
@@ -625,7 +626,7 @@ ALWAYS_INLINE void PartitionRoot<thread_safe>::Free(void* ptr) {
   auto* root = PartitionRoot<thread_safe>::FromPage(page);
   if (root->allow_extras && !page->bucket->is_direct_mapped()) {
     // TODO(tasak): clear partition tag. Temporarily set the tag to be 0.
-    internal::PartitionTagClearValue(ptr);
+    internal::PartitionTagClearValue(ptr, page->bucket->slot_size);
   }
   ptr = internal::PartitionPointerAdjustSubtract(root->allow_extras, ptr);
   internal::DeferredUnmap deferred_unmap;

base/allocator/partition_allocator/partition_alloc_constants.h

@@ -8,6 +8,7 @@
 #include <limits.h>
 #include <cstddef>
 
+#include "base/allocator/partition_allocator/checked_ptr_support.h"
 #include "base/allocator/partition_allocator/page_allocator_constants.h"
 
 #include "build/build_config.h"
@@ -152,8 +153,13 @@ static_assert(alignof(std::max_align_t) <= 16,
 //
 // In practice, this means 8 bytes alignment on 32 bit architectures, and 16
 // bytes on 64 bit ones.
+#if ENABLE_TAG_FOR_MTE_CHECKED_PTR
+// MTECheckedPtr requires 16B-alignment because kBytesPerPartitionTag is 16.
+static const size_t kGenericMinBucketedOrder = 5;
+#else
 static const size_t kGenericMinBucketedOrder =
     alignof(std::max_align_t) == 16 ? 5 : 4;  // 2^(order - 1), that is 16 or 8.
+#endif
 // The largest bucketed order is 1 << (20 - 1), storing [512 KiB, 1 MiB):
 static const size_t kGenericMaxBucketedOrder = 20;
 static const size_t kGenericNumBucketedOrders =

base/allocator/partition_allocator/partition_alloc_unittest.cc

@@ -2515,9 +2515,9 @@ TEST_F(PartitionAllocTest, TagBasic) {
   constexpr PartitionTag kTag1 = static_cast<PartitionTag>(0xBADA);
   constexpr PartitionTag kTag2 = static_cast<PartitionTag>(0xDB8A);
   constexpr PartitionTag kTag3 = static_cast<PartitionTag>(0xA3C4);
-  PartitionTagSetValue(ptr1, kTag1);
-  PartitionTagSetValue(ptr2, kTag2);
-  PartitionTagSetValue(ptr3, kTag3);
+  PartitionTagSetValue(ptr1, page->bucket->slot_size, kTag1);
+  PartitionTagSetValue(ptr2, page->bucket->slot_size, kTag2);
+  PartitionTagSetValue(ptr3, page->bucket->slot_size, kTag3);
 
   memset(ptr1, 0, kTestAllocSize);
   memset(ptr2, 0, kTestAllocSize);

base/allocator/partition_allocator/partition_tag.h

@@ -5,6 +5,8 @@
 #ifndef BASE_ALLOCATOR_PARTITION_ALLOCATOR_PARTITION_TAG_H_
 #define BASE_ALLOCATOR_PARTITION_ALLOCATOR_PARTITION_TAG_H_
 
+#include <string.h>
+
 #include "base/allocator/partition_allocator/partition_alloc_constants.h"
 #include "base/allocator/partition_allocator/partition_cookie.h"
 #include "base/allocator/partition_allocator/partition_tag_bitmap.h"
@@ -61,7 +63,7 @@ ALWAYS_INLINE void* PartitionTagPointerAdjustAdd(void* ptr) {
                                  kInSlotTagBufferSize);
 }
 
-ALWAYS_INLINE void PartitionTagSetValue(void* ptr, PartitionTag value) {
+ALWAYS_INLINE void PartitionTagSetValue(void* ptr, size_t, PartitionTag value) {
   *PartitionTagPointer(ptr) = value;
 }
 
@@ -69,7 +71,7 @@ ALWAYS_INLINE PartitionTag PartitionTagGetValue(void* ptr) {
   return *PartitionTagPointer(ptr);
 }
 
-ALWAYS_INLINE void PartitionTagClearValue(void* ptr) {
+ALWAYS_INLINE void PartitionTagClearValue(void* ptr, size_t) {
   PA_DCHECK(PartitionTagGetValue(ptr));
   *PartitionTagPointer(ptr) = 0;
 }
@@ -122,20 +124,32 @@ ALWAYS_INLINE void* PartitionTagPointerAdjustAdd(void* ptr) {
   return ptr;
 }
 
-ALWAYS_INLINE void PartitionTagSetValue(void* ptr, PartitionTag value) {
-  *PartitionTagPointer(ptr) = value;
+ALWAYS_INLINE void PartitionTagSetValue(void* ptr,
+                                        size_t size,
+                                        PartitionTag value) {
+  PA_DCHECK((size % tag_bitmap::kBytesPerPartitionTag) == 0);
+  size_t tag_count = size >> tag_bitmap::kBytesPerPartitionTagShift;
+  PartitionTag* tag_ptr = PartitionTagPointer(ptr);
+  if (sizeof(PartitionTag) == 1) {
+    memset(tag_ptr, value, tag_count);
+  } else {
+    while (tag_count-- > 0)
+      *tag_ptr++ = value;
+  }
 }
 
 ALWAYS_INLINE PartitionTag PartitionTagGetValue(void* ptr) {
   return *PartitionTagPointer(ptr);
 }
 
-ALWAYS_INLINE void PartitionTagClearValue(void* ptr) {
-  PA_DCHECK(PartitionTagGetValue(ptr));
-  *PartitionTagPointer(ptr) = 0;
+ALWAYS_INLINE void PartitionTagClearValue(void* ptr, size_t size) {
+  size_t tag_region_size = size >> tag_bitmap::kBytesPerPartitionTagShift
+                                       << tag_bitmap::kPartitionTagSizeShift;
+  PA_DCHECK(!memchr(PartitionTagPointer(ptr), 0, tag_region_size));
+  memset(PartitionTagPointer(ptr), 0, tag_region_size);
 }
 
-#else  // !ENABLE_TAG_FOR_CHECKED_PTR2
+#else  // !ENABLE_TAG_FOR_CHECKED_PTR2 && !ENABLE_TAG_FOR_MTE_CHECKED_PTR
 
 using PartitionTag = uint16_t;
 
@@ -163,13 +177,13 @@ ALWAYS_INLINE void* PartitionTagPointerAdjustAdd(void* ptr) {
   return ptr;
 }
 
-ALWAYS_INLINE void PartitionTagSetValue(void*, PartitionTag) {}
+ALWAYS_INLINE void PartitionTagSetValue(void*, size_t, PartitionTag) {}
 
 ALWAYS_INLINE PartitionTag PartitionTagGetValue(void*) {
   return 0;
 }
 
-ALWAYS_INLINE void PartitionTagClearValue(void* ptr) {}
+ALWAYS_INLINE void PartitionTagClearValue(void* ptr, size_t) {}
 
 #endif  // !ENABLE_TAG_FOR_CHECKED_PTR2
 

base/allocator/partition_allocator/partition_tag_bitmap.h

@@ -5,17 +5,14 @@
 #ifndef BASE_ALLOCATOR_PARTITION_ALLOCATOR_PARTITION_TAG_BITMAP_H_
 #define BASE_ALLOCATOR_PARTITION_ALLOCATOR_PARTITION_TAG_BITMAP_H_
 
+#include "base/allocator/partition_allocator/checked_ptr_support.h"
 #include "base/allocator/partition_allocator/partition_alloc_constants.h"
 
-#define ENABLE_TAG_FOR_CHECKED_PTR2 0
-#define ENABLE_TAG_FOR_MTE_CHECKED_PTR 0
-
 namespace base {
 
 namespace internal {
 
 #if ENABLE_TAG_FOR_MTE_CHECKED_PTR
-#undef ENABLE_TAG_FOR_CHECKED_PTR2
 
 // Normal bucket layout
 // +----------------+ super_page_base
@@ -62,6 +59,9 @@ static constexpr size_t kBytesPerPartitionTagShift = 4;
 //  +-----------+ slot_size
 static constexpr size_t kBytesPerPartitionTag = 1U
                                                 << kBytesPerPartitionTagShift;
+static_assert(
+    kGenericMinBucketedOrder >= kBytesPerPartitionTagShift + 1,
+    "MTECheckedPtr requires kBytesPerPartitionTagShift-bytes alignment.");
 
 static constexpr size_t kBytesPerPartitionTagRatio =
     kBytesPerPartitionTag / kPartitionTagSize;