Use mach_override to intercept all newly registered malloc zones.

base/allocator/allocator_interception_mac.mm

Index: base/allocator/allocator_interception_mac.mm
diff --git a/base/allocator/allocator_interception_mac.mm b/base/allocator/allocator_interception_mac.mm
index 687100e1b81d7d62cca0e8fa45371e050f8fb38c..6b01b9112fed739a5ea3c23b1477ca45894443cf 100644
--- a/base/allocator/allocator_interception_mac.mm
+++ b/base/allocator/allocator_interception_mac.mm
@@ -37,6 +37,7 @@
 #include "base/scoped_clear_errno.h"
 #include "build/build_config.h"
 #include "third_party/apple_apsl/CFBase.h"
+#include "third_party/mach_override/mach_override.h"
 
 namespace base {
 namespace allocator {
@@ -302,6 +303,50 @@ void ReplaceZoneFunctions(ChromeMallocZone* zone,
   }
 }
 
+typedef void (*MallocZoneRegisterType)(malloc_zone_t* zone);
+MallocZoneRegisterType g_original_malloc_zone_register = nullptr;
+MallocZoneFunctions g_malloc_zone_functions;
+bool g_update_malloc_zone_during_interception = true;
+
+void ChromeMallocZoneRegister(malloc_zone_t* new_zone) {
+  DCHECK(g_original_malloc_zone_register);
+  if (!g_update_malloc_zone_during_interception) {
+    g_original_malloc_zone_register(new_zone);
+    return;
+  }
+
+  ChromeMallocZone* zone = reinterpret_cast<ChromeMallocZone*>(new_zone);
+
+  if (!StoreMallocZone(zone))
+    return;
+
+  ReplaceZoneFunctions(zone, &g_malloc_zone_functions);
+  g_original_malloc_zone_register(new_zone);
+}
+
+typedef struct _malloc_zone_t* (*MallocCreateZoneType)(vm_size_t start_size,
+                                                       unsigned flags);
+MallocCreateZoneType g_original_malloc_create_zone = nullptr;
+
+struct _malloc_zone_t* ChromeMallocCreateZone(vm_size_t start_size,
+                                              unsigned flags) {
+  DCHECK(g_original_malloc_create_zone);
+
+  struct _malloc_zone_t* new_zone =
+      g_original_malloc_create_zone(start_size, flags);
+  if (!g_update_malloc_zone_during_interception) {
+    return new_zone;
+  }
+
+  ChromeMallocZone* zone = reinterpret_cast<ChromeMallocZone*>(new_zone);
+
+  if (!StoreMallocZone(zone))
+    return new_zone;
+
+  ReplaceZoneFunctions(zone, &g_malloc_zone_functions);
+  return new_zone;
+}
+
 void UninterceptMallocZoneForTesting(struct _malloc_zone_t* zone) {
   ChromeMallocZone* chrome_zone = reinterpret_cast<ChromeMallocZone*>(zone);
   if (!IsMallocZoneAlreadyStored(chrome_zone))
@@ -513,6 +558,43 @@ void InterceptAllocationsMac() {
                            reinterpret_cast<IMP>(oom_killer_allocWithZone));
 }
 
+void InterceptNewlyRegisteredMallocZones(const MallocZoneFunctions* functions) {
+  g_update_malloc_zone_during_interception = true;
+
+  // In a typical Chrome build, InterceptNewlyRegisteredMallocZones should only
+  // be called once. In a testing environment, it might be called more often.
+  if (g_original_malloc_zone_register)
+    return;
+
+  memcpy(&g_malloc_zone_functions, functions, sizeof(MallocZoneFunctions));
+  mach_error_t err = mach_override_ptr(
+      reinterpret_cast<void*>(malloc_zone_register),
+      reinterpret_cast<void*>(ChromeMallocZoneRegister),
+      reinterpret_cast<void**>(&g_original_malloc_zone_register));
+  if (err != err_none) {
+    DLOG(WARNING) << "mach_override malloc_zone_register: " << err;
+  }
+
+  // On macOS 10.9, the implementation of malloc_create_zone calls
+  // malloc_zone_register_while_locked, a private symbol in libmalloc, so we
+  // also need to interpose it. Starting from macOS 10.10.0, libmalloc was open
+  // sourced, and the implementation of malloc_create_zone calls
+  // malloc_zone_register.
+  if (base::mac::IsOS10_9()) {
+    err = mach_override_ptr(
+        reinterpret_cast<void*>(malloc_create_zone),
+        reinterpret_cast<void*>(ChromeMallocCreateZone),
+        reinterpret_cast<void**>(&g_original_malloc_create_zone));
+    if (err != err_none) {
+      DLOG(WARNING) << "mach_override malloc_create_zone: " << err;
+    }
+  }
+}
+
+void StopMallocZoneRegistrationInterceptionForTesting() {
+  g_update_malloc_zone_during_interception = false;
+}
+
 void UninterceptMallocZonesForTesting() {
   UninterceptMallocZoneForTesting(malloc_default_zone());
   vm_address_t* zones;
@@ -525,6 +607,7 @@ void UninterceptMallocZonesForTesting() {
   }
 
   ClearAllMallocZonesForTesting();
+  StopMallocZoneRegistrationInterceptionForTesting();
 }
 
 }  // namespace allocator