mac: Several minor fixes to allocator shim.

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 c484830da5e2538904cd1b8218ca8fb97e354536..687100e1b81d7d62cca0e8fa45371e050f8fb38c 100644
--- a/base/allocator/allocator_interception_mac.mm
+++ b/base/allocator/allocator_interception_mac.mm
@@ -264,36 +264,6 @@ id oom_killer_allocWithZone(id self, SEL _cmd, NSZone* zone) {
   return result;
 }
 
-}  // namespace
-
-bool UncheckedMallocMac(size_t size, void** result) {
-#if defined(ADDRESS_SANITIZER)
-  *result = malloc(size);
-#else
-  if (g_old_zone.malloc) {
-    *result = g_old_zone.malloc(malloc_default_zone(), size);
-  } else {
-    *result = malloc(size);
-  }
-#endif  // defined(ADDRESS_SANITIZER)
-
-  return *result != NULL;
-}
-
-bool UncheckedCallocMac(size_t num_items, size_t size, void** result) {
-#if defined(ADDRESS_SANITIZER)
-  *result = calloc(num_items, size);
-#else
-  if (g_old_zone.calloc) {
-    *result = g_old_zone.calloc(malloc_default_zone(), num_items, size);
-  } else {
-    *result = calloc(num_items, size);
-  }
-#endif  // defined(ADDRESS_SANITIZER)
-
-  return *result != NULL;
-}
-
 void ReplaceZoneFunctions(ChromeMallocZone* zone,
                           const MallocZoneFunctions* functions) {
   // Remove protection.
@@ -332,6 +302,44 @@ void ReplaceZoneFunctions(ChromeMallocZone* 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
+
+bool UncheckedMallocMac(size_t size, void** result) {
+#if defined(ADDRESS_SANITIZER)
+  *result = malloc(size);
+#else
+  if (g_old_zone.malloc) {
+    *result = g_old_zone.malloc(malloc_default_zone(), size);
+  } else {
+    *result = malloc(size);
+  }
+#endif  // defined(ADDRESS_SANITIZER)
+
+  return *result != NULL;
+}
+
+bool UncheckedCallocMac(size_t num_items, size_t size, void** result) {
+#if defined(ADDRESS_SANITIZER)
+  *result = calloc(num_items, size);
+#else
+  if (g_old_zone.calloc) {
+    *result = g_old_zone.calloc(malloc_default_zone(), num_items, size);
+  } else {
+    *result = calloc(num_items, size);
+  }
+#endif  // defined(ADDRESS_SANITIZER)
+
+  return *result != NULL;
+}
+
 void StoreFunctionsForDefaultZone() {
   ChromeMallocZone* default_zone = reinterpret_cast<ChromeMallocZone*>(
       malloc_default_zone());
@@ -355,6 +363,13 @@ void StoreFunctionsForAllZones() {
 }
 
 void ReplaceFunctionsForStoredZones(const MallocZoneFunctions* functions) {
+  ChromeMallocZone* default_zone =

[Primiano Tucci (use gerrit), 2017/03/21 19:39:23]

can you add a comment explaining that the malloc zone is not listed in the
get_all_zones below?

[erikchen, 2017/03/24 21:59:03]

Done.

+      reinterpret_cast<ChromeMallocZone*>(malloc_default_zone());
+  if (IsMallocZoneAlreadyStored(default_zone) &&

[Primiano Tucci (use gerrit), 2017/03/21 19:39:23]

minor comment: you are now repeating this:
IsMallocZoneAlreadyStored(default_zone) && ...!= ...malloc in three places.
Maybe you could factor this, either as a flag of ReplaceZoneFunctions, or by
introducing a MaybeReplaceZoneFunction (or whatever other name)

[erikchen, 2017/03/24 21:59:03]

Done.

+      default_zone->malloc != functions->malloc) {
+    ReplaceZoneFunctions(default_zone, functions);
+  }
+
   vm_address_t* zones;
   unsigned int count;
   kern_return_t kr =
@@ -393,6 +408,7 @@ void InterceptAllocationsMac() {
   if (!IsMallocZoneAlreadyStored(default_zone)) {
     StoreZoneFunctions(default_zone, &g_old_zone);
     MallocZoneFunctions new_functions;
+    memset(&new_functions, 0, sizeof(MallocZoneFunctions));

[Primiano Tucci (use gerrit), 2017/03/21 19:39:23]

Just MallocZoneFunctions new_functions = {} should be enough (and save a
precious line :P)?
same below

[erikchen, 2017/03/24 21:59:03]

Done.

     new_functions.malloc = oom_killer_malloc;
     new_functions.calloc = oom_killer_calloc;
     new_functions.valloc = oom_killer_valloc;
@@ -409,6 +425,7 @@ void InterceptAllocationsMac() {
   if (purgeable_zone && !IsMallocZoneAlreadyStored(purgeable_zone)) {
     StoreZoneFunctions(purgeable_zone, &g_old_purgeable_zone);
     MallocZoneFunctions new_functions;
+    memset(&new_functions, 0, sizeof(MallocZoneFunctions));
     new_functions.malloc = oom_killer_malloc_purgeable;
     new_functions.calloc = oom_killer_calloc_purgeable;
     new_functions.valloc = oom_killer_valloc_purgeable;
@@ -496,5 +513,19 @@ void InterceptAllocationsMac() {
                            reinterpret_cast<IMP>(oom_killer_allocWithZone));
 }
 
+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();
+}
+
 }  // namespace allocator
 }  // namespace base