/*
 * heaplist is meant to be an LD_PRELOAD library
 * that tracks all memory allocations and is locatable
 * by a signature in memory.
 */
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <dlfcn.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <pthread.h>

#ifndef SPARC
extern void *_dl_open(const char *name, int mode, int caller) __attribute__((regparm(3)));
extern void *_dl_sym(void *handle, const char *sym) __attribute__((regparm(2)));
#else
extern void *_dl_open(const char *name, int mode, int caller);
extern void *_dl_sym(void *handle, const char *sym);
#endif

typedef struct _pointer_list {

	void                 *ptr;
	size_t               size;
	struct _pointer_list *next;

} PointerList;

typedef struct _heap_list_state {

	unsigned char sig[16];
	PointerList   *list;

} HeapListState;

void memory_initialize();
void memory_push(void *ptr, size_t size);
void memory_remove(void *ptr);

void *(*orig_malloc)(size_t)         = NULL;
void *(*orig_realloc)(void *,size_t) = NULL;
void (*orig_free)(void *ptr)         = NULL;

int init                   = 0;
int disable                = 0;
HeapListState *state       = NULL;
pthread_mutex_t list_mutex = PTHREAD_MUTEX_INITIALIZER;

void __attribute__((constructor)) __init()
{
}

void __attribute__((destructor)) __fini()
{
}

void *malloc(size_t size)
{
	void *ptr;

	if (!init && !disable)
		memory_initialize();

	if (disable)
	{
#ifdef MAP_ANON
		ptr = mmap(NULL, size + sizeof(size_t), PROT_READ|PROT_WRITE, MAP_PRIVATE | MAP_ANON, -1, 0);
#else
		{
			int fd = open("/dev/null", O_RDWR);
			ptr = mmap(NULL, size + sizeof(size_t), PROT_READ|PROT_WRITE, MAP_PRIVATE, fd, 0);
			close(fd);
		}
#endif

		*(size_t *)ptr = size;

		ptr = (unsigned char *)ptr + sizeof(size_t);
	}
	else
	{
		ptr = orig_malloc(size);

		if (ptr)
			memory_push(ptr, size);
	}

	return ptr;
}

void *realloc(void *ptr, size_t size)
{
	void *newPtr;

	if (!init && !disable)
		memory_initialize();

	if (disable)
	{
		newPtr = malloc(size);

		if (newPtr && ptr)
		{
			size_t origSize = *(size_t *)((char *)ptr - sizeof(size_t));

			memcpy(newPtr, ptr, origSize);

			free(ptr);
		}
	}
	else
	{
		newPtr = orig_realloc(ptr, size);

		if ((newPtr) &&
		    (newPtr != ptr))
			memory_push(newPtr, size);
	}

	return newPtr;
}

void free(void *ptr)
{
	if (!init && !disable)
		memory_initialize();

	if (disable)
	{
		void *real = (unsigned char *)ptr - sizeof(size_t);

		munmap(real, *(size_t *)real);
	}
	else
	{
		orig_free(ptr);

		memory_remove(ptr);
	}
}

void memory_initialize()
{
	void *base = NULL;

	disable = 1;
	
	base = _dl_open("/lib/libc.so.6", RTLD_LAZY | 0x80000000, 0);

	orig_malloc  = (void *(*)(size_t))_dl_sym(base, "malloc");
	orig_realloc = (void *(*)(void *, size_t))_dl_sym(base, "realloc");
	orig_free    = (void (*)(void *))_dl_sym(base, "free");

	disable = 0;
	init    = 1;

	state = (HeapListState *)orig_malloc(sizeof(HeapListState));

	if (state)
	{
		memset(state, 0, sizeof(HeapListState));
		memcpy(state->sig, "HEAPLISTSIG", 11);
	
		printf("[+] heaplist.so loaded.\n\n");
	}
	else
		printf("[-] heaplist.so NOT loaded successfully.\n\n");
}

void memory_push(void *ptr, size_t size)
{
	PointerList *newEntry;

	if (!init || !state)
		return;
	
	newEntry = (PointerList *)orig_malloc(sizeof(PointerList));

	newEntry->ptr  = ptr;
	newEntry->size = size;
	newEntry->next = NULL;

	pthread_mutex_lock(&list_mutex);
	newEntry->next = state->list;
	state->list    = newEntry;
	pthread_mutex_unlock(&list_mutex);
}

void memory_remove(void *ptr)
{
	PointerList *curr, *prev = NULL;

	if (!init || !state)
		return;
	
	pthread_mutex_lock(&list_mutex);

	for (curr = state->list, prev = NULL;
	     curr;
		  prev = curr, curr = curr->next)
	{
		if (curr->ptr == ptr)
			break;
	}

	if (!curr)
	{
		pthread_mutex_unlock(&list_mutex);
		return;
	}	

	if (prev)
		prev->next  = curr->next;
	else
		state->list = curr->next;

	orig_free(curr);

	pthread_mutex_unlock(&list_mutex);
}