objc object search

main.m

#import <objc/runtime.h>
#import <mach/mach.h>
#import <malloc/malloc.h>
#import <dlfcn.h>

typedef struct {
    Class targetClass;
    void (^block)(id object, BOOL isSubclass, BOOL *stop);
    dispatch_semaphore_t semaphore;
    BOOL shouldStop;
} EnumerationContext;

static kern_return_t fast_read_memory(task_t task, vm_address_t remote_address, vm_size_t size, void **local_memory) {
    *local_memory = (void *)remote_address;
    return KERN_SUCCESS;
}

static kern_return_t safe_read_memory(vm_address_t address, vm_size_t size, void *buffer) {
    vm_size_t dataCount;
    return vm_read_overwrite(mach_task_self_, address, size, (vm_address_t)buffer, &dataCount);
}

static BOOL is_valid_objc_address(vm_address_t address) {

    uintptr_t potentialIsa = 0;
    if (safe_read_memory(address, sizeof(uintptr_t), &potentialIsa) != KERN_SUCCESS) {
        return NO;
    }

    Dl_info info;
    if (dladdr((const void *)potentialIsa, &info) == 0) {
        return NO;
    }

    const char *className = info.dli_sname;
    if (className != NULL && strlen(className) > 13 && objc_getClass(className + 13) != NULL) {
        return YES;
    }

    return NO;
}

static void range_callback(task_t task, void *context, unsigned type, vm_range_t *ranges, unsigned count) {
    
    EnumerationContext *enumContext = (EnumerationContext *)context;
    for (unsigned i = 0; i < count; i++) {
        void *ptr = (void *)ranges[i].address;

        if (is_valid_objc_address((vm_address_t)ptr) == 0) {
            continue;
        }
        
        id obj = (__bridge id)ptr;
        Class cls = object_getClass(obj);
        
        BOOL isSubclass = (cls != enumContext->targetClass) && class_getSuperclass(cls) == enumContext->targetClass;

        if (cls == enumContext->targetClass || isSubclass) {
            BOOL stop = NO;
            enumContext->block(obj, isSubclass, &stop);
            if (stop) {
                enumContext->shouldStop = YES;
                dispatch_semaphore_signal(enumContext->semaphore);
                return;
            }
        }
    }
}

dispatch_semaphore_t forEveryInstanceOfClass(Class targetClass, void (^block)(id object, BOOL isSubclass, BOOL *stop)) {
    dispatch_semaphore_t semaphore = dispatch_semaphore_create(0);
    
    dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
        vm_address_t *zones = NULL;
        unsigned int zoneCount = 0;
        
        EnumerationContext context = {
            .targetClass = targetClass,
            .block = block,
            .semaphore = semaphore,
            .shouldStop = NO
        };
        
        kern_return_t result = malloc_get_all_zones(mach_task_self_, fast_read_memory, &zones, &zoneCount);
        if (result == KERN_SUCCESS) {

            for (unsigned int i = 0; i < zoneCount && !context.shouldStop; i++) {
                malloc_zone_t *zone = (malloc_zone_t *)zones[i];
                if (zone->introspect && zone->introspect->enumerator) {
                    zone->introspect->enumerator(TASK_NULL, &context, MALLOC_PTR_IN_USE_RANGE_TYPE, (vm_address_t)zone, fast_read_memory, range_callback);
                }
            }
        } else {
            NSLog(@"Failed to get malloc zones");
        }
        
        if (!context.shouldStop) {
            dispatch_semaphore_signal(semaphore);
        }
    });
    
    return semaphore;
}