#include <asc.h>
#include <fcntl.h>
#include <sys/param.h>
#include <sys/ptrace.h>
#include <asm/ptrace-abi.h>
#include <asm/unistd_64.h>
#include <sys/reg.h>
#include <sys/user.h>
#include <sys/wait.h>

#define HAVE_HW_BREAKPOINT_H 1

#ifdef HAVE_HW_BREAKPOINT_H
#include <linux/hw_breakpoint.h>
#endif

static struct perf_event_attr pe;

long transition(struct machine *m, enum mode o)
{
    long i, c[3];
    int status;

    /* Reset error code.  */
    errno = 0;

    /* Fetch instruction count.  */
    if (read(m->fd, c + 0, sizeof(long)) < 0) {
        diagnostic("read");
        goto failure;
    }

    if (o == BASELINE) {
        /* Let the process fly until it makes a system call.  */
        if (ptrace(PTRACE_SYSEMU, m->pid, 0, 0) < 0) {
            diagnostic("ptrace sysemu");
            goto failure;
        }

        /* Re-attach to child process.  */
        if (waitpid(m->pid, &status, 0) == -1) {
            diagnostic("waitpid");
            goto failure;
        }
    } else if (o == SINGLESTEP) {
        /* Execute one instruction.  */
        if (ptrace(PTRACE_SYSEMU_SINGLESTEP, m->pid, 0, 0) < 0) {
            diagnostic("ptrace singlestep");
            goto failure;
        }

        /* Re-attach to child process.  */
        if (waitpid(m->pid, &status, 0) == -1) {
            diagnostic("waitpid");
            goto failure;
        }
    } else if (o == BLOCKSTEP) {
        /* Execute one basic block.  */
        if (ptrace(PTRACE_SINGLEBLOCK, m->pid, 0, 0) < 0) {
            diagnostic("ptrace singleblock");
            goto failure;
        }

        /* Re-attach to child process.  */
        if (waitpid(m->pid, &status, 0) == -1) {
            diagnostic("waitpid");
            goto failure;
        }
    } else if (o == BREAKPOINT) {
        int fd;

        pe.size = sizeof(struct perf_event_attr);
        pe.type = PERF_TYPE_BREAKPOINT;
        pe.bp_type = HW_BREAKPOINT_X;
        pe.bp_len = sizeof(long);

        pe.bp_addr = m->address;
        pe.sample_period = m->period;

        if ((fd = syscall(__NR_perf_event_open, &pe, m->pid, -1, -1, 0)) < 0) {
            diagnostic("perf_event_open");
            goto failure;
        }

        if (fcntl(fd, F_SETFL, O_ASYNC) < 0) {
            diagnostic("fcntl");
            goto failure;
        }

        if (fcntl(fd, F_SETSIG, SIGSTOP) < 0) {
            diagnostic("fcntl");
            goto failure;
        }

        if (fcntl(fd, F_SETOWN, m->pid) < 0) {
            diagnostic("fcntl");
            goto failure;
        }

        /* Let the process fly until it trips the breakpoint K times.  */
        if (ptrace(PTRACE_SYSEMU, m->pid, 0, 0) < 0) {
            diagnostic("ptrace sysemu");
            goto failure;
        }

        /* Re-attach to child process.  */
        if (waitpid(m->pid, &status, 0) == -1) {
            diagnostic("waitpid");
            goto failure;
        }

        if (close(fd) < 0) {
            diagnostic("close");
            goto failure;
        }
    } else if (o == INSTRUCTIONS) {
        /* Let the process fly until overflow interrupt fires.  */
        if (ptrace(PTRACE_SYSEMU, m->pid, 0, 0) < 0) {
            diagnostic("ptrace sysemu");
            goto failure;
        }

        /* Re-attach to child process.  */
        if (waitpid(m->pid, &status, 0) == -1) {
            diagnostic("waitpid");
            goto failure;
        }
    } else if (o == (INSTRUCTIONS | BREAKPOINT)) {
        /* Let the process fly until overflow interrupt fires.  */
        if (ptrace(PTRACE_SYSEMU, m->pid, 0, 0) < 0) {
            diagnostic("ptrace sysemu");
            goto failure;
        }

        /* Re-attach to child process.  */
        if (waitpid(m->pid, &status, 0) == -1) {
            diagnostic("waitpid");
            goto failure;
        }

        /* Fetch virtual register containing system call number if any.  */
        i = ptrace(PTRACE_PEEKUSER, m->pid, 8 * ORIG_RAX, 0);

        /* Bail out on system call or error.  */
        if (i != -1 || errno)
            goto count;

        /* Calculate offset of debug register 0.  */
        i = offsetof(struct user, u_debugreg[0]);

        /* Store designated breakpoint address in debug register 0.  */
        if (ptrace(PTRACE_POKEUSER, m->pid, i, m->address) < 0) {
            diagnostic("ptrace pokeuser");
            goto failure;
        }

        /* Calculate offset of debug register 7.  */
        i = offsetof(struct user, u_debugreg[7]);

        /* Enable designated hardware breakpoint.  */
        if (ptrace(PTRACE_POKEUSER, m->pid, i, 1) < 0) {
            diagnostic("ptrace pokeuser");
            goto failure;
        }

        /* Let the process fly until it hits the breakpoint.  */
        if (ptrace(PTRACE_SYSEMU, m->pid, 0, 0) < 0) {
            diagnostic("ptrace sysemu");
            goto failure;
        }

        /* Re-attach to child process.  */
        if (waitpid(m->pid, &status, 0) == -1) {
            diagnostic("waitpid");
            goto failure;
        }

        /* Disable designated hardware breakpoint.  */
        if (ptrace(PTRACE_POKEUSER, m->pid, i, 0) < 0) {
            diagnostic("ptrace pokeuser");
            goto failure;
        }
    } else {
        diagnostic("illegal transition rule mode combination: 0x%x", o);
        goto failure;
    }

    /* Fetch new instruction count.  */
  count:
    if (read(m->fd, c + 1, sizeof(long)) < 0) {
        diagnostic("read");
        goto failure;
    }

    /* Compute number of instructions executed in this invocation.  */
    c[2] = c[1] - c[0];

    /* Compute cumulative instruction count.  */
    m->retired += c[2];

    /* Fetch virtual register containing system call number if any.  */
    i = ptrace(PTRACE_PEEKUSER, m->pid, 8 * ORIG_RAX, 0);

    /* Bail out if ptrace error.  */
    if (errno) {
        diagnostic("ptrace peekuser");
        goto failure;
    }

    /* Check for system call.  */
    switch (i) {
    case -1:                   /* No system call attempted.  */
        goto success;
    case __NR_exit:            /* Emulate exit system call.  */
        goto halting;
    default:                   /* Reject every other system call.  */
        diagnostic("illegal system call: 0x%lx", i);
        goto failure;
    }

  success:
    return c[2];

  halting:
    return MIN(-1, c[2]);

  failure:
    return 0;
}