#include <signal.h>
#include <stdio.h>
#include <errno.h>
#include <sys/mman.h>
#include <asm/unistd.h>

struct my_sig_dbg_op {
	int dbg_type;
	unsigned long dbg_value;
};

/* Enable or disable single-stepping.  The value sets the state. */
#define MY_SIG_DBG_SINGLE_STEPPING		1

/* Enable or disable branch tracing.  The value sets the state. */
#define MY_SIG_DBG_BRANCH_TRACING		2

#ifndef __NR_dbg_sigreturn
#define __NR_dbg_sigreturn 256
#endif

/* Create the debug return syscall. */
_syscall3(int, dbg_sigreturn,
	  void *, ucontext,
	  int,    ndbg,
	  struct my_sig_dbg_op *, op);


volatile int called = 0;
volatile int called2 = 0;
volatile int called3 = 0;
volatile int trap_called = 0;

void sighand(int sig, siginfo_t *info, void *ucontext)
{
	struct my_sig_dbg_op op;

	called++;

	kill(getpid(), SIGUSR2);

	op.dbg_type = MY_SIG_DBG_SINGLE_STEPPING;
	op.dbg_value = 1;
	dbg_sigreturn(ucontext, 1, &op);
}

void sighand2(int sig, siginfo_t *info, void *ucontext)
{
	kill(getpid(), SIGPWR);
	called2++;
}

void sighand3(int sig, siginfo_t *info, void *ucontext)
{
	struct my_sig_dbg_op op;
	called3++;

	op.dbg_type = MY_SIG_DBG_SINGLE_STEPPING;
	op.dbg_value = 1;
	dbg_sigreturn(ucontext, 1, &op);
}

#define PAGE_SIZE 4096
#define TO_PAGEBASE(a) (((unsigned int) (a)) & (~(PAGE_SIZE-1)))
#define TRAP_INSTRUCTION  0x0ce00097

static int
write_instruction(unsigned char *address,
		  unsigned long new_instr,
		  unsigned long *old_instr)
{
	char         *pagebase = (char *) TO_PAGEBASE(address);

	/* FIXME - currently assuming read-only executable memory, need a
	   way to fetch the old memory protection. */
	if (mprotect(pagebase,
		     PAGE_SIZE,
		     PROT_READ | PROT_WRITE | PROT_EXEC) != 0) {
		/* Couldn't change memory protection, return an error */
		return -1;
	}

	if (old_instr)
		*old_instr = *((volatile unsigned long *) address);

	*((volatile unsigned long *) address) = new_instr;

	mprotect(pagebase, PAGE_SIZE, PROT_READ | PROT_EXEC);

	/* Flush the cache for the instruction. */
	__asm__ ("dcbst 0,%0\n\ticbi 0,%0" : : "r" (address));

	return 0;
}

unsigned char *instr_addr;
unsigned long old_instr;
int restore = 0;
int tracing = 0;

/* The "old" ucontext. */
struct old_sigcontext_struct {
	unsigned long	_unused[4];
	int		signal;
	unsigned long	handler;
	unsigned long	oldmask;
	struct pt_regs 	*regs;
};
struct old_ucontext {
	unsigned long	  uc_flags;
	struct ucontext  *uc_link;
	stack_t		  uc_stack;
	struct sigcontext_struct uc_mcontext;
	sigset_t	  uc_sigmask;	/* mask last for extensibility */
};

unsigned long
dbg_get_instruction_ptr_from_ucontext(void *ucontext)
{
    struct old_ucontext *uc = ucontext;
    struct pt_regs      *regs = uc->uc_mcontext.regs;
    return regs->nip;
}

void sigtrap(int sig, siginfo_t *info, void *ucontext)
{
	int old_errno = errno;
	trap_called++;
	if (restore) {
		write_instruction(instr_addr, old_instr, NULL);
		restore = 0;
	}
	if (tracing) {
		char buf[100];
		sprintf(buf, "Trap at %8.8x\n", dbg_get_instruction_ptr_from_ucontext(ucontext));
		write(2, buf, strlen(buf));
		if (trap_called > 20)
			tracing = 0;
		else {
			struct my_sig_dbg_op op;
			errno = old_errno;

			op.dbg_type = MY_SIG_DBG_BRANCH_TRACING;
			op.dbg_value = 1;
			dbg_sigreturn(ucontext, 1, &op);
		}
	}
	errno = old_errno;
}

void call_printf(void)
{
	printf("test\n");
}

int main(int argc, char *argv)
{
	struct sigaction act;
	int rv;

	act.sa_sigaction = sighand;
	act.sa_flags = SA_SIGINFO;
	sigemptyset(&act.sa_mask);
	sigaddset(&act.sa_mask, SIGUSR2);
	rv = sigaction(SIGUSR1, &act, NULL);
	if (rv == -1) {
		perror("sigaction");
		exit(1);
	}

	act.sa_sigaction = sighand2;
	rv = sigaction(SIGUSR2, &act, NULL);
	if (rv == -1) {
		perror("sigaction");
		exit(1);
	}

	act.sa_sigaction = sighand3;
	rv = sigaction(SIGPWR, &act, NULL);
	if (rv == -1) {
		perror("sigaction");
		exit(1);
	}

	act.sa_sigaction = sigtrap;
	rv = sigaction(SIGTRAP, &act, NULL);
	if (rv == -1) {
		perror("sigaction");
		exit(1);
	}

	kill(getpid(), SIGUSR1);

	if (!called)
		printf("Didn't get called\n");
	else
		printf("Got called %d times\n", called);

	if (!called2)
		printf("Didn't get called 2\n");
	else
		printf("Got called 2 %d times\n", called2);

	if (!called3)
		printf("Didn't get called 3\n");
	else
		printf("Got called 3 %d times\n", called3);

	if (!trap_called) {
		printf("ERROR: Didn't get trapped\n");
		exit(1);
	} else
		printf("Got trapped %d times\n", trap_called);

	instr_addr = (unsigned char *) call_printf;

	write_instruction(instr_addr, TRAP_INSTRUCTION, &old_instr);
	restore = 1;
	tracing = 1;
	call_printf();

	exit(0);
}