macOS: Shim all malloc zones.

base/allocator/malloc_zone_functions_mac.h

Index: base/allocator/malloc_zone_functions_mac.h
diff --git a/base/allocator/allocator_interception_mac.h b/base/allocator/malloc_zone_functions_mac.h
similarity index 52%
copy from base/allocator/allocator_interception_mac.h
copy to base/allocator/malloc_zone_functions_mac.h
index ff9b0b1b85c76e19c498bc89cb2e9671ac07a87a..c2ec2fcd3535fee4952a7c6a91048601fb337e54 100644
--- a/base/allocator/allocator_interception_mac.h
+++ b/base/allocator/malloc_zone_functions_mac.h
@@ -2,13 +2,14 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
-#ifndef BASE_ALLOCATOR_ALLOCATOR_INTERCEPTION_MAC_H_
-#define BASE_ALLOCATOR_ALLOCATOR_INTERCEPTION_MAC_H_
+#ifndef BASE_ALLOCATOR_MALLOC_ZONE_FUNCTIONS_MAC_H_
+#define BASE_ALLOCATOR_MALLOC_ZONE_FUNCTIONS_MAC_H_
 
 #include <malloc/malloc.h>
 #include <stddef.h>
 
 #include "base/base_export.h"
+#include "base/logging.h"
 #include "third_party/apple_apsl/malloc.h"
 
 namespace base {
@@ -50,27 +51,49 @@ struct MallocZoneFunctions {
   batch_free_type batch_free = nullptr;
   free_definite_size_type free_definite_size = nullptr;
   size_fn_type size = nullptr;
+  const ChromeMallocZone* context = nullptr;
 };
 
-// Saves the function pointers currently used by default zone into |functions|.
-void StoreFunctionsForDefaultZone(MallocZoneFunctions* functions);
+void StoreZoneFunctions(const ChromeMallocZone* zone,
+                        MallocZoneFunctions* functions);
+static constexpr int kMaxZoneCount = 30;
+BASE_EXPORT extern MallocZoneFunctions* g_malloc_zones;
 
-// Updates the default malloc zone to use the functions specified by
-// |functions|.
-void ReplaceFunctionsForDefaultZone(const MallocZoneFunctions* functions);
+// The array g_malloc_zones stores all information about malloc zones before
+// they are shimmed. This information needs to be accessed during dispatch back
+// into the zone, and additional zones may be added later in the execution fo
+// the program, so the array needs to be both thread-safe and high-performance.
+//
+// We begin by creating an array of MallocZoneFunctions of fixed size. We will
+// never modify the container, which provides thread-safety to iterators.  When
+// we want to add a MallocZoneFunctions to the container, we:
+//   1. Fill in all the fields.
+//   2. Update the total zone count.
+//   3. Insert a memory barrier.
+//   4. Insert our shim.
+//
+// Each MallocZoneFunctions is uniquely identified by |context|, which is a
+// pointer to the original malloc zone. When we wish to dispatch back to the
+// original malloc zones, we iterate through the array, looking for a matching
+// |context|.
+//
+// Most allocations go through the default allocator. We will ensure that the
+// default allocator is stored as the first MallocZoneFunctions.
+BASE_EXPORT void StoreMallocZone(ChromeMallocZone* zone);
+BASE_EXPORT bool IsMallocZoneAlreadyStored(ChromeMallocZone* zone);
 
-extern bool g_replaced_default_zone;
+BASE_EXPORT int GetMallocZoneCountForTesting();
+BASE_EXPORT void ClearAllMallocZonesForTesting();
 
-// Calls the original implementation of malloc/calloc prior to interception.
-bool UncheckedMallocMac(size_t size, void** result);
-bool UncheckedCallocMac(size_t num_items, size_t size, void** result);
+inline MallocZoneFunctions& GetFunctionsForZone(void* zone) {
+  for (unsigned int i = 0; i < kMaxZoneCount; ++i) {
+    if (g_malloc_zones[i].context == zone)
+      return g_malloc_zones[i];
+  }
+  IMMEDIATE_CRASH();
+}
 
-// Intercepts calls to default and purgeable malloc zones. Intercepts Core
-// Foundation and Objective-C allocations.
-// Has no effect on the default malloc zone if the allocator shim already
-// performs that interception.
-BASE_EXPORT void InterceptAllocationsMac();
 }  // namespace allocator
 }  // namespace base
 
-#endif  // BASE_ALLOCATOR_ALLOCATOR_INTERCEPTION_MAC_H_
+#endif  // BASE_ALLOCATOR_MALLOC_ZONE_FUNCTIONS_MAC_H_