Set partition tags at an initial allocation and free for MTECheckedPtr.

This is a sub-variant of ENABLE_TAG_FOR_MTE_CHECKED_PTR. At first, the tags in the bitmap have to be initialized for an entire slot span upon allocation (rare operation) and then subsequently set to a new value for an entire slot (not slot span) every time that slot is freed.

Bug: 1092288
Change-Id: I0c2f06ab580e941f3db1a4c51dfa8e4617acd1e8
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2308554
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@{#791722}

base/allocator/partition_allocator/checked_ptr_support.h

@@ -9,4 +9,7 @@
 #define ENABLE_TAG_FOR_MTE_CHECKED_PTR 0
 #define ENABLE_TAG_FOR_SINGLE_TAG_CHECKED_PTR 0
 
+// This is a sub-variant of ENABLE_TAG_FOR_MTE_CHECKED_PTR
+#define MTE_CHECKED_PTR_SET_TAG_AT_FREE 1
+
 #endif  // BASE_ALLOCATOR_PARTITION_ALLOCATOR_CHECKED_PTR_SUPPORT_H_

base/allocator/partition_allocator/partition_alloc.h

@@ -74,7 +74,6 @@
 #include "base/no_destructor.h"
 #include "base/notreached.h"
 #include "base/partition_alloc_buildflags.h"
-#include "base/rand_util.h"
 #include "base/stl_util.h"
 #include "base/synchronization/lock.h"
 #include "base/sys_byteorder.h"
@@ -504,14 +503,6 @@ struct BASE_EXPORT PartitionRoot {
 
   internal::PartitionBucket<thread_safe>* SizeToBucket(size_t size) const;
 
- private:
-  ALWAYS_INLINE void* AllocFromBucket(Bucket* bucket, int flags, size_t size)
-      EXCLUSIVE_LOCKS_REQUIRED(lock_);
-  bool ReallocDirectMappedInPlace(internal::PartitionPage<thread_safe>* page,
-                                  size_t raw_size)
-      EXCLUSIVE_LOCKS_REQUIRED(lock_);
-  void DecommitEmptyPages() EXCLUSIVE_LOCKS_REQUIRED(lock_);
-
   ALWAYS_INLINE internal::PartitionTag GetNewPartitionTag()
       EXCLUSIVE_LOCKS_REQUIRED(lock_) {
 #if ENABLE_TAG_FOR_CHECKED_PTR2 || ENABLE_TAG_FOR_MTE_CHECKED_PTR
@@ -522,6 +513,14 @@ struct BASE_EXPORT PartitionRoot {
     return 0;
 #endif
   }
+
+ private:
+  ALWAYS_INLINE void* AllocFromBucket(Bucket* bucket, int flags, size_t size)
+      EXCLUSIVE_LOCKS_REQUIRED(lock_);
+  bool ReallocDirectMappedInPlace(internal::PartitionPage<thread_safe>* page,
+                                  size_t raw_size)
+      EXCLUSIVE_LOCKS_REQUIRED(lock_);
+  void DecommitEmptyPages() EXCLUSIVE_LOCKS_REQUIRED(lock_);
 };
 
 template <bool thread_safe>
@@ -604,12 +603,16 @@ ALWAYS_INLINE void* PartitionRoot<thread_safe>::AllocFromBucket(Bucket* bucket,
     memset(ret, 0, size_with_no_extras);
   }
 
+  // Do not set tag for MTECheckedPtr in the set-tag-at-free case.
+  // It is set only at Free() time and at slot span allocation time.
+#if !ENABLE_TAG_FOR_MTE_CHECKED_PTR || !MTE_CHECKED_PTR_SET_TAG_AT_FREE
   if (allow_extras && !bucket->is_direct_mapped()) {
     size_t slot_size_with_no_extras = internal::PartitionSizeAdjustSubtract(
         allow_extras, page->bucket->slot_size);
     internal::PartitionTagSetValue(ret, slot_size_with_no_extras,
                                    GetNewPartitionTag());
   }
+#endif
 
   return ret;
 }
@@ -636,8 +639,11 @@ ALWAYS_INLINE void PartitionRoot<thread_safe>::Free(void* ptr) {
   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.
+#if ENABLE_TAG_FOR_MTE_CHECKED_PTR && MTE_CHECKED_PTR_SET_TAG_AT_FREE
+    internal::PartitionTagIncrementValue(ptr, size_with_no_extras);
+#else
     internal::PartitionTagClearValue(ptr, size_with_no_extras);
+#endif
   }
   ptr = internal::PartitionPointerAdjustSubtract(root->allow_extras, ptr);
   internal::DeferredUnmap deferred_unmap;

base/allocator/partition_allocator/partition_bucket.cc

@@ -259,6 +259,10 @@ ALWAYS_INLINE void* PartitionBucket<thread_safe>::AllocNewSlotSpan(
                            PageReadWrite);
       root->next_tag_bitmap_page = next_tag_bitmap_page;
     }
+#if MTE_CHECKED_PTR_SET_TAG_AT_FREE
+    // TODO(tasak): Consider initializing each slot with a different tag.
+    PartitionTagSetValue(ret, total_size, root->GetNewPartitionTag());
+#endif
 #endif
     return ret;
   }
@@ -323,6 +327,10 @@ ALWAYS_INLINE void* PartitionBucket<thread_safe>::AllocNewSlotSpan(
   PA_DCHECK(next_tag_bitmap_page <= ret);
   SetSystemPagesAccess(next_tag_bitmap_page, ret - next_tag_bitmap_page,
                        PageInaccessible);
+#if MTE_CHECKED_PTR_SET_TAG_AT_FREE
+  // TODO(tasak): Consider initializing each slot with a different tag.
+  PartitionTagSetValue(ret, total_size, root->GetNewPartitionTag());
+#endif
   root->next_tag_bitmap_page = next_tag_bitmap_page;
 #endif
 

base/allocator/partition_allocator/partition_tag.h

@@ -147,6 +147,24 @@ ALWAYS_INLINE void PartitionTagClearValue(void* ptr, size_t 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);
+}
+
 #elif ENABLE_TAG_FOR_SINGLE_TAG_CHECKED_PTR
 
 using PartitionTag = uint8_t;