PartitionAlloc: Add MTECheckedPtr buildflag

This change adds a buildflag gating the use of `MTECheckedPtr<T>` and
restores the original implementation of the tags and tag bitmaps.

Bug: 1298696
Change-Id: Ied78574e9d22f0e55e94d2a48376deda118c67ae
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/3495724
Reviewed-by: Bartek Nowierski <bartekn@chromium.org>
Commit-Queue: Kalvin Lee <kdlee@chromium.org>
Cr-Commit-Position: refs/heads/main@{#983525}

base/BUILD.gn

@@ -2054,6 +2054,8 @@ mixed_component("base") {
         "allocator/partition_allocator/partition_root.h",
         "allocator/partition_allocator/partition_stats.cc",
         "allocator/partition_allocator/partition_stats.h",
+        "allocator/partition_allocator/partition_tag.h",
+        "allocator/partition_allocator/partition_tag_bitmap.h",
         "allocator/partition_allocator/partition_tls.h",
         "allocator/partition_allocator/random.cc",
         "allocator/partition_allocator/random.h",

base/allocator/BUILD.gn

@@ -22,6 +22,10 @@ buildflag_header("buildflags") {
   _enable_dangling_raw_ptr_checks =
       enable_dangling_raw_ptr_checks && _use_backup_ref_ptr
 
+  # MTECheckedPtr is exclusive against BRP (asserted at declaration).
+  # MTECheckedPtr requires 64-bit pointers (not available in NaCl).
+  _use_mte_checked_ptr = use_mte_checked_ptr && !is_nacl
+
   _record_alloc_info = false
 
   flags = [
@@ -35,6 +39,9 @@ buildflag_header("buildflags") {
     "ENABLE_DANGLING_RAW_PTR_CHECKS=$_enable_dangling_raw_ptr_checks",
     "PUT_REF_COUNT_IN_PREVIOUS_SLOT=$_put_ref_count_in_previous_slot",
 
+    # Not to be used directly - see `partition_alloc_config.h`.
+    "USE_MTE_CHECKED_PTR=$_use_mte_checked_ptr",
+
     "USE_FAKE_BINARY_EXPERIMENT=$use_fake_binary_experiment",
 
     "RECORD_ALLOC_INFO=$_record_alloc_info",

base/allocator/allocator.gni

@@ -87,8 +87,13 @@ declare_args() {
   # Set use_backup_ref_ptr true to use BackupRefPtr (BRP) as the implementation
   # of raw_ptr<T>, and enable PartitionAlloc support for it.
   use_backup_ref_ptr = _is_brp_supported
+
+  use_mte_checked_ptr = false
 }
 
+assert(!(use_backup_ref_ptr && use_mte_checked_ptr),
+       "MTECheckedPtr conflicts with BRP.")
+
 declare_args() {
   # If BRP is enabled, additional options are available:
   # - put_ref_count_in_previous_slot: place the ref-count at the end of the

base/allocator/partition_allocator/partition_alloc_config.h

@@ -216,4 +216,14 @@ constexpr bool kUseLazyCommit = false;
 #define PA_PREFER_SMALLER_SLOT_SPANS
 #endif  // BUILDFLAG(IS_LINUX)
 
+// Build MTECheckedPtr code.
+//
+// Only applicable to code with 64-bit pointers. Currently conflicts with true
+// hardware MTE.
+#if BUILDFLAG(USE_MTE_CHECKED_PTR) && defined(PA_HAS_64_BITS_POINTERS) && \
+    !defined(PA_HAS_MEMORY_TAGGING)
+#define PA_USE_MTE_CHECKED_PTR_WITH_64_BITS_POINTERS
+#endif  // BUILDFLAG(USE_MTE_CHECKED_PTR) && defined(PA_HAS_64_BITS_POINTERS) &&
+        // !defined(PA_HAS_MEMORY_TAGGING)
+
 #endif  // BASE_ALLOCATOR_PARTITION_ALLOCATOR_PARTITION_ALLOC_CONFIG_H_

base/allocator/partition_allocator/partition_tag.h

@@ -0,0 +1,124 @@
+// 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_PARTITION_TAG_H_
+#define BASE_ALLOCATOR_PARTITION_ALLOCATOR_PARTITION_TAG_H_
+
+// This file defines types and functions for `MTECheckedPtr<T>` (cf.
+// `tagging.h`, which deals with real ARM MTE).
+
+#include <string.h>
+
+#include "base/allocator/buildflags.h"
+#include "base/allocator/partition_allocator/partition_alloc_constants.h"
+#include "base/allocator/partition_allocator/partition_alloc_notreached.h"
+#include "base/allocator/partition_allocator/partition_cookie.h"
+#include "base/allocator/partition_allocator/partition_tag_bitmap.h"
+#include "base/base_export.h"
+#include "build/build_config.h"
+
+namespace partition_alloc::internal {
+
+#if defined(PA_USE_MTE_CHECKED_PTR_WITH_64_BITS_POINTERS)
+
+// Use 8 bits for the partition tag.
+// TODO(tasak): add a description about the partition tag.
+using PartitionTag = uint8_t;
+
+static_assert(
+    sizeof(PartitionTag) == tag_bitmap::kPartitionTagSize,
+    "sizeof(PartitionTag) must be equal to bitmap::kPartitionTagSize.");
+
+static constexpr size_t kInSlotTagBufferSize = 0;
+
+ALWAYS_INLINE PartitionTag* PartitionTagPointer(void* ptr) {
+  // See the comment explaining the layout in partition_tag_bitmap.h.
+  uintptr_t pointer_as_uintptr = reinterpret_cast<uintptr_t>(ptr);
+  uintptr_t bitmap_base =
+      (pointer_as_uintptr & kSuperPageBaseMask) + PartitionPageSize();
+  uintptr_t offset =
+      (pointer_as_uintptr & kSuperPageOffsetMask) - PartitionPageSize();
+  // Not to depend on partition_address_space.h and PartitionAllocGigaCage
+  // feature, use "offset" to see whether the given ptr is_direct_mapped or not.
+  // DirectMap object should cause this PA_DCHECK's failure, as tags aren't
+  // currently supported there.
+  PA_DCHECK(offset >= ReservedTagBitmapSize());
+  size_t bitmap_offset = (offset - ReservedTagBitmapSize()) >>
+                         tag_bitmap::kBytesPerPartitionTagShift
+                             << tag_bitmap::kPartitionTagSizeShift;
+  return reinterpret_cast<PartitionTag* const>(bitmap_base + bitmap_offset);
+}
+
+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, 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);
+}
+
+ALWAYS_INLINE void PartitionTagIncrementValue(void* ptr, size_t size) {
+  PartitionTag tag = PartitionTagGetValue(ptr);
+  PartitionTag new_tag = tag;
+  ++new_tag;
+  new_tag += !new_tag;  // Avoid 0.
+#if DCHECK_IS_ON()
+  // This verifies that tags for the entire slot have the same value and that
+  // |size| doesn't exceed the slot size.
+  size_t tag_count = size >> tag_bitmap::kBytesPerPartitionTagShift;
+  PartitionTag* tag_ptr = PartitionTagPointer(ptr);
+  while (tag_count-- > 0) {
+    PA_DCHECK(tag == *tag_ptr);
+    tag_ptr++;
+  }
+#endif
+  PartitionTagSetValue(ptr, size, new_tag);
+}
+
+#else  // No-op versions
+
+using PartitionTag = uint8_t;
+
+static constexpr size_t kInSlotTagBufferSize = 0;
+
+ALWAYS_INLINE PartitionTag* PartitionTagPointer(void* ptr) {
+  PA_NOTREACHED();
+  return nullptr;
+}
+
+ALWAYS_INLINE void PartitionTagSetValue(void*, size_t, PartitionTag) {}
+
+ALWAYS_INLINE PartitionTag PartitionTagGetValue(void*) {
+  return 0;
+}
+
+ALWAYS_INLINE void PartitionTagClearValue(void* ptr, size_t) {}
+
+ALWAYS_INLINE void PartitionTagIncrementValue(void* ptr, size_t size) {}
+
+#endif  // defined(PA_USE_MTE_CHECKED_PTR_WITH_64_BITS_POINTERS)
+
+constexpr size_t kPartitionTagSizeAdjustment = kInSlotTagBufferSize;
+constexpr size_t kPartitionTagOffsetAdjustment = kInSlotTagBufferSize;
+
+}  // namespace partition_alloc::internal
+
+#endif  // BASE_ALLOCATOR_PARTITION_ALLOCATOR_PARTITION_TAG_H_

base/allocator/partition_allocator/partition_tag_bitmap.h

@@ -0,0 +1,145 @@
+// 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_PARTITION_TAG_BITMAP_H_
+#define BASE_ALLOCATOR_PARTITION_ALLOCATOR_PARTITION_TAG_BITMAP_H_
+
+#include "base/allocator/buildflags.h"
+#include "base/allocator/partition_allocator/page_allocator_constants.h"
+#include "base/allocator/partition_allocator/partition_alloc_constants.h"
+
+namespace partition_alloc::internal {
+
+#if defined(PA_USE_MTE_CHECKED_PTR_WITH_64_BITS_POINTERS)
+
+namespace tag_bitmap {
+// kPartitionTagSize should be equal to sizeof(PartitionTag).
+// PartitionTag is defined in partition_tag.h and static_assert there
+// checks the condition.
+static constexpr size_t kPartitionTagSizeShift = 0;
+static constexpr size_t kPartitionTagSize = 1U << kPartitionTagSizeShift;
+
+static constexpr size_t kBytesPerPartitionTagShift = 4;
+// One partition tag is assigned per |kBytesPerPartitionTag| bytes in the slot
+// spans.
+//  +-----------+ 0
+//  |           |  ====> 1 partition tag
+//  +-----------+ kBytesPerPartitionTag
+//  |           |  ====> 1 partition tag
+//  +-----------+ 2*kBytesPerPartitionTag
+// ...
+//  +-----------+ slot_size
+static constexpr size_t kBytesPerPartitionTag = 1U
+                                                << kBytesPerPartitionTagShift;
+static_assert(
+    kMinBucketedOrder >= kBytesPerPartitionTagShift + 1,
+    "MTECheckedPtr requires kBytesPerPartitionTagShift-bytes alignment.");
+
+static constexpr size_t kBytesPerPartitionTagRatio =
+    kBytesPerPartitionTag / kPartitionTagSize;
+
+static_assert(kBytesPerPartitionTag > 0,
+              "kBytesPerPartitionTag should be larger than 0");
+static_assert(
+    kBytesPerPartitionTag % kPartitionTagSize == 0,
+    "kBytesPerPartitionTag should be multiples of sizeof(PartitionTag).");
+
+constexpr size_t CeilCountOfUnits(size_t size, size_t unit_size) {
+  return (size + unit_size - 1) / unit_size;
+}
+
+}  // namespace tag_bitmap
+
+// kTagBitmapSize is calculated in the following way:
+// (1) kSuperPageSize - 2 * PartitionPageSize() = kTagBitmapSize +
+// SlotSpanSize()
+// (2) kTagBitmapSize >= SlotSpanSize() / kBytesPerPartitionTag *
+// sizeof(PartitionTag)
+//--
+// (1)' SlotSpanSize() = kSuperPageSize - 2 * PartitionPageSize() -
+// kTagBitmapSize
+// (2)' SlotSpanSize() <= kTagBitmapSize * Y
+// (3)' Y = kBytesPerPartitionTag / sizeof(PartitionTag) =
+// kBytesPerPartitionTagRatio
+//
+//   kTagBitmapSize * Y >= kSuperPageSize - 2 * PartitionPageSize() -
+//   kTagBitmapSize (1 + Y) * kTagBimapSize >= kSuperPageSize - 2 *
+//   PartitionPageSize()
+// Finally,
+//   kTagBitmapSize >= (kSuperPageSize - 2 * PartitionPageSize()) / (1 + Y)
+PAGE_ALLOCATOR_CONSTANTS_DECLARE_CONSTEXPR ALWAYS_INLINE size_t
+NumPartitionPagesPerTagBitmap() {
+  return tag_bitmap::CeilCountOfUnits(
+      kSuperPageSize / PartitionPageSize() - 2,
+      tag_bitmap::kBytesPerPartitionTagRatio + 1);
+}
+
+// To make guard pages between the tag bitmap and the slot span, calculate the
+// number of SystemPages of TagBitmap. If kNumSystemPagesPerTagBitmap *
+// SystemPageSize() < kTagBitmapSize, guard pages will be created. (c.f. no
+// guard pages if sizeof(PartitionTag) == 2.)
+PAGE_ALLOCATOR_CONSTANTS_DECLARE_CONSTEXPR ALWAYS_INLINE size_t
+NumSystemPagesPerTagBitmap() {
+  return tag_bitmap::CeilCountOfUnits(
+      kSuperPageSize / SystemPageSize() -
+          2 * PartitionPageSize() / SystemPageSize(),
+      tag_bitmap::kBytesPerPartitionTagRatio + 1);
+}
+
+PAGE_ALLOCATOR_CONSTANTS_DECLARE_CONSTEXPR ALWAYS_INLINE size_t
+ActualTagBitmapSize() {
+  return NumSystemPagesPerTagBitmap() * SystemPageSize();
+}
+
+// PartitionPageSize-aligned tag bitmap size.
+PAGE_ALLOCATOR_CONSTANTS_DECLARE_CONSTEXPR ALWAYS_INLINE size_t
+ReservedTagBitmapSize() {
+  return PartitionPageSize() * NumPartitionPagesPerTagBitmap();
+}
+
+#if PAGE_ALLOCATOR_CONSTANTS_ARE_CONSTEXPR
+static_assert(ActualTagBitmapSize() <= ReservedTagBitmapSize(),
+              "kActualTagBitmapSize should be smaller than or equal to "
+              "kReservedTagBitmapSize.");
+static_assert(ReservedTagBitmapSize() - ActualTagBitmapSize() <
+                  PartitionPageSize(),
+              "Unused space in the tag bitmap should be smaller than "
+              "PartitionPageSize()");
+
+// The region available for slot spans is the reminder of the super page, after
+// taking away the first and last partition page (for metadata and guard pages)
+// and partition pages reserved for the tag bitmap.
+PAGE_ALLOCATOR_CONSTANTS_DECLARE_CONSTEXPR ALWAYS_INLINE size_t
+SlotSpansSize() {
+  return kSuperPageSize - 2 * PartitionPageSize() - ReservedTagBitmapSize();
+}
+
+static_assert(ActualTagBitmapSize() * tag_bitmap::kBytesPerPartitionTagRatio >=
+                  SlotSpansSize(),
+              "bitmap is large enough to cover slot spans");
+static_assert((ActualTagBitmapSize() - PartitionPageSize()) *
+                      tag_bitmap::kBytesPerPartitionTagRatio <
+                  SlotSpansSize(),
+              "any smaller bitmap wouldn't suffice to cover slots spans");
+#endif  // PAGE_ALLOCATOR_CONSTANTS_ARE_CONSTEXPR
+
+#else
+
+constexpr ALWAYS_INLINE size_t NumPartitionPagesPerTagBitmap() {
+  return 0;
+}
+
+constexpr ALWAYS_INLINE size_t ActualTagBitmapSize() {
+  return 0;
+}
+
+constexpr ALWAYS_INLINE size_t ReservedTagBitmapSize() {
+  return 0;
+}
+
+#endif  // defined(PA_USE_MTE_CHECKED_PTR_WITH_64_BITS_POINTERS)
+
+}  // namespace partition_alloc::internal
+
+#endif  // BASE_ALLOCATOR_PARTITION_ALLOCATOR_PARTITION_TAG_BITMAP_H_