Add extra 16 bytes per each allocation for PartitionTag

Use the first 16bytes region of the allocated memory for tag. So both data region, and PartitionCookie region will be moved by +16bytes.
This behavior is disabled behind ENABLE_CHECKED_PTR.

Change-Id: I09a0241f61ca1b186b239097601ea498f6d9f00c
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2263417
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@{#782887}

base/BUILD.gn

@@ -1765,6 +1765,7 @@ jumbo_component("base") {
         "allocator/partition_allocator/partition_oom.h",
         "allocator/partition_allocator/partition_page.cc",
         "allocator/partition_allocator/partition_page.h",
+        "allocator/partition_allocator/partition_tag.h",
         "allocator/partition_allocator/random.cc",
         "allocator/partition_allocator/random.h",
         "allocator/partition_allocator/spin_lock.cc",

base/allocator/partition_allocator/partition_alloc.cc

@@ -303,6 +303,7 @@ bool PartitionRoot<thread_safe>::ReallocDirectMappedInPlace(
   PA_DCHECK(page->bucket->is_direct_mapped());
 
   raw_size = internal::PartitionCookieSizeAdjustAdd(raw_size);
+  raw_size = internal::PartitionTagSizeAdjustAdd(raw_size);
 
   // Note that the new size might be a bucketed size; this function is called
   // whenever we're reallocating a direct mapped allocation.
@@ -388,8 +389,8 @@ void* PartitionRoot<thread_safe>::ReallocFlags(int flags,
         &actual_old_size, ptr);
   }
   if (LIKELY(!overridden)) {
-    auto* page =
-        Page::FromPointer(internal::PartitionCookieFreePointerAdjust(ptr));
+    auto* page = Page::FromPointer(internal::PartitionTagFreePointerAdjust(
+        internal::PartitionCookieFreePointerAdjust(ptr)));
     bool success = false;
     {
       internal::ScopedGuard<thread_safe> guard{lock_};
@@ -421,7 +422,9 @@ void* PartitionRoot<thread_safe>::ReallocFlags(int flags,
       // Trying to allocate a block of size |new_size| would give us a block of
       // the same size as the one we've already got, so re-use the allocation
       // after updating statistics (and cookies, if present).
-      page->set_raw_size(internal::PartitionCookieSizeAdjustAdd(new_size));
+      size_t new_raw_size = internal::PartitionCookieSizeAdjustAdd(new_size);
+      new_raw_size = internal::PartitionTagSizeAdjustAdd(new_raw_size);
+      page->set_raw_size(new_raw_size);
 #if DCHECK_IS_ON()
       // Write a new trailing cookie when it is possible to keep track of
       // |new_size| via the raw size pointer.

base/allocator/partition_allocator/partition_alloc.h

@@ -64,6 +64,7 @@
 #include "base/allocator/partition_allocator/partition_cookie.h"
 #include "base/allocator/partition_allocator/partition_direct_map_extent.h"
 #include "base/allocator/partition_allocator/partition_page.h"
+#include "base/allocator/partition_allocator/partition_tag.h"
 #include "base/allocator/partition_allocator/spin_lock.h"
 #include "base/base_export.h"
 #include "base/bits.h"
@@ -460,26 +461,42 @@ ALWAYS_INLINE void* PartitionRoot<thread_safe>::AllocFromBucket(Bucket* bucket,
     PA_DCHECK(raw_size == size);
     new_slot_size = raw_size;
   }
-  size_t no_cookie_size =
-      internal::PartitionCookieSizeAdjustSubtract(new_slot_size);
-  char* char_ret = static_cast<char*>(ret);
+  // 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.
+  size_t size_with_no_extras = internal::PartitionTagSizeAdjustSubtract(
+      internal::PartitionCookieSizeAdjustSubtract(new_slot_size));
+  char* char_ret = static_cast<char*>(ret) + internal::kPartitionTagSize;
   // The value given to the application is actually just after the cookie.
   ret = char_ret + internal::kCookieSize;
 
   // Fill the region kUninitializedByte or 0, and surround it with 2 cookies.
   internal::PartitionCookieWriteValue(char_ret);
   if (!zero_fill) {
-    memset(ret, kUninitializedByte, no_cookie_size);
+    memset(ret, kUninitializedByte, size_with_no_extras);
   } else if (!is_already_zeroed) {
-    memset(ret, 0, no_cookie_size);
+    memset(ret, 0, size_with_no_extras);
   }
   internal::PartitionCookieWriteValue(char_ret + internal::kCookieSize +
-                                      no_cookie_size);
+                                      size_with_no_extras);
 #else
-  if (ret && zero_fill && !is_already_zeroed) {
-    memset(ret, 0, size);
+  if (!ret)
+    return nullptr;
+
+  ret = static_cast<char*>(ret) + internal::kPartitionTagSize;
+  if (zero_fill && !is_already_zeroed) {
+    memset(ret, 0, internal::PartitionTagSizeAdjustSubtract(size));
   }
 #endif
+  // TODO(tasak): initialize tag randomly. Temporarily use
+  // kTagTemporaryInitialValue to initialize the tag.
+  internal::PartitionTagSetValue(ret, internal::kTagTemporaryInitialValue);
 
   return ret;
 }
@@ -500,7 +517,10 @@ ALWAYS_INLINE void PartitionRoot<thread_safe>::Free(void* ptr) {
       return;
   }
 
+  // TODO(tasak): clear partition tag. Temporarily set the tag to be 0.
+  internal::PartitionTagSetValue(ptr, 0);
   ptr = internal::PartitionCookieFreePointerAdjust(ptr);
+  ptr = internal::PartitionTagFreePointerAdjust(ptr);
   Page* page = Page::FromPointer(ptr);
   // TODO(palmer): See if we can afford to make this a CHECK.
   PA_DCHECK(IsValidPage(page));
@@ -616,8 +636,12 @@ PartitionAllocGetPageForSize(void* ptr) {
 template <bool thread_safe>
 ALWAYS_INLINE size_t PartitionAllocGetSize(void* ptr) {
   ptr = internal::PartitionCookieFreePointerAdjust(ptr);
+  ptr = internal::PartitionTagFreePointerAdjust(ptr);
   auto* page = internal::PartitionAllocGetPageForSize<thread_safe>(ptr);
-  return internal::PartitionCookieSizeAdjustSubtract(page->bucket->slot_size);
+  size_t size =
+      internal::PartitionCookieSizeAdjustSubtract(page->bucket->slot_size);
+  size = internal::PartitionTagSizeAdjustSubtract(size);
+  return size;
 }
 
 // Gets the offset from the beginning of the allocated slot, adjusted for cookie
@@ -628,6 +652,7 @@ template <bool thread_safe>
 ALWAYS_INLINE size_t PartitionAllocGetSlotOffset(void* ptr) {
   PA_DCHECK(IsManagedByPartitionAllocAndNotDirectMapped(ptr));
   ptr = internal::PartitionCookieFreePointerAdjust(ptr);
+  ptr = internal::PartitionTagFreePointerAdjust(ptr);
   auto* page = internal::PartitionAllocGetPageForSize<thread_safe>(ptr);
   size_t slot_size = page->bucket->slot_size;
 
@@ -687,6 +712,10 @@ ALWAYS_INLINE void* PartitionRoot<thread_safe>::AllocFlags(
   }
   size_t requested_size = size;
   size = internal::PartitionCookieSizeAdjustAdd(size);
+  size = internal::PartitionTagSizeAdjustAdd(size);
+#if ENABLE_CHECKED_PTR
+  PA_CHECK(size >= requested_size);
+#endif
   auto* bucket = SizeToBucket(size);
   PA_DCHECK(bucket);
   {
@@ -730,6 +759,7 @@ ALWAYS_INLINE size_t PartitionRoot<thread_safe>::ActualSize(size_t size) {
 #else
   PA_DCHECK(PartitionRoot<thread_safe>::initialized);
   size = internal::PartitionCookieSizeAdjustAdd(size);
+  size = internal::PartitionTagSizeAdjustAdd(size);
   auto* bucket = SizeToBucket(size);
   if (LIKELY(!bucket->is_direct_mapped())) {
     size = bucket->slot_size;
@@ -738,7 +768,9 @@ ALWAYS_INLINE size_t PartitionRoot<thread_safe>::ActualSize(size_t size) {
   } else {
     size = Bucket::get_direct_map_size(size);
   }
-  return internal::PartitionCookieSizeAdjustSubtract(size);
+  size = internal::PartitionCookieSizeAdjustSubtract(size);
+  size = internal::PartitionTagSizeAdjustSubtract(size);
+  return size;
 #endif
 }
 

base/allocator/partition_allocator/partition_page.h

@@ -13,6 +13,7 @@
 #include "base/allocator/partition_allocator/partition_bucket.h"
 #include "base/allocator/partition_allocator/partition_cookie.h"
 #include "base/allocator/partition_allocator/partition_freelist_entry.h"
+#include "base/allocator/partition_allocator/partition_tag.h"
 #include "base/allocator/partition_allocator/random.h"
 #include "base/check_op.h"
 #include "base/thread_annotations.h"
@@ -247,7 +248,7 @@ ALWAYS_INLINE DeferredUnmap PartitionPage<thread_safe>::Free(void* ptr) {
   }
 
   // If these asserts fire, you probably corrupted memory.
-  PartitionCookieCheckValue(ptr);
+  PartitionCookieCheckValue(reinterpret_cast<char*>(ptr) + kPartitionTagSize);
   PartitionCookieCheckValue(reinterpret_cast<char*>(ptr) + slot_size -
                             kCookieSize);
 

base/allocator/partition_allocator/partition_tag.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_PARTITION_TAG_H_
+#define BASE_ALLOCATOR_PARTITION_ALLOCATOR_PARTITION_TAG_H_
+
+#include "base/allocator/partition_allocator/partition_alloc_constants.h"
+#include "base/allocator/partition_allocator/partition_cookie.h"
+#include "base/notreached.h"
+#include "build/build_config.h"
+
+#define ENABLE_CHECKED_PTR 0
+
+namespace base {
+namespace internal {
+
+// Use 16 bits for the partition tag.
+// TODO(tasak): add a description about the partition tag.
+using PartitionTag = uint16_t;
+
+static constexpr PartitionTag kTagTemporaryInitialValue = 0x0BAD;
+
+#if ENABLE_CHECKED_PTR
+
+// Allocate extra 16 bytes for the partition tag. 14 bytes are unused
+// (reserved).
+static constexpr size_t kPartitionTagSize = 16;
+
+#if DCHECK_IS_ON()
+// The layout inside the slot is |tag|cookie|object|(empty)|cookie|.
+static constexpr size_t kPartitionTagOffset = kPartitionTagSize + kCookieSize;
+#else
+// The layout inside the slot is |tag|object|(empty)|.
+static constexpr size_t kPartitionTagOffset = kPartitionTagSize;
+#endif
+
+ALWAYS_INLINE size_t PartitionTagSizeAdjustAdd(size_t size) {
+  PA_DCHECK(size + kPartitionTagSize > size);
+  return size + kPartitionTagSize;
+}
+
+ALWAYS_INLINE size_t PartitionTagSizeAdjustSubtract(size_t size) {
+  PA_DCHECK(size >= kPartitionTagSize);
+  return size - kPartitionTagSize;
+}
+
+ALWAYS_INLINE PartitionTag* PartitionTagPointer(void* ptr) {
+  return reinterpret_cast<PartitionTag*>(reinterpret_cast<char*>(ptr) -
+                                         kPartitionTagOffset);
+}
+
+ALWAYS_INLINE void* PartitionTagFreePointerAdjust(void* ptr) {
+  return reinterpret_cast<void*>(reinterpret_cast<char*>(ptr) -
+                                 kPartitionTagSize);
+}
+
+ALWAYS_INLINE void PartitionTagSetValue(void* ptr, uint16_t value) {
+  *PartitionTagPointer(ptr) = value;
+}
+
+ALWAYS_INLINE PartitionTag PartitionTagGetValue(void* ptr) {
+  return *PartitionTagPointer(ptr);
+}
+
+#else  // !ENABLE_CHECKED_PTR
+
+// No tag added.
+static constexpr size_t kPartitionTagSize = 0;
+
+ALWAYS_INLINE size_t PartitionTagSizeAdjustAdd(size_t size) {
+  return size;
+}
+
+ALWAYS_INLINE size_t PartitionTagSizeAdjustSubtract(size_t size) {
+  return size;
+}
+
+ALWAYS_INLINE PartitionTag* PartitionTagPointer(void* ptr) {
+  NOTREACHED();
+  return nullptr;
+}
+
+ALWAYS_INLINE void* PartitionTagFreePointerAdjust(void* ptr) {
+  return ptr;
+}
+
+ALWAYS_INLINE void PartitionTagSetValue(void*, uint16_t) {}
+
+ALWAYS_INLINE PartitionTag PartitionTagGetValue(void*) {
+  return 0;
+}
+
+#endif  // !ENABLE_CHECKED_PTR
+
+}  // namespace internal
+}  // namespace base
+
+#endif  // BASE_ALLOCATOR_PARTITION_ALLOCATOR_PARTITION_TAG_H_