/*
 * SPDX-License-Identifier: CDDL 1.0
 *
 * Copyright (c) 2023 The FreeBSD Foundation
 *
 * This software was developed by Christos Margiolis <christos@FreeBSD.org>
 * under sponsorship from the FreeBSD Foundation.
 */

#include <sys/param.h>

#include <sys/dtrace.h>
#include <cddl/dev/dtrace/dtrace_cddl.h>

#include "kinst.h"

DPCPU_DEFINE_STATIC(struct kinst_cpu_state, kinst_state);

#define _MATCH_REG(reg) \
        (offsetof(struct trapframe, tf_ ## reg) / sizeof(register_t))

static int
kinst_regoff(struct trapframe *frame, int n)
{
        switch (n) {
        case 0:
                /* There is no zero register in the trapframe structure. */
                return (-1);
        case 1:
                return (_MATCH_REG(ra));
        case 2:
                return (_MATCH_REG(sp));
        case 3:
                return (_MATCH_REG(gp));
        case 4:
                return (_MATCH_REG(tp));
        case 5 ... 7:
                return (_MATCH_REG(t[n - 5]));
        case 8 ... 9:
                return (_MATCH_REG(s[n - 8]));
        case 10 ... 17:
                return (_MATCH_REG(a[n - 10]));
        case 18 ... 27:
                return (_MATCH_REG(s[n - 18 + 2]));
        case 28 ... 31:
                return (_MATCH_REG(t[n - 28 + 3]));
        default:
                panic("%s: unhandled register index %d", __func__, n);
        }
}

static int
kinst_c_regoff(struct trapframe *frame, int n)
{
        switch (n) {
        case 0 ... 1:
                return (_MATCH_REG(s[n]));
        case 2 ... 7:
                return (_MATCH_REG(a[n - 2]));
        default:
                panic("%s: unhandled register index %d", __func__, n);
        }
}

#undef _MATCH_REG

static int
kinst_emulate(struct trapframe *frame, const struct kinst_probe *kp)
{
        kinst_patchval_t instr = kp->kp_savedval;
        register_t prevpc;
        uint64_t imm;
        uint16_t off;
        uint8_t funct;

        if (kp->kp_md.instlen == INSN_SIZE) {
#define rs1_index       ((instr & RS1_MASK) >> RS1_SHIFT)
#define rs2_index       ((instr & RS2_MASK) >> RS2_SHIFT)
#define rd_index        ((instr & RD_MASK) >> RD_SHIFT)
#define rs1             ((register_t *)frame)[kinst_regoff(frame, rs1_index)]
#define rs2             ((register_t *)frame)[kinst_regoff(frame, rs2_index)]
#define rd              ((register_t *)frame)[kinst_regoff(frame, rd_index)]
#define rs1_lval        (rs1_index != 0 ? rs1 : 0)
#define rs2_lval        (rs2_index != 0 ? rs2 : 0)
                switch (instr & 0x7f) {
                case 0b1101111: /* jal */
                        imm = 0;
                        imm |= ((instr >> 21) & 0x03ff) << 1;
                        imm |= ((instr >> 20) & 0x0001) << 11;
                        imm |= ((instr >> 12) & 0x00ff) << 12;
                        imm |= ((instr >> 31) & 0x0001) << 20;
                        if (imm & 0x0000000000100000)
                                imm |= 0xfffffffffff00000;
                        if (rd_index != 0)
                                rd = frame->tf_sepc + INSN_SIZE;
                        frame->tf_sepc += imm;
                        break;
                case 0b1100111: /* jalr */
                        prevpc = frame->tf_sepc;
                        imm = (instr & IMM_MASK) >> IMM_SHIFT;
                        if (imm & 0x0000000000000800)
                                imm |= 0xfffffffffffff000;
                        frame->tf_sepc = (rs1_lval + imm) & ~1;
                        if (rd_index != 0)
                                rd = prevpc + INSN_SIZE;
                        break;
                case 0b1100011: /* branch */
                        imm = 0;
                        imm |= ((instr >> 8) & 0x000f) << 1;
                        imm |= ((instr >> 25) & 0x003f) << 5;
                        imm |= ((instr >> 7) & 0x0001) << 11;
                        imm |= ((instr >> 31) & 0x0001) << 12;
                        if (imm & 0x0000000000001000)
                                imm |= 0xfffffffffffff000;
                        funct = (instr >> 12) & 0x07;
                        switch (funct) {
                        case 0b000:     /* beq */
                                if (rs1_lval == rs2_lval)
                                        frame->tf_sepc += imm;
                                else
                                        frame->tf_sepc += INSN_SIZE;
                                break;
                        case 0b001:     /* bne */
                                if (rs1_lval != rs2_lval)
                                        frame->tf_sepc += imm;
                                else
                                        frame->tf_sepc += INSN_SIZE;
                                break;
                        case 0b100:     /* blt */
                                if ((int64_t)rs1_lval < (int64_t)rs2_lval)
                                        frame->tf_sepc += imm;
                                else
                                        frame->tf_sepc += INSN_SIZE;
                                break;
                        case 0b110:     /* bltu */
                                if ((uint64_t)rs1_lval < (uint64_t)rs2_lval)
                                        frame->tf_sepc += imm;
                                else
                                        frame->tf_sepc += INSN_SIZE;
                                break;
                        case 0b101:     /* bge */
                                if ((int64_t)rs1_lval >= (int64_t)rs2_lval)
                                        frame->tf_sepc += imm;
                                else
                                        frame->tf_sepc += INSN_SIZE;
                                break;
                        case 0b111:     /* bgeu */
                                if ((uint64_t)rs1_lval >= (uint64_t)rs2_lval)
                                        frame->tf_sepc += imm;
                                else
                                        frame->tf_sepc += INSN_SIZE;
                                break;
                        }
                        break;
                case 0b0010111: /* auipc */
                        imm = instr & 0xfffff000;
                        rd = frame->tf_sepc +
                            (imm & 0x0000000080000000 ?
                            imm | 0xffffffff80000000 : imm);
                        frame->tf_sepc += INSN_SIZE;
                        break;
                }
#undef rs1_lval
#undef rs2_lval
#undef rs1
#undef rs2
#undef rd
#undef rs1_index
#undef rs2_index
#undef rd_index
        } else {
                switch (instr & 0x03) {
#define rs1     \
        ((register_t *)frame)[kinst_c_regoff(frame, (instr >> 7) & 0x07)]
                case 0b01:
                        funct = (instr >> 13) & 0x07;
                        switch (funct) {
                        case 0b101:     /* c.j */
                                off = (instr >> 2) & 0x07ff;
                                imm = 0;
                                imm |= ((off >> 1) & 0x07) << 1;
                                imm |= ((off >> 9) & 0x01) << 4;
                                imm |= ((off >> 0) & 0x01) << 5;
                                imm |= ((off >> 5) & 0x01) << 6;
                                imm |= ((off >> 4) & 0x01) << 7;
                                imm |= ((off >> 7) & 0x03) << 8;
                                imm |= ((off >> 6) & 0x01) << 10;
                                imm |= ((off >> 10) & 0x01) << 11;
                                if (imm & 0x0000000000000800)
                                        imm |= 0xfffffffffffff000;
                                frame->tf_sepc += imm;
                                break;
                        case 0b110:     /* c.beqz */
                        case 0b111:     /* c.bnez */
                                imm = 0;
                                imm |= ((instr >> 3) & 0x03) << 1;
                                imm |= ((instr >> 10) & 0x03) << 3;
                                imm |= ((instr >> 2) & 0x01) << 5;
                                imm |= ((instr >> 5) & 0x03) << 6;
                                imm |= ((instr >> 12) & 0x01) << 8;
                                if (imm & 0x0000000000000100)
                                        imm |= 0xffffffffffffff00;
                                if (funct == 0b110 && rs1 == 0)
                                        frame->tf_sepc += imm;
                                else if (funct == 0b111 && rs1 != 0)
                                        frame->tf_sepc += imm;
                                else
                                        frame->tf_sepc += INSN_C_SIZE;
                                break;
                        }
                        break;
#undef rs1
#define rs1_index       ((instr & RD_MASK) >> RD_SHIFT)
#define rs1             ((register_t *)frame)[kinst_regoff(frame, rs1_index)]
                case 0b10:
                        funct = (instr >> 13) & 0x07;
                        if (funct == 0b100 && rs1_index != 0) {
                                /* c.jr/c.jalr */
                                prevpc = frame->tf_sepc;
                                frame->tf_sepc = rs1;
                                if (((instr >> 12) & 0x01) != 0)
                                        frame->tf_ra = prevpc + INSN_C_SIZE;
                        }
                        break;
#undef rs1
#undef rs1_index
                }
        }

        return (MATCH_C_NOP);
}

static int
kinst_jump_next_instr(struct trapframe *frame, const struct kinst_probe *kp)
{
        frame->tf_sepc = (register_t)((const uint8_t *)kp->kp_patchpoint +
            kp->kp_md.instlen);

        return (MATCH_C_NOP);
}

static void
kinst_trampoline_populate(struct kinst_probe *kp)
{
        static uint16_t nop = MATCH_C_NOP;
        static uint32_t ebreak = MATCH_EBREAK;
        int ilen;

        ilen = kp->kp_md.instlen;
        kinst_memcpy(kp->kp_tramp, &kp->kp_savedval, ilen);

        /*
         * Since we cannot encode large displacements in a single instruction
         * in order to encode a far-jump back to the next instruction, and we
         * also cannot clobber a register inside the trampoline, we execute a
         * breakpoint after the copied instruction. kinst_invop() is
         * responsible for detecting this special case and performing the
         * "jump" manually.
         *
         * Add a NOP after a compressed instruction for padding.
         */
        if (ilen == INSN_C_SIZE)
                kinst_memcpy(&kp->kp_tramp[ilen], &nop, INSN_C_SIZE);

        kinst_memcpy(&kp->kp_tramp[INSN_SIZE], &ebreak, INSN_SIZE);

        fence_i();
}

/*
 * There are two ways by which an instruction is traced:
 *
 * - By using the trampoline.
 * - By emulating it in software (see kinst_emulate()).
 *
 * The trampoline is used for instructions that can be copied and executed
 * as-is without additional modification. However, instructions that use
 * PC-relative addressing have to be emulated, because RISC-V doesn't allow
 * encoding of large displacements in a single instruction, and since we cannot
 * clobber a register in order to encode the two-instruction sequence needed to
 * create large displacements, we cannot use the trampoline at all.
 * Fortunately, the instructions are simple enough to be emulated in just a few
 * lines of code.
 *
 * The problem discussed above also means that, unlike amd64, we cannot encode
 * a far-jump back from the trampoline to the next instruction. The mechanism
 * employed to achieve this functionality, is to use a breakpoint instead of a
 * jump after the copied instruction. This breakpoint is detected and handled
 * by kinst_invop(), which performs the jump back to the next instruction
 * manually (see kinst_jump_next_instr()).
 */
int
kinst_invop(uintptr_t addr, struct trapframe *frame, uintptr_t scratch)
{
        solaris_cpu_t *cpu;
        struct kinst_cpu_state *ks;
        const struct kinst_probe *kp;

        ks = DPCPU_PTR(kinst_state);

        /*
         * Detect if the breakpoint was triggered by the trampoline, and
         * manually set the PC to the next instruction.
         */
        if (ks->state == KINST_PROBE_FIRED &&
            addr == (uintptr_t)(ks->kp->kp_tramp + INSN_SIZE)) {
                /*
                 * Restore interrupts if they were enabled prior to the first
                 * breakpoint.
                 */
                if ((ks->status & SSTATUS_SPIE) != 0)
                        frame->tf_sstatus |= SSTATUS_SPIE;
                ks->state = KINST_PROBE_ARMED;
                return (kinst_jump_next_instr(frame, ks->kp));
        }

        LIST_FOREACH(kp, KINST_GETPROBE(addr), kp_hashnext) {
                if ((uintptr_t)kp->kp_patchpoint == addr)
                        break;
        }
        if (kp == NULL)
                return (0);

        cpu = &solaris_cpu[curcpu];
        cpu->cpu_dtrace_caller = addr;
        dtrace_probe(kp->kp_id, 0, 0, 0, 0, 0);
        cpu->cpu_dtrace_caller = 0;

        if (kp->kp_md.emulate)
                return (kinst_emulate(frame, kp));

        ks->state = KINST_PROBE_FIRED;
        ks->kp = kp;

        /*
         * Cache the current SSTATUS and clear interrupts for the
         * duration of the double breakpoint.
         */
        ks->status = frame->tf_sstatus;
        frame->tf_sstatus &= ~SSTATUS_SPIE;
        frame->tf_sepc = (register_t)kp->kp_tramp;

        return (MATCH_C_NOP);
}

void
kinst_patch_tracepoint(struct kinst_probe *kp, kinst_patchval_t val)
{
        switch (kp->kp_patchval) {
        case KINST_C_PATCHVAL:
                *(uint16_t *)kp->kp_patchpoint = (uint16_t)val;
                fence_i();
                break;
        case KINST_PATCHVAL:
                *kp->kp_patchpoint = val;
                fence_i();
                break;
        }
}

static void
kinst_instr_dissect(struct kinst_probe *kp, int instrsize)
{
        struct kinst_probe_md *kpmd;
        kinst_patchval_t instr = kp->kp_savedval;
        uint8_t funct;

        kpmd = &kp->kp_md;
        kpmd->instlen = instrsize;
        kpmd->emulate = false;

        /*
         * The following instructions use PC-relative addressing and need to be
         * emulated in software.
         */
        if (kpmd->instlen == INSN_SIZE) {
                switch (instr & 0x7f) {
                case 0b1101111: /* jal */
                case 0b1100111: /* jalr */
                case 0b1100011: /* branch */
                case 0b0010111: /* auipc */
                        kpmd->emulate = true;
                        break;
                }
        } else {
                switch (instr & 0x03) {
                case 0b01:
                        funct = (instr >> 13) & 0x07;
                        switch (funct) {
                        case 0b101:     /* c.j */
                        case 0b110:     /* c.beqz */
                        case 0b111:     /* c.bnez */
                                kpmd->emulate = true;
                                break;
                        }
                        break;
                case 0b10:
                        funct = (instr >> 13) & 0x07;
                        if (funct == 0b100 &&
                            ((instr >> 7) & 0x1f) != 0 &&
                            ((instr >> 2) & 0x1f) == 0)
                                kpmd->emulate = true;   /* c.jr/c.jalr */
                        break;
                }
        }

        if (!kpmd->emulate)
                kinst_trampoline_populate(kp);
}

static bool
kinst_instr_system(kinst_patchval_t instr)
{
        if (dtrace_match_opcode(instr, MATCH_C_EBREAK, MASK_C_EBREAK) ||
            (instr & 0x7f) == 0b1110011)
                return (true);

        return (false);
}

static bool
kinst_instr_lr(kinst_patchval_t instr)
{
        if (dtrace_match_opcode(instr, MATCH_LR_W, MASK_LR_W) ||
            dtrace_match_opcode(instr, MATCH_LR_D, MASK_LR_D))
                return (true);

        return (false);
}

static bool
kinst_instr_sc(kinst_patchval_t instr)
{
        if (dtrace_match_opcode(instr, MATCH_SC_W, MASK_SC_W) ||
            dtrace_match_opcode(instr, MATCH_SC_D, MASK_SC_D))
                return (true);

        return (false);
}

int
kinst_make_probe(linker_file_t lf, int symindx, linker_symval_t *symval,
    void *opaque)
{
        struct kinst_probe *kp;
        dtrace_kinst_probedesc_t *pd;
        const char *func;
        kinst_patchval_t *insn, v;
        uint8_t *instr, *limit;
        int instrsize, n, off;
        bool lrsc_block, store_found;

        pd = opaque;
        func = symval->name;

        if (kinst_excluded(func))
                return (0);
        if (strcmp(func, pd->kpd_func) != 0)
                return (0);

        instr = (uint8_t *)(symval->value);
        limit = (uint8_t *)(symval->value + symval->size);
        if (instr >= limit)
                return (0);

        /* Check for the usual function prologue. */
        store_found = false;
        for (insn = (kinst_patchval_t *)instr;
            insn < (kinst_patchval_t *)limit; insn++) {
                if (dtrace_instr_sdsp(&insn) || dtrace_instr_c_sdsp(&insn)) {
                        store_found = true;
                        break;
                }
        }
        if (!store_found)
                return (0);

        n = 0;
        lrsc_block = false;
        while (instr < limit) {
                instrsize = dtrace_instr_size(instr);
                off = (int)(instr - (uint8_t *)symval->value);

                /*
                 * Avoid undefined behavior (i.e simply casting `*instr` to
                 * `kinst_patchval_t`) in case the pointer is unaligned.
                 * memcpy() can safely operate on unaligned pointers.
                 */
                memcpy(&v, instr, sizeof(kinst_patchval_t));

                /* Skip SYSTEM instructions. */
                if (kinst_instr_system(v))
                        goto cont;

                /*
                 * Skip LR/SC blocks used to build atomic operations. If a
                 * breakpoint is placed in a LR/SC block, the loop becomes
                 * unconstrained. In this case we violate the operation and the
                 * loop might fail on some implementations (see section 8.3 of
                 * the RISC-V unprivileged spec).
                 */
                if (kinst_instr_lr(v))
                        lrsc_block = true;
                else if (kinst_instr_sc(v)) {
                        lrsc_block = false;
                        goto cont;
                }
                if (lrsc_block)
                        goto cont;

                if (pd->kpd_off != -1 && off != pd->kpd_off)
                        goto cont;

                /*
                 * Prevent separate dtrace(1) instances from creating copies of
                 * the same probe.
                 */
                LIST_FOREACH(kp, KINST_GETPROBE(instr), kp_hashnext) {
                        if (strcmp(kp->kp_func, func) == 0 &&
                            strtol(kp->kp_name, NULL, 10) == off)
                                return (0);
                }
                if (++n > KINST_PROBETAB_MAX) {
                        KINST_LOG("probe list full: %d entries", n);
                        return (ENOMEM);
                }
                kp = malloc(sizeof(struct kinst_probe), M_KINST,
                    M_WAITOK | M_ZERO);
                kp->kp_func = func;
                snprintf(kp->kp_name, sizeof(kp->kp_name), "%d", off);
                kp->kp_patchpoint = (kinst_patchval_t *)instr;
                kp->kp_savedval = v;
                if (instrsize == INSN_SIZE)
                        kp->kp_patchval = KINST_PATCHVAL;
                else
                        kp->kp_patchval = KINST_C_PATCHVAL;
                if ((kp->kp_tramp = kinst_trampoline_alloc(M_WAITOK)) == NULL) {
                        KINST_LOG("cannot allocate trampoline for %p", instr);
                        return (ENOMEM);
                }

                kinst_instr_dissect(kp, instrsize);
                kinst_probe_create(kp, lf);
cont:
                instr += instrsize;
        }
        if (lrsc_block)
                KINST_LOG("warning: unterminated LR/SC block");

        return (0);
}

int
kinst_md_init(void)
{
        struct kinst_cpu_state *ks;
        int cpu;

        CPU_FOREACH(cpu) {
                ks = DPCPU_PTR(kinst_state);
                ks->state = KINST_PROBE_ARMED;
        }

        return (0);
}

void
kinst_md_deinit(void)
{
}

/*
 * Exclude machine-dependent functions that are not safe-to-trace.
 */
bool
kinst_md_excluded(const char *name)
{
        if (strcmp(name, "cpu_exception_handler") == 0 ||
            strcmp(name, "cpu_exception_handler_supervisor") == 0 ||
            strcmp(name, "cpu_exception_handler_user") == 0 ||
            strcmp(name, "do_trap_supervisor") == 0 ||
            strcmp(name, "do_trap_user") == 0)
                return (true);

        return (false);
}