macOS: Shim all 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 e945a293beb58fbe36e9e4f2c17f1525293b28cc..87c882a19cb8e5214610857237c50f8a6db41659 100644
--- a/base/allocator/allocator_interception_mac.mm
+++ b/base/allocator/allocator_interception_mac.mm
@@ -29,6 +29,7 @@
 
 #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"
@@ -42,8 +43,6 @@ namespace allocator {
 
 bool g_replaced_default_zone = false;
 
-MallocZoneFunctions::MallocZoneFunctions() {}
-
 namespace {
 
 bool g_oom_killer_enabled;
@@ -295,31 +294,6 @@ bool UncheckedCallocMac(size_t num_items, size_t size, void** result) {
   return *result != NULL;
 }
 
-void StoreZoneFunctions(ChromeMallocZone* zone,
-                        MallocZoneFunctions* functions) {
-  functions->malloc = zone->malloc;
-  functions->calloc = zone->calloc;
-  functions->valloc = zone->valloc;
-  functions->free = zone->free;
-  functions->realloc = zone->realloc;
-  functions->size = zone->size;
-  CHECK(functions->malloc && functions->calloc && functions->valloc &&
-        functions->free && functions->realloc && functions->size);
-
-  // These functions might be nullptr.
-  functions->batch_malloc = zone->batch_malloc;
-  functions->batch_free = zone->batch_free;
-
-  if (zone->version >= 5) {
-    functions->memalign = zone->memalign;
-    CHECK(functions->memalign);
-  }
-  if (zone->version >= 6) {
-    // This may be nullptr.
-    functions->free_definite_size = zone->free_definite_size;
-  }
-}
-
 void ReplaceZoneFunctions(ChromeMallocZone* zone,
                           const MallocZoneFunctions* functions) {
   // Remove protection.
@@ -358,21 +332,44 @@ void ReplaceZoneFunctions(ChromeMallocZone* zone,
   }
 }
 
-void StoreFunctionsForDefaultZone(MallocZoneFunctions* functions) {
+void StoreFunctionsForDefaultZone() {
   ChromeMallocZone* default_zone = reinterpret_cast<ChromeMallocZone*>(
       malloc_default_zone());
-  StoreZoneFunctions(default_zone, functions);
+  StoreMallocZone(default_zone);
+}
+
+void StoreFunctionsForAllZones() {
+  // This ensures that the default zone is always at the front of the array,
+  // which is important for performance.
+  StoreFunctionsForDefaultZone();
+
+  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) {

[Primiano Tucci (use gerrit), 2017/02/23 00:40:17]

just checking: what happens if a zone goes away while we iterate? there is no
locking here.
Asking because we had some similar problem on windows and the conclusion was
that the winapi don't have any way to safely enumerate the heaps (similar
concept to these zones, really), as the winapi  would crash if a heap is
destroyed while iterating. (sorry too late to dig in crbug, but if you git blame
the malloc_dump_provider under the win section I'm pretty sure you'll find the
CL with all the context, if interested)

I hope we're not in the same situation here :)

[erikchen, 2017/02/23 01:43:14]

same problem, sorry. 

[Primiano Tucci (use gerrit), 2017/02/23 11:19:41]

Are you saying that we know upfront that this code will likely make chrome
crash?

[erikchen, 2017/02/23 20:06:44]

there is a race condition that could cause a chrome crash, if a library destroys
a malloc zone while we iterate over all zones. 

malloc_get_all_zones is almost really helpful, as it takes a memory_reader_t
argument that theoretically can copy into a safe region, but it doesn't acquire
the malloc zone lock! Ugh.

I'm not worried about this CL, since it shims very early on in the life cycle. I
am however worried about the next CL. Will have a follow-up discussion there.

[Primiano Tucci (use gerrit), 2017/02/23 21:06:07]

Acknowledged.

+    ChromeMallocZone* zone = reinterpret_cast<ChromeMallocZone*>(zones[i]);
+    StoreMallocZone(zone);
+  }
 }
 
-void ReplaceFunctionsForDefaultZone(const MallocZoneFunctions* functions) {
-  CHECK(!g_replaced_default_zone);
+void ReplaceFunctionsForStoredZones(const MallocZoneFunctions* functions) {
+  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]);
+    if (!IsMallocZoneAlreadyStored(zone)) {
+      continue;
+    }
+    if (zone->malloc == functions->malloc)
+      continue;
+    ReplaceZoneFunctions(zone, functions);
+  }
   g_replaced_default_zone = true;
-#if !defined(ADDRESS_SANITIZER)
-  StoreFunctionsForDefaultZone(&g_old_zone);
-#endif
-  ChromeMallocZone* default_zone = reinterpret_cast<ChromeMallocZone*>(
-      malloc_default_zone());
-  ReplaceZoneFunctions(default_zone, functions);
 }
 
 void InterceptAllocationsMac() {
@@ -393,8 +390,10 @@ void InterceptAllocationsMac() {
 #if !defined(ADDRESS_SANITIZER)
   // Don't do anything special on OOM for the malloc zones replaced by
   // AddressSanitizer, as modifying or protecting them may not work correctly.
-  if (!g_replaced_default_zone) {
-    StoreFunctionsForDefaultZone(&g_old_zone);
+  ChromeMallocZone* default_zone =
+      reinterpret_cast<ChromeMallocZone*>(malloc_default_zone());
+  if (!IsMallocZoneAlreadyStored(default_zone)) {
+    StoreZoneFunctions(default_zone, &g_old_zone);
     MallocZoneFunctions new_functions;
     new_functions.malloc = oom_killer_malloc;
     new_functions.calloc = oom_killer_calloc;
@@ -402,12 +401,14 @@ void InterceptAllocationsMac() {
     new_functions.free = oom_killer_free;
     new_functions.realloc = oom_killer_realloc;
     new_functions.memalign = oom_killer_memalign;
-    ReplaceFunctionsForDefaultZone(&new_functions);
+
+    ReplaceZoneFunctions(default_zone, &new_functions);
+    g_replaced_default_zone = true;
   }
 
   ChromeMallocZone* purgeable_zone =
       reinterpret_cast<ChromeMallocZone*>(malloc_default_purgeable_zone());
-  if (purgeable_zone) {
+  if (purgeable_zone && !IsMallocZoneAlreadyStored(purgeable_zone)) {
     StoreZoneFunctions(purgeable_zone, &g_old_purgeable_zone);
     MallocZoneFunctions new_functions;
     new_functions.malloc = oom_killer_malloc_purgeable;