Use mach_override to intercept all newly registered malloc zones.

base/allocator/allocator_shim.cc

 // Copyright 2016 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.
 
 #include "base/allocator/allocator_shim.h"
 
 #include <errno.h>
 
 #include <new>
 
 #include "base/atomicops.h"
 #include "base/logging.h"
 #include "base/macros.h"
 #include "base/process/process_metrics.h"
 #include "base/threading/platform_thread.h"
 #include "build/build_config.h"
 
 #if !defined(OS_WIN)
 #include <unistd.h>
 #else
 #include "base/allocator/winheap_stubs_win.h"
 #endif
 
 #if defined(OS_MACOSX)
 #include <malloc/malloc.h>
 
 #include "base/allocator/allocator_interception_mac.h"
+#include "base/base_switches.h"
+#include "base/command_line.h"
 #endif
 
 // No calls to malloc / new in this file. They would would cause re-entrancy of
 // the shim, which is hard to deal with. Keep this code as simple as possible
 // and don't use any external C++ object here, not even //base ones. Even if
 // they are safe to use today, in future they might be refactored.
 
 namespace {
 
 using namespace base;
@@ skipping 283 matching lines @@
 #if defined(OS_MACOSX)
 namespace base {
 namespace allocator {
 void InitializeAllocatorShim() {
   // Prepares the default dispatch. After the intercepted malloc calls have
   // traversed the shim this will route them to the default malloc zone.
   InitializeDefaultDispatchToMacAllocator();
 
   MallocZoneFunctions functions = MallocZoneFunctionsToReplaceDefault();
 
+  // Intercept and override newly registered malloc zones.
+  if (CommandLine::InitializedForCurrentProcess() &&
+      CommandLine::ForCurrentProcess()->HasSwitch(
+          switches::kEnableHeapProfiling)) {
+    InterceptNewlyRegisteredMallocZones(&functions);
+  }
+
   // This replaces the default malloc zone, causing calls to malloc & friends
   // from the codebase to be routed to ShimMalloc() above.
   base::allocator::ReplaceFunctionsForStoredZones(&functions);
 }
 }  // namespace allocator
 }  // namespace base
 #endif
 
 // Cross-checks.
 
 #if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
 #error The allocator shim should not be compiled when building for memory tools.
 #endif
 
 #if (defined(__GNUC__) && defined(__EXCEPTIONS)) || \
     (defined(_HAS_EXCEPTIONS) && _HAS_EXCEPTIONS)
 #error This code cannot be used when exceptions are turned on.
 #endif