Hook malloc

mallocwrap.c

// profiling_zone.c
#include <malloc/malloc.h>
#include <dlfcn.h>
#include <stdlib.h>
#include <unistd.h>

// Profiling callback function pointers
static void (*profiling_alloc_callback)(void* ptr, size_t reserved_size, size_t used_size) = NULL;
static void (*profiling_free_callback)(void* ptr, size_t reserved_size, size_t used_size) = NULL;

// System zone pointer
static malloc_zone_t* system_zone = NULL;

// System allocator function pointers
static void* (*system_malloc)(malloc_zone_t*, size_t) = NULL;
static void (*system_free)(malloc_zone_t*, void*) = NULL;
static void* (*system_realloc)(malloc_zone_t*, void*, size_t) = NULL;
static void* (*system_calloc)(malloc_zone_t*, size_t, size_t) = NULL;
static size_t (*system_malloc_size)(malloc_zone_t*, const void*) = NULL;

// Our profiling zone
static malloc_zone_t profiling_zone;
static malloc_introspection_t zone_introspect;

// Introspection function wrappers
static kern_return_t zone_enumerator(task_t task, void* context, unsigned type_mask,
                                      vm_address_t zone_address, memory_reader_t reader,
                                      vm_range_recorder_t recorder) {
    return system_zone->introspect->enumerator(task, context, type_mask, zone_address, reader, recorder);
}

static size_t zone_good_size(malloc_zone_t* zone, size_t size) {
    return system_zone->introspect->good_size(system_zone, size);
}

static boolean_t zone_check(malloc_zone_t* zone) {
    return system_zone->introspect->check(system_zone);
}

static void zone_print(malloc_zone_t* zone, boolean_t verbose) {
    system_zone->introspect->print(system_zone, verbose);
}

static void zone_log(malloc_zone_t* zone, void* address) {
    system_zone->introspect->log(system_zone, address);
}

static void zone_force_lock(malloc_zone_t* zone) {
    system_zone->introspect->force_lock(system_zone);
}

static void zone_force_unlock(malloc_zone_t* zone) {
    system_zone->introspect->force_unlock(system_zone);
}

static void zone_statistics(malloc_zone_t* zone, malloc_statistics_t* stats) {
    system_zone->introspect->statistics(system_zone, stats);
}

static boolean_t zone_locked(malloc_zone_t* zone) {
    return system_zone->introspect->zone_locked(system_zone);
}

static void zone_reinit_lock(malloc_zone_t* zone) {
    system_zone->introspect->reinit_lock(system_zone);
}

// Zone function implementations
static void* zone_malloc(malloc_zone_t* zone, size_t size) {
    write(1, "mallc\n", sizeof("malloc\n")-1);
    void* ptr = system_malloc(system_zone, size);
    if (ptr && profiling_alloc_callback) {
        size_t actual_size = system_malloc_size(system_zone, ptr);
        profiling_alloc_callback(ptr, actual_size, size);
    }
    return ptr;
}

static void zone_free(malloc_zone_t* zone, void* ptr) {
    if (ptr && profiling_free_callback) {
        size_t actual_size = system_malloc_size(system_zone, ptr);
        profiling_free_callback(ptr, actual_size, actual_size);
    }
    system_free(system_zone, ptr);
}

static void* zone_realloc(malloc_zone_t* zone, void* ptr, size_t size) {
    size_t old_size = 0;
    if (ptr && profiling_free_callback) {
        old_size = system_malloc_size(system_zone, ptr);
        profiling_free_callback(ptr, old_size, old_size);
    }

    void* new_ptr = system_realloc(system_zone, ptr, size);
    if (new_ptr && profiling_alloc_callback) {
        size_t actual_size = system_malloc_size(system_zone, new_ptr);
        profiling_alloc_callback(new_ptr, actual_size, size);
    }
    return new_ptr;
}

static void* zone_calloc(malloc_zone_t* zone, size_t num, size_t size) {
    void* ptr = system_calloc(system_zone, num, size);
    if (ptr && profiling_alloc_callback) {
        size_t actual_size = system_malloc_size(system_zone, ptr);
        profiling_alloc_callback(ptr, actual_size, num * size);
    }
    return ptr;
}

static size_t zone_size(malloc_zone_t* zone, const void* ptr) {
    return system_malloc_size(system_zone, ptr);
}

static void* zone_valloc(malloc_zone_t* zone, size_t size) {
    void* ptr = system_zone->valloc(system_zone, size);
    if (ptr && profiling_alloc_callback) {
        size_t actual_size = system_malloc_size(system_zone, ptr);
        profiling_alloc_callback(ptr, actual_size, size);
    }
    return ptr;
}

static void zone_destroy(malloc_zone_t* zone) {
    system_zone->destroy(system_zone);
}

static unsigned zone_batch_malloc(malloc_zone_t* zone, size_t size, void** results, unsigned num_requested) {
    unsigned allocated = system_zone->batch_malloc(system_zone, size, results, num_requested);
    if (profiling_alloc_callback) {
        for (unsigned i = 0; i < allocated; i++) {
            size_t actual_size = system_malloc_size(system_zone, results[i]);
            profiling_alloc_callback(results[i], actual_size, size);
        }
    }
    return allocated;
}

static void zone_batch_free(malloc_zone_t* zone, void** to_be_freed, unsigned num) {
    if (profiling_free_callback) {
        for (unsigned i = 0; i < num; i++) {
            if (to_be_freed[i]) {
                size_t actual_size = system_malloc_size(system_zone, to_be_freed[i]);
                profiling_free_callback(to_be_freed[i], actual_size, actual_size);
            }
        }
    }
    system_zone->batch_free(system_zone, to_be_freed, num);
}

static void* zone_memalign(malloc_zone_t* zone, size_t alignment, size_t size) {
    void* ptr = system_zone->memalign(system_zone, alignment, size);
    if (ptr && profiling_alloc_callback) {
        size_t actual_size = system_malloc_size(system_zone, ptr);
        profiling_alloc_callback(ptr, actual_size, size);
    }
    return ptr;
}

static void zone_free_definite_size(malloc_zone_t* zone, void* ptr, size_t size) {
    if (ptr && profiling_free_callback) {
        size_t actual_size = system_malloc_size(system_zone, ptr);
        profiling_free_callback(ptr, actual_size, actual_size);
    }
    system_zone->free_definite_size(system_zone, ptr, size);
}

static size_t zone_pressure_relief(malloc_zone_t* zone, size_t goal) {
    return system_zone->pressure_relief(system_zone, goal);
}

static malloc_zone_t* get_default_zone() {
  malloc_zone_t** zones = NULL;
  unsigned int num_zones = 0;

  // On OSX 10.12, malloc_default_zone returns a special zone that is not
  // present in the list of registered zones. That zone uses a "lite zone"
  // if one is present (apparently enabled when malloc stack logging is
  // enabled), or the first registered zone otherwise. In practice this
  // means unless malloc stack logging is enabled, the first registered
  // zone is the default.
  // So get the list of zones to get the first one, instead of relying on
  // malloc_default_zone.
  if (KERN_SUCCESS !=
      malloc_get_all_zones(0, NULL, (vm_address_t**)&zones, &num_zones)) {
    // Reset the value in case the failure happened after it was set.
    num_zones = 0;
  }
  if (num_zones) {
    return zones[0];
  }
  return malloc_default_zone();
}

__attribute__((constructor))
static void initialize_profiling_zone(void) {
    // Load profiling library if specified
    const char* profiling_lib = getenv("ALLOC_PROFILING_LIB");
    if (profiling_lib) {
        void* handle = dlopen(profiling_lib, RTLD_LAZY);
        if (handle) {
            profiling_alloc_callback = dlsym(handle, "profiling_alloc");
            profiling_free_callback = dlsym(handle, "profiling_free");
        }
    }

    // Get system allocator functions
    system_zone = get_default_zone();
    system_malloc = system_zone->malloc;
    system_free = system_zone->free;
    system_realloc = system_zone->realloc;
    system_calloc = system_zone->calloc;
    system_malloc_size = system_zone->size;

    // Set up introspection
    zone_introspect.enumerator = zone_enumerator;
    zone_introspect.good_size = zone_good_size;
    zone_introspect.check = zone_check;
    zone_introspect.print = zone_print;
    zone_introspect.log = zone_log;
    zone_introspect.force_lock = zone_force_lock;
    zone_introspect.force_unlock = zone_force_unlock;
    zone_introspect.statistics = zone_statistics;
    zone_introspect.zone_locked = zone_locked;
    zone_introspect.reinit_lock = zone_reinit_lock;

    // Set up our zone
    profiling_zone.size = zone_size;
    profiling_zone.malloc = zone_malloc;
    profiling_zone.calloc = zone_calloc;
    profiling_zone.valloc = zone_valloc;
    profiling_zone.free = zone_free;
    profiling_zone.realloc = zone_realloc;
    profiling_zone.destroy = zone_destroy;
    profiling_zone.zone_name = "allocation_profiling_zone";
    profiling_zone.batch_malloc = zone_batch_malloc;
    profiling_zone.batch_free = zone_batch_free;
    profiling_zone.introspect = &zone_introspect;
    profiling_zone.version = 9;
    profiling_zone.memalign = zone_memalign;
    profiling_zone.free_definite_size = zone_free_definite_size;
    profiling_zone.pressure_relief = zone_pressure_relief;

    // Register zone and make it default (following Zone.c pattern)
    malloc_zone_t* default_zone = get_default_zone();
    malloc_zone_t* purgeable_zone = malloc_default_purgeable_zone();

    malloc_zone_register(&profiling_zone);

    do {
        malloc_zone_unregister(default_zone);
        malloc_zone_register(default_zone);
        malloc_zone_unregister(purgeable_zone);
        malloc_zone_register(purgeable_zone);
        default_zone = malloc_default_zone();
    } while (default_zone != &profiling_zone && false);
}