[PA][BRP] Allow RefCount in aligned partition if it's after allocation

This makes it independent of a cookie (present only in DCHECK builds)
thus bringing the Debug build one step closer to the Release build.

Change-Id: I4c0f441f752b11b87ae355f68f64349e6b6575be
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2763992
Commit-Queue: Bartek Nowierski <bartekn@chromium.org>
Reviewed-by: Benoit L <lizeb@chromium.org>
Reviewed-by: Kentaro Hara <haraken@chromium.org>
Cr-Commit-Position: refs/heads/master@{#864083}

base/allocator/allocator_shim_default_dispatch_to_partition_alloc.cc

@@ -144,12 +144,20 @@ base::ThreadSafePartitionRoot* AlignedAllocator() {
 #else
   // Since the general-purpose allocator uses the thread cache, this one cannot.
   static base::NoDestructor<base::ThreadSafePartitionRoot> aligned_allocator(
-      base::PartitionOptions{base::PartitionOptions::Alignment::kAlignedAlloc,
-                             base::PartitionOptions::ThreadCache::kDisabled,
-                             base::PartitionOptions::Quarantine::kAllowed,
-                             base::PartitionOptions::RefCount::kDisabled});
-  return aligned_allocator.get();
+      base::PartitionOptions {
+        base::PartitionOptions::Alignment::kAlignedAlloc,
+            base::PartitionOptions::ThreadCache::kDisabled,
+            base::PartitionOptions::Quarantine::kAllowed,
+#if BUILDFLAG(REF_COUNT_AT_END_OF_ALLOCATION)
+            // Given the outer #if, this is possible only when DCHECK_IS_ON().
+            base::PartitionOptions::RefCount::kEnabled
+#else
+            base::PartitionOptions::RefCount::kDisabled
 #endif
+      });
+  return aligned_allocator.get();
+#endif  // !DCHECK_IS_ON() && (!BUILDFLAG(USE_BACKUP_REF_PTR) ||
+        //                     BUILDFLAG(REF_COUNT_AT_END_OF_ALLOCATION))
 }
 
 #if defined(OS_WIN) && defined(ARCH_CPU_X86)

base/allocator/partition_allocator/partition_ref_count.cc

@@ -38,9 +38,11 @@ PartitionRefCount* PartitionRefCountPointer(void* slot_start) {
   // reserved). Instead, refcount is stored in the subsequent page metadata.
 
   auto* slot_span = SlotSpanMetadata<ThreadSafe>::FromSlotStartPtr(slot_start);
-  PA_DCHECK(slot_span);
 #if DCHECK_IS_ON()
-  PartitionCookieCheckValue(slot_start);
+  PA_DCHECK(slot_span);
+  auto* root = PartitionRoot<ThreadSafe>::FromSlotSpan(slot_span);
+  if (root->allow_cookies)
+    PartitionCookieCheckValue(slot_start);
 #endif
   uint8_t* partition_ref_count_ptr;
   if (UNLIKELY(slot_span->CanStoreRawSize())) {

base/allocator/partition_allocator/partition_root.cc

@@ -477,24 +477,14 @@ void PartitionRoot<thread_safe>::Init(PartitionOptions opts) {
       internal::PartitionAddressSpace::Init();
 #endif
 
-    // If alignment needs to be enforced, disallow adding a cookie and/or
-    // ref-count at the beginning of the slot.
-    if (opts.alignment == PartitionOptions::Alignment::kAlignedAlloc) {
-      allow_cookies = false;
-      allow_ref_count = false;
-      // There should be no configuration where aligned root and ref-count are
-      // requested at the same time. In theory REF_COUNT_AT_END_OF_ALLOCATION
-      // allows these to co-exist, but in this case aligned root is not even
-      // created.
-      PA_CHECK(opts.ref_count == PartitionOptions::RefCount::kDisabled);
-    } else {
-      allow_cookies = true;
-      // Allow ref-count if it's explicitly requested *and* GigaCage is enabled.
-      // Without GigaCage it'd be unused, thus wasteful.
-      allow_ref_count =
-          (opts.ref_count == PartitionOptions::RefCount::kEnabled) &&
-          features::IsPartitionAllocGigaCageEnabled();
-    }
+    // If alignment needs to be enforced, disallow adding a cookie.
+    allow_cookies =
+        opts.alignment != PartitionOptions::Alignment::kAlignedAlloc;
+    // Allow ref-count if it's explicitly requested *and* GigaCage is enabled.
+    // Without GigaCage it'd be unused, thus wasteful.
+    allow_ref_count =
+        (opts.ref_count == PartitionOptions::RefCount::kEnabled) &&
+        features::IsPartitionAllocGigaCageEnabled();
 
 #if PARTITION_EXTRAS_REQUIRED
     extras_size = 0;
@@ -514,6 +504,10 @@ void PartitionRoot<thread_safe>::Init(PartitionOptions opts) {
     }
 #endif
 
+    // Partitions that allow AlignedAlloc can't have pre-allocation extras.
+    PA_CHECK(opts.alignment != PartitionOptions::Alignment::kAlignedAlloc ||
+             !extras_offset);
+
     quarantine_mode =
 #if PA_ALLOW_PCSCAN
         (opts.quarantine == PartitionOptions::Quarantine::kDisallowed