#include <stdio.h>
#include <stdlib.h>
#include "x86_codegen.h"
#include <unistd.h>
#include <fcntl.h>
#include <sys/mman.h>

typedef unsigned char byte;
typedef void* (*malloc_fun_t)(size_t size);

/* to work around NX bit restrictions on most
   secure distros */
void *allocate_executable_mem(size_t size)
{
	static int zero_fd = -1;	
	void * addr;
	if(zero_fd == -1)
	{
		/* thread safety is not easy */
		zero_fd = open("/dev/zero", O_RDWR, 0);
	}
	addr = mmap(NULL, size, PROT_READ | PROT_WRITE | PROT_EXEC,
			MAP_SHARED | MAP_ANON, zero_fd, 0);
	return addr;
}

malloc_fun_t make_closure(malloc_fun_t original, 
							const char * filename, int line)
{
	byte * code = allocate_executable_mem(4096);
	byte * method = code;
	byte * target;
	x86_push_reg(code, X86_EBP);
	x86_mov_reg_reg(code, X86_EBP, X86_ESP, 4);

	target = code;
	x86_jump8(code, 0);
	
	/* magic to identify closures */
	*(code++) = 0x42;
	*(code++) = 0x13;
	*(code++) = 0x37;
	*(code++) = 0x42;

	x86_patch(target, code); 

	/* add 8 (two words) to size */
	x86_alu_membase_imm(code, X86_ADD, X86_EBP, 8, 8);
	
	/* frame for next function */
	x86_alu_reg_imm(code, X86_SUB, X86_ESP, 12);

	/* push size */	
	x86_push_membase(code, X86_EBP, 0x8);
	x86_call_code(code, original);
	
	/* pop frame */
	x86_alu_reg_imm(code, X86_ADD, X86_ESP, 12);

	/* return value in eax */
	x86_mov_membase_imm(code, X86_EAX, 0, (int )filename, 4);
	x86_mov_membase_imm(code, X86_EAX, 4, (int)line, 4);
	x86_alu_reg_imm(code, X86_ADD, X86_EAX, 8);

	x86_leave(code);
	x86_ret(code);
	return method;
}

int main()
{
	int i;
	malloc_fun_t f = make_closure(malloc, __FILE__, __LINE__);
	malloc_fun_t f2 = make_closure(malloc, __FILE__, __LINE__);
	void * data = f(9);
	void * data2 = f2(9);
	printf("%p traced to %s:%d\n", data, ((char**)data)[-2], ((int*)data)[-1]);
	printf("%p traced to %s:%d\n", data2, ((char**)data2)[-2], ((int*)data2)[-1]);
}