Use mach_override to intercept all newly registered malloc zones.

base/allocator/allocator_interception_mac.mm

 // Copyright 2017 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 // This file contains all the logic necessary to intercept allocations on
 // macOS. "malloc zones" are an abstraction that allows the process to intercept
 // all malloc-related functions.  There is no good mechanism [short of
 // interposition] to determine new malloc zones are added, so there's no clean
 // mechanism to intercept all malloc zones. This file contains logic to
 // intercept the default and purgeable zones, which always exist. A cursory
@@ skipping 19 matching lines @@
 #include "base/allocator/allocator_shim.h"
 #include "base/allocator/features.h"
 #include "base/allocator/malloc_zone_functions_mac.h"
 #include "base/logging.h"
 #include "base/mac/mac_util.h"
 #include "base/mac/mach_logging.h"
 #include "base/process/memory.h"
 #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 {
 
 bool g_replaced_default_zone = false;
 
 namespace {
 
 bool g_oom_killer_enabled;
 
@@ skipping 245 matching lines @@
 
   // Restore protection if it was active.
   if (reprotection_start) {
     kern_return_t result =
         mach_vm_protect(mach_task_self(), reprotection_start,
                         reprotection_length, false, reprotection_value);
     MACH_CHECK(result == KERN_SUCCESS, result) << "mach_vm_protect";
   }
 }
 
+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))
     return;
   MallocZoneFunctions& functions = GetFunctionsForZone(zone);
   ReplaceZoneFunctions(chrome_zone, &functions);
 }
 
 }  // namespace
 
@@ skipping 191 matching lines @@
   Method orig_method =
       class_getClassMethod(nsobject_class, @selector(allocWithZone:));
   g_old_allocWithZone =
       reinterpret_cast<allocWithZone_t>(method_getImplementation(orig_method));
   CHECK(g_old_allocWithZone)
       << "Failed to get allocWithZone allocation function.";
   method_setImplementation(orig_method,
                            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;
   unsigned int count;
   kern_return_t kr = malloc_get_all_zones(mach_task_self(), 0, &zones, &count);
   CHECK(kr == KERN_SUCCESS);
   for (unsigned int i = 0; i < count; ++i) {
     UninterceptMallocZoneForTesting(
         reinterpret_cast<struct _malloc_zone_t*>(zones[i]));
   }
 
   ClearAllMallocZonesForTesting();
+  StopMallocZoneRegistrationInterceptionForTesting();
 }
 
 }  // namespace allocator
 }  // namespace base