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 55%
copy from base/allocator/allocator_interception_mac.h
copy to base/allocator/malloc_zone_functions_mac.h
index ff9b0b1b85c76e19c498bc89cb2e9671ac07a87a..5b3a9665b4926dcae75cf5ac082b38351936911f 100644
--- a/base/allocator/allocator_interception_mac.h
+++ b/base/allocator/malloc_zone_functions_mac.h
@@ -2,8 +2,8 @@
 // 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>
@@ -50,27 +50,41 @@ 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[kMaxZoneCount];
 
-// 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);
-
-// 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_