kext_memleak.c

/**
 * Created 18H18  lynnl
 */
#include <mach/mach_types.h>
#include <sys/systm.h>
#include <libkern/OSAtomic.h>

#define INAME "kext-memleak"
#define LOG(fmt, ...) printf(INAME ": " fmt "\n", ##__VA_ARGS__)

static void kern_mref(int opt)
{
    static volatile SInt64 cnt = 0;
    switch (opt) {
    case 0:
        if (OSDecrementAtomic64(&cnt) > 0) return;
        break;
    case 1:
        if (OSIncrementAtomic64(&cnt) >= 0) return;
        break;
    case 2:
        if (cnt == 0) return;
        break;
    }
    panic("\n%s#L%d (potential memleak)  opt: %d cnt: %llu\n",
            __func__, __LINE__, opt, cnt);
}

void *kern_malloc(size_t size, int flags)
{
    /* _MALLOC `type' parameter is a joke */
    void *addr = _MALLOC(size, M_TEMP, flags);
    if (addr != NULL) kern_mref(1);
    return addr;
}

static void *kern_realloc(void *addr, size_t size, int flags)
{
    /* XXX: _REALLOC unavailable in most archs  implement you own */
    void *newaddr = _REALLOC(addr, size, M_TEMP, flags);
    if (!!addr ^ !!newaddr) kern_mref(!!newaddr);
    return newaddr;
}

void kern_mfree(void *addr)
{
    if (addr != NULL) kern_mref(0);
    _FREE(addr, M_TEMP);
}

/* XXX: call when all memory freed */
void kern_memck(void)
{
    kern_mref(2);
}

static unsigned char *nil;
static unsigned char *p;

kern_return_t kext_memleak_start(kmod_info_t *ki, void *d)
{
    nil = kern_malloc(INT_MAX, M_NOWAIT); /* Should fail */
    p = kern_malloc(1, M_WAITOK);
    LOG("nil address %#llx", (uint64_t) nil);
    LOG("byte allocated %#llx %#x", (uint64_t) p, *p);
    return KERN_SUCCESS;
}

kern_return_t kext_memleak_stop(kmod_info_t *ki, void *d)
{
    kern_mfree(nil);
    /* Say, you forget to free the p */

    kern_memck(); /* Panic if memleaked */
    LOG("byte leaked!!! %#llx %#x", (uint64_t) p, *p);
    return KERN_SUCCESS;
}