mac: Fix StoreFunctionsForAllZones and ReplaceFunctionsForStoredZones.

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 334 matching lines @@
 
   vm_address_t* zones;
   unsigned int count;
   kern_return_t kr = malloc_get_all_zones(mach_task_self(), 0, &zones, &count);
   if (kr != KERN_SUCCESS)
     return;
   for (unsigned int i = 0; i < count; ++i) {
     ChromeMallocZone* zone = reinterpret_cast<ChromeMallocZone*>(zones[i]);
     StoreMallocZone(zone);
   }
+
+  // This ensures that the purgeable zone is at the back of the array.
+  ChromeMallocZone* purgeable_zone =
+      reinterpret_cast<ChromeMallocZone*>(malloc_default_purgeable_zone());
+  StoreMallocZone(purgeable_zone);
 }
 
 void ReplaceFunctionsForStoredZones(const MallocZoneFunctions* functions) {
+  ChromeMallocZone* default_zone =
+      reinterpret_cast<ChromeMallocZone*>(malloc_default_zone());
+  if (IsMallocZoneAlreadyStored(default_zone) &&
+      default_zone->malloc != functions->malloc) {
+    ReplaceZoneFunctions(default_zone, functions);
+  }
+
   vm_address_t* zones;
   unsigned int count;
   kern_return_t kr =
       malloc_get_all_zones(mach_task_self(), nullptr, &zones, &count);
   if (kr != KERN_SUCCESS)
     return;
   for (unsigned int i = 0; i < count; ++i) {
     ChromeMallocZone* zone = reinterpret_cast<ChromeMallocZone*>(zones[i]);
     if (IsMallocZoneAlreadyStored(zone) && zone->malloc != functions->malloc) {
       ReplaceZoneFunctions(zone, functions);
     }
   }
+
+  ChromeMallocZone* purgeable_zone =
+      reinterpret_cast<ChromeMallocZone*>(malloc_default_purgeable_zone());
+  if (IsMallocZoneAlreadyStored(purgeable_zone) &&
+      purgeable_zone->malloc != functions->malloc) {
+    ReplaceZoneFunctions(purgeable_zone, functions);
+  }
+
   g_replaced_default_zone = true;
 }
 
 void InterceptAllocationsMac() {
   if (g_oom_killer_enabled)
     return;
 
   g_oom_killer_enabled = true;
 
 // === C malloc/calloc/valloc/realloc/posix_memalign ===
@@ skipping 111 matching lines @@
   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));
 }
 
 }  // namespace allocator
 }  // namespace base