#undef _GNU_SOURCE
#define _GNU_SOURCE

#include <stdio.h>
#include <sys/mman.h>
#include <stdint.h>
#include <unistd.h>
#include <signal.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ucontext.h>


volatile uint32_t *page = NULL;
static int write_flag = 0;
static void *address_buffer;

static uint32_t read_memory(uint32_t *addr)
{
	static uint32_t foo = 0;
	/* Always return incremented number */
	foo++;
	return foo;
}

static uint32_t write_memory(uint32_t* addr, uint32_t val)
{
	printf("Write @ %p: %u\n", (void *)addr, val);
}


static void segfault_handler(int signal, siginfo_t* sig_info, void *uap)
{
	ucontext_t *context;
	int err;
	uint32_t read_val;

	/* Unhandled error */
	if (sig_info->si_code != SEGV_ACCERR)
		exit(-2);

	context = (ucontext_t *)uap;

	/* Allow access */
	mprotect(page, getpagesize(), PROT_READ | PROT_WRITE);

	/* Read or Write instruction? */
	err = context->uc_mcontext.gregs[REG_ERR];
	write_flag = ((err & 2) ? 1 : 0);

	address_buffer = (void *)sig_info->si_addr;

	/* Handle read */
	if (!write_flag) {
		read_val = read_memory((uint32_t *)address_buffer);
		/* Fill value to memory */
		*((uint32_t *)address_buffer) = read_val;
	}

	/* Set Trap flag. Leads to single step */
	context->uc_mcontext.gregs[REG_EFL] |= 0x100;

}

static void trap_handler(int signal, siginfo_t* sig_info, void *uap)
{
	/* Disable trap */
	((ucontext_t *)uap)->uc_mcontext.gregs[REG_EFL] &= ~(0x100U);

	/* Read written value and call callback */
	if (write_flag)
		write_memory((uint32_t *)address_buffer, *((uint32_t *)address_buffer));

	/* Lock memory */
	mprotect(page, getpagesize(), PROT_NONE);
}

static void register_handler()
{
	struct sigaction action;
	struct sigaction trap;

	memset(&action, 0, sizeof(struct sigaction));
	sigemptyset(&action.sa_mask);
	action.sa_flags = SA_SIGINFO;
	action.sa_sigaction = segfault_handler;

	sigaction(SIGSEGV, &action, NULL);

	memset(&trap, 0, sizeof(struct sigaction));
	sigemptyset(&trap.sa_mask);
	trap.sa_flags = SA_SIGINFO;
	trap.sa_sigaction = trap_handler;

	sigaction(SIGTRAP, &trap, NULL);
}

int main()
{

	/* Allocate memory */
	page = (uint32_t *)mmap(NULL, getpagesize(), PROT_READ | PROT_WRITE, MAP_ANONYMOUS | MAP_SHARED, -1, 0);
	if (!page)
		return -1;
	printf("VMA: %p\n", page);
	*page=12;
	/* Register SIGSEGV handler */
	register_handler();

	/* Protect memory */
	mprotect(page, getpagesize(), PROT_NONE);


	printf("RD1: %d\n", *page);
	printf("RD2: %d\n", *page);
	printf("RD3: %d\n", *page);

	*page = 100;

	return 0;
}