Merge pull request #182 from NSExceptional/fix-heap

Fix #119

Classes/Utility/FLEXHeapEnumerator.m

@@ -22,8 +22,7 @@ typedef struct {
 
 static void range_callback(task_t task, void *context, unsigned type, vm_range_t *ranges, unsigned rangeCount)
 {
-    flex_object_enumeration_block_t block = (__bridge flex_object_enumeration_block_t)context;
-    if (!block) {
+    if (!context) {
         return;
     }
     
@@ -40,7 +39,7 @@ static void range_callback(task_t task, void *context, unsigned type, vm_range_t
 #endif
         // If the class pointer matches one in our set of class pointers from the runtime, then we should have an object.
         if (CFSetContainsValue(registeredClasses, (__bridge const void *)(tryClass))) {
-            block((__bridge id)tryObject, tryClass);
+            (*(flex_object_enumeration_block_t __unsafe_unretained *)context)((__bridge id)tryObject, tryClass);
         }
     }
 }
@@ -71,8 +70,42 @@ static kern_return_t reader(__unused task_t remote_task, vm_address_t remote_add
     if (result == KERN_SUCCESS) {
         for (unsigned int i = 0; i < zoneCount; i++) {
             malloc_zone_t *zone = (malloc_zone_t *)zones[i];
-            if (zone->introspect && zone->introspect->enumerator) {
-                zone->introspect->enumerator(TASK_NULL, (__bridge void *)block, MALLOC_PTR_IN_USE_RANGE_TYPE, (vm_address_t)zone, reader, &range_callback);
+            malloc_introspection_t *introspection = zone->introspect;
+
+            if (!introspection) {
+                continue;
+            }
+
+            void (*lock_zone)(malloc_zone_t *zone)   = introspection->force_lock;
+            void (*unlock_zone)(malloc_zone_t *zone) = introspection->force_unlock;
+
+            // Callback has to unlock the zone so we freely allocate memory inside the given block
+            flex_object_enumeration_block_t callback = ^(__unsafe_unretained id object, __unsafe_unretained Class actualClass) {
+                unlock_zone(zone);
+                block(object, actualClass);
+                lock_zone(zone);
+            };
+
+            // The largest realistic memory address varies by platform.
+            // Only 48 bits are used by 64 bit machines while
+            // 32 bit machines use all bits.
+#if __arm64__
+            static uintptr_t MAX_REALISTIC_ADDRESS = 0xFFFFFFFFFFFF;
+#else
+            static uintptr_t MAX_REALISTIC_ADDRESS = INT_MAX;
+#endif
+
+            // There is little documentation on when and why
+            // any of these function pointers might be NULL
+            // or garbage, so we resort to checking for NULL
+            // and impossible memory addresses at least
+            if (lock_zone && unlock_zone &&
+                introspection->enumerator &&
+                (uintptr_t)lock_zone < MAX_REALISTIC_ADDRESS &&
+                (uintptr_t)unlock_zone < MAX_REALISTIC_ADDRESS) {
+                lock_zone(zone);
+                introspection->enumerator(TASK_NULL, (void *)&callback, MALLOC_PTR_IN_USE_RANGE_TYPE, (vm_address_t)zone, reader, &range_callback);
+                unlock_zone(zone);
             }
         }
     }