Add simple correctness checks to PartitionAlloc.

ajwong pointed out that `PartitionRootBase::max_allocation` was used only in
tests. This CL additionally `DCHECK`s that it is not being mis-used in the most
simple ways. In the future, we might be able to change the API such that
`SizeSpecificPartitionAllocator` 'knows' its allocation size and just allocates
it, without needing a `size` parameter.

This CL also does some code cleanup. This involved making sure we always observe
`PartitionAllocReturnNull` in every case where `MEMORY_TOOL_REPLACES_ALLOCATOR`
is defined; previously we did not.

Bug: None
Change-Id: I04aa797176851e15ef3a0b803b85ddd038e0f772
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1607149Reviewed-by: Kentaro Hara <haraken@chromium.org>
Reviewed-by: Albert J. Wong <ajwong@chromium.org>
Commit-Queue: Chris Palmer <palmer@chromium.org>
Cr-Commit-Position: refs/heads/master@{#659314}

base/allocator/partition_allocator/partition_alloc.cc

@@ -198,13 +198,9 @@ void PartitionRoot::Init(size_t num_buckets, size_t max_allocation) {
 
   this->num_buckets = num_buckets;
   this->max_allocation = max_allocation;
-  size_t i;
-  for (i = 0; i < this->num_buckets; ++i) {
+  for (size_t i = 0; i < this->num_buckets; ++i) {
     internal::PartitionBucket* bucket = &this->buckets()[i];
-    if (!i)
-      bucket->Init(kAllocationGranularity);
-    else
-      bucket->Init(i << kBucketShift);
+    bucket->Init(i == 0 ? kAllocationGranularity : (i << kBucketShift));
   }
 }
 

base/allocator/partition_allocator/partition_alloc.h

@@ -82,6 +82,16 @@
 #include <stdlib.h>
 #endif
 
+// We use this to make MEMORY_TOOL_REPLACES_ALLOCATOR behave the same for max
+// size as other alloc code.
+#define CHECK_MAX_SIZE_OR_RETURN_NULLPTR(size, flags) \
+  if (size > kGenericMaxDirectMapped) {               \
+    if (flags & PartitionAllocReturnNull) {           \
+      return nullptr;                                 \
+    }                                                 \
+    CHECK(false);                                     \
+  }
+
 namespace base {
 
 class PartitionStatsDumper;
@@ -289,14 +299,12 @@ ALWAYS_INLINE void* PartitionRoot::AllocFlags(int flags,
                                               size_t size,
                                               const char* type_name) {
 #if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
-  // Make MEMORY_TOOL_REPLACES_ALLOCATOR behave the same for max size
-  // as other alloc code.
-  if (size > kGenericMaxDirectMapped)
-    return nullptr;
+  CHECK_MAX_SIZE_OR_RETURN_NULLPTR(size, flags);
   void* result = malloc(size);
   CHECK(result);
   return result;
 #else
+  DCHECK(max_allocation == 0 || size <= max_allocation);
   void* result;
   const bool hooks_enabled = PartitionAllocHooks::AreHooksEnabled();
   if (UNLIKELY(hooks_enabled)) {
@@ -388,16 +396,15 @@ ALWAYS_INLINE void* PartitionAllocGenericFlags(PartitionRootGeneric* root,
   DCHECK_LT(flags, PartitionAllocLastFlag << 1);
 
 #if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
-  // Make MEMORY_TOOL_REPLACES_ALLOCATOR behave the same for max size
-  // as other alloc code.
-  if (size > kGenericMaxDirectMapped)
-    return nullptr;
+  CHECK_MAX_SIZE_OR_RETURN_NULLPTR(size, flags);
   const bool zero_fill = flags & PartitionAllocZeroFill;
   void* result = zero_fill ? calloc(1, size) : malloc(size);
   CHECK(result || flags & PartitionAllocReturnNull);
   return result;
 #else
   DCHECK(root->initialized);
+  // Only SizeSpecificPartitionAllocator should use max_allocation.
+  DCHECK(root->max_allocation == 0);
   void* result;
   const bool hooks_enabled = PartitionAllocHooks::AreHooksEnabled();
   if (UNLIKELY(hooks_enabled)) {