/*
 * SPDX-License-Identifier: CDDL 1.0
 *
 * Copyright (c) 2022 Christos Margiolis <christos@FreeBSD.org>
 * Copyright (c) 2022 Mark Johnston <markj@FreeBSD.org>
 * Copyright (c) 2023 The FreeBSD Foundation
 *
 * Portions of this software were developed by Christos Margiolis
 * <christos@FreeBSD.org> under sponsorship from the FreeBSD Foundation.
 */

#include <sys/param.h>
#include <sys/pcpu.h>

#include <machine/cpufunc.h>
#include <machine/md_var.h>

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

#include "kinst.h"

#define KINST_PUSHL_RBP         0x55
#define KINST_STI               0xfb
#define KINST_POPF              0x9d

#define KINST_MODRM_MOD(b)      (((b) & 0xc0) >> 6)
#define KINST_MODRM_REG(b)      (((b) & 0x38) >> 3)
#define KINST_MODRM_RM(b)       ((b) & 0x07)

#define KINST_SIB_SCALE(s)      (((s) & 0xc0) >> 6)
#define KINST_SIB_INDEX(s)      (((s) & 0x38) >> 3)
#define KINST_SIB_BASE(s)       (((s) & 0x07) >> 0)

#define KINST_REX_W(r)          (((r) & 0x08) >> 3)
#define KINST_REX_R(r)          (((r) & 0x04) >> 2)
#define KINST_REX_X(r)          (((r) & 0x02) >> 1)
#define KINST_REX_B(r)          (((r) & 0x01) >> 0)

#define KINST_F_CALL            0x0001  /* instruction is a "call" */
#define KINST_F_DIRECT_CALL     0x0002  /* instruction is a direct call */
#define KINST_F_RIPREL          0x0004  /* instruction is position-dependent */
#define KINST_F_JMP             0x0008  /* instruction is a %rip-relative jmp */
#define KINST_F_MOD_DIRECT      0x0010  /* operand is not a memory address */

/*
 * Per-CPU trampolines used when the interrupted thread is executing with
 * interrupts disabled.  If an interrupt is raised while executing a trampoline,
 * the interrupt thread cannot safely overwrite its trampoline if it hits a
 * kinst probe while executing the interrupt handler.
 */
DPCPU_DEFINE_STATIC(uint8_t *, intr_tramp);

/*
 * Map ModR/M register bits to a trapframe offset.
 */
static int
kinst_regoff(int reg)
{
#define _MATCH_REG(i, reg)                      \
        case i:                                 \
                return (offsetof(struct trapframe, tf_ ## reg) / \
                    sizeof(register_t))
        switch (reg) {
        _MATCH_REG( 0, rax);
        _MATCH_REG( 1, rcx);
        _MATCH_REG( 2, rdx);
        _MATCH_REG( 3, rbx);
        _MATCH_REG( 4, rsp); /* SIB when mod != 3 */
        _MATCH_REG( 5, rbp);
        _MATCH_REG( 6, rsi);
        _MATCH_REG( 7, rdi);
        _MATCH_REG( 8, r8); /* REX.R is set */
        _MATCH_REG( 9, r9);
        _MATCH_REG(10, r10);
        _MATCH_REG(11, r11);
        _MATCH_REG(12, r12);
        _MATCH_REG(13, r13);
        _MATCH_REG(14, r14);
        _MATCH_REG(15, r15);
        }
#undef _MATCH_REG
        panic("%s: unhandled register index %d", __func__, reg);
}

/*
 * Obtain the specified register's value.
 */
static uint64_t
kinst_regval(struct trapframe *frame, int reg)
{
        if (reg == -1)
                return (0);
        return (((register_t *)frame)[kinst_regoff(reg)]);
}

static uint32_t
kinst_riprel_disp(struct kinst_probe *kp, void *dst)
{
        return ((uint32_t)((intptr_t)kp->kp_patchpoint + kp->kp_md.disp -
            (intptr_t)dst));
}

static void
kinst_trampoline_populate(struct kinst_probe *kp, uint8_t *tramp)
{
        uint8_t *instr;
        uint32_t disp;
        int ilen;

        ilen = kp->kp_md.tinstlen;

        kinst_memcpy(tramp, kp->kp_md.template, ilen);
        if ((kp->kp_md.flags & KINST_F_RIPREL) != 0) {
                disp = kinst_riprel_disp(kp, tramp);
                kinst_memcpy(&tramp[kp->kp_md.dispoff], &disp, sizeof(uint32_t));
        }

        /*
         * The following position-independent jmp takes us back to the
         * original code.  It is encoded as "jmp *0(%rip)" (six bytes),
         * followed by the absolute address of the instruction following
         * the one that was traced (eight bytes).
         */
        tramp[ilen + 0] = 0xff;
        tramp[ilen + 1] = 0x25;
        tramp[ilen + 2] = 0x00;
        tramp[ilen + 3] = 0x00;
        tramp[ilen + 4] = 0x00;
        tramp[ilen + 5] = 0x00;
        instr = kp->kp_patchpoint + kp->kp_md.instlen;
        kinst_memcpy(&tramp[ilen + 6], &instr, sizeof(uintptr_t));
}

int
kinst_invop(uintptr_t addr, struct trapframe *frame, uintptr_t scratch)
{
        solaris_cpu_t *cpu;
        uintptr_t *stack, retaddr;
        struct kinst_probe *kp;
        struct kinst_probe_md *kpmd;
        uint8_t *tramp;

        stack = (uintptr_t *)frame->tf_rsp;
        cpu = &solaris_cpu[curcpu];

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

        /*
         * Report the address of the breakpoint for the benefit of consumers
         * fetching register values with regs[].
         */
        frame->tf_rip--;

        DTRACE_CPUFLAG_SET(CPU_DTRACE_NOFAULT);
        cpu->cpu_dtrace_caller = stack[0];
        DTRACE_CPUFLAG_CLEAR(CPU_DTRACE_NOFAULT | CPU_DTRACE_BADADDR);
        dtrace_probe(kp->kp_id, 0, 0, 0, 0, 0);
        cpu->cpu_dtrace_caller = 0;

        kpmd = &kp->kp_md;
        if ((kpmd->flags & KINST_F_CALL) != 0) {
                /*
                 * dtrace_invop_start() reserves space on the stack to
                 * store the return address of the call instruction.
                 */
                retaddr = (uintptr_t)(kp->kp_patchpoint + kpmd->instlen);
                *(uintptr_t *)scratch = retaddr;

                if ((kpmd->flags & KINST_F_DIRECT_CALL) != 0) {
                        frame->tf_rip = (uintptr_t)(kp->kp_patchpoint +
                            kpmd->disp + kpmd->instlen);
                } else {
                        register_t rval;

                        if (kpmd->reg1 == -1 && kpmd->reg2 == -1) {
                                /* rip-relative */
                                rval = frame->tf_rip + kpmd->instlen;
                        } else {
                                /* indirect */
                                rval = kinst_regval(frame, kpmd->reg1) +
                                    (kinst_regval(frame, kpmd->reg2) <<
                                    kpmd->scale);
                        }

                        if ((kpmd->flags & KINST_F_MOD_DIRECT) != 0) {
                                frame->tf_rip = rval + kpmd->disp;
                        } else {
                                frame->tf_rip =
                                    *(uintptr_t *)(rval + kpmd->disp);
                        }
                }
                return (DTRACE_INVOP_CALL);
        } else {
                if ((frame->tf_rflags & PSL_I) == 0)
                        tramp = DPCPU_GET(intr_tramp);
                else
                        tramp = curthread->t_kinst_tramp;
                if (tramp == NULL) {
                        /*
                         * A trampoline allocation failed, so this probe is
                         * effectively disabled.  Restore the original
                         * instruction.
                         *
                         * We can't safely print anything here, but the
                         * trampoline allocator should have left a breadcrumb in
                         * the dmesg.
                         */
                        kinst_patch_tracepoint(kp, kp->kp_savedval);
                        frame->tf_rip = (register_t)kp->kp_patchpoint;
                } else {
                        kinst_trampoline_populate(kp, tramp);
                        frame->tf_rip = (register_t)tramp;
                }
                return (DTRACE_INVOP_NOP);
        }
}

void
kinst_patch_tracepoint(struct kinst_probe *kp, kinst_patchval_t val)
{
        register_t reg;
        int oldwp;

        reg = intr_disable();
        oldwp = disable_wp();
        *kp->kp_patchpoint = val;
        restore_wp(oldwp);
        intr_restore(reg);
}

static void
kinst_set_disp8(struct kinst_probe *kp, uint8_t byte)
{
        kp->kp_md.disp = (int64_t)(int8_t)byte;
}

static void
kinst_set_disp32(struct kinst_probe *kp, uint8_t *bytes)
{
        int32_t disp32;

        memcpy(&disp32, bytes, sizeof(disp32));
        kp->kp_md.disp = (int64_t)disp32;
}

/*
 * Set up all of the state needed to faithfully execute a probed instruction.
 *
 * In the simple case, we copy the instruction unmodified to a per-thread
 * trampoline, wherein it is followed by a jump back to the original code.
 * - Instructions can have %rip as an operand:
 *   - with %rip-relative addressing encoded in ModR/M, or
 *   - implicitly as a part of the instruction definition (jmp, call).
 * - Call instructions (which may be %rip-relative) need to push the correct
 *   return address onto the stack.
 *
 * Call instructions are simple enough to be emulated in software, so we simply
 * do not use the trampoline mechanism in that case.  kinst_invop() will compute
 * the branch target using the address info computed here (register operands and
 * displacement).
 *
 * %rip-relative operands encoded using the ModR/M byte always use a 32-bit
 * displacement; when populating the trampoline the displacement is adjusted to
 * be relative to the trampoline address.  Trampolines are always allocated
 * above KERNBASE for this reason.
 *
 * For other %rip-relative operands (just jumps) we take the same approach.
 * Instructions which specify an 8-bit displacement must be rewritten to use a
 * 32-bit displacement.
 */
static int
kinst_instr_dissect(struct kinst_probe *kp, uint8_t **instr)
{
        struct kinst_probe_md *kpmd;
        dis86_t d86;
        uint8_t *bytes, modrm, rex;
        int dispoff, i, ilen, opcidx;

        kpmd = &kp->kp_md;

        d86.d86_data = instr;
        d86.d86_get_byte = dtrace_dis_get_byte;
        d86.d86_check_func = NULL;
        if (dtrace_disx86(&d86, SIZE64) != 0) {
                KINST_LOG("failed to disassemble instruction at: %p", *instr);
                return (EINVAL);
        }
        bytes = d86.d86_bytes;
        kpmd->instlen = kpmd->tinstlen = d86.d86_len;

        /*
         * Skip over prefixes, save REX.
         */
        rex = 0;
        for (i = 0; i < kpmd->instlen; i++) {
                switch (bytes[i]) {
                case 0xf0 ... 0xf3:
                        /* group 1 */
                        continue;
                case 0x26:
                case 0x2e:
                case 0x36:
                case 0x3e:
                case 0x64:
                case 0x65:
                        /* group 2 */
                        continue;
                case 0x66:
                        /* group 3 */
                        continue;
                case 0x67:
                        /* group 4 */
                        continue;
                case 0x40 ... 0x4f:
                        /* REX */
                        rex = bytes[i];
                        continue;
                }
                break;
        }
        KASSERT(i < kpmd->instlen,
            ("%s: failed to disassemble instruction at %p", __func__, bytes));
        opcidx = i;

        /*
         * Identify instructions of interest by opcode: calls and jumps.
         * Extract displacements.
         */
        dispoff = -1;
        switch (bytes[opcidx]) {
        case 0x0f:
                switch (bytes[opcidx + 1]) {
                case 0x80 ... 0x8f:
                        /* conditional jmp near */
                        kpmd->flags |= KINST_F_JMP | KINST_F_RIPREL;
                        dispoff = opcidx + 2;
                        kinst_set_disp32(kp, &bytes[dispoff]);
                        break;
                }
                break;
        case 0xe3:
                /*
                 * There is no straightforward way to translate this instruction
                 * to use a 32-bit displacement.  Fortunately, it is rarely
                 * used.
                 */
                return (EINVAL);
        case 0x70 ... 0x7f:
                /* conditional jmp short */
                kpmd->flags |= KINST_F_JMP | KINST_F_RIPREL;
                dispoff = opcidx + 1;
                kinst_set_disp8(kp, bytes[dispoff]);
                break;
        case 0xe9:
                /* unconditional jmp near */
                kpmd->flags |= KINST_F_JMP | KINST_F_RIPREL;
                dispoff = opcidx + 1;
                kinst_set_disp32(kp, &bytes[dispoff]);
                break;
        case 0xeb:
                /* unconditional jmp short */
                kpmd->flags |= KINST_F_JMP | KINST_F_RIPREL;
                dispoff = opcidx + 1;
                kinst_set_disp8(kp, bytes[dispoff]);
                break;
        case 0xe8:
        case 0x9a:
                /* direct call */
                kpmd->flags |= KINST_F_CALL | KINST_F_DIRECT_CALL;
                dispoff = opcidx + 1;
                kinst_set_disp32(kp, &bytes[dispoff]);
                break;
        case 0xff:
                KASSERT(d86.d86_got_modrm,
                    ("no ModR/M byte for instr at %p", *instr - kpmd->instlen));
                switch (KINST_MODRM_REG(bytes[d86.d86_rmindex])) {
                case 0x02:
                case 0x03:
                        /* indirect call */
                        kpmd->flags |= KINST_F_CALL;
                        break;
                case 0x04:
                case 0x05:
                        /* indirect jump */
                        kpmd->flags |= KINST_F_JMP;
                        break;
                }
        }

        /*
         * If there's a ModR/M byte, we need to check it to see if the operand
         * is %rip-relative, and rewrite the displacement if so.  If not, we
         * might still have to extract operand info if this is a call
         * instruction.
         */
        if (d86.d86_got_modrm) {
                uint8_t mod, rm, sib;

                kpmd->reg1 = kpmd->reg2 = -1;

                modrm = bytes[d86.d86_rmindex];
                mod = KINST_MODRM_MOD(modrm);
                rm = KINST_MODRM_RM(modrm);
                if (mod == 0 && rm == 5) {
                        kpmd->flags |= KINST_F_RIPREL;
                        dispoff = d86.d86_rmindex + 1;
                        kinst_set_disp32(kp, &bytes[dispoff]);
                } else if ((kpmd->flags & KINST_F_CALL) != 0) {
                        bool havesib;

                        havesib = (mod != 3 && rm == 4);
                        dispoff = d86.d86_rmindex + (havesib ? 2 : 1);
                        if (mod == 1)
                                kinst_set_disp8(kp, bytes[dispoff]);
                        else if (mod == 2)
                                kinst_set_disp32(kp, &bytes[dispoff]);
                        else if (mod == 3)
                                kpmd->flags |= KINST_F_MOD_DIRECT;

                        if (havesib) {
                                sib = bytes[d86.d86_rmindex + 1];
                                if (KINST_SIB_BASE(sib) != 5) {
                                        kpmd->reg1 = KINST_SIB_BASE(sib) |
                                            (KINST_REX_B(rex) << 3);
                                }
                                kpmd->scale = KINST_SIB_SCALE(sib);
                                kpmd->reg2 = KINST_SIB_INDEX(sib) |
                                    (KINST_REX_X(rex) << 3);
                        } else {
                                kpmd->reg1 = rm | (KINST_REX_B(rex) << 3);
                        }
                }
        }

        /*
         * Calls are emulated in software; once operands are decoded we have
         * nothing else to do.
         */
        if ((kpmd->flags & KINST_F_CALL) != 0)
                return (0);

        /*
         * Allocate and populate an instruction trampoline template.
         *
         * Position-independent instructions can simply be copied, but
         * position-dependent instructions require some surgery: jump
         * instructions with an 8-bit displacement need to be converted to use a
         * 32-bit displacement, and the adjusted displacement needs to be
         * computed.
         */
        ilen = kpmd->instlen;
        if ((kpmd->flags & KINST_F_RIPREL) != 0) {
                if ((kpmd->flags & KINST_F_JMP) == 0 ||
                    bytes[opcidx] == 0x0f ||
                    bytes[opcidx] == 0xe9 ||
                    bytes[opcidx] == 0xff) {
                        memcpy(kpmd->template, bytes, dispoff);
                        memcpy(&kpmd->template[dispoff + 4],
                            &bytes[dispoff + 4], ilen - (dispoff + 4));
                        kpmd->dispoff = dispoff;
                } else if (bytes[opcidx] == 0xeb) {
                        memcpy(kpmd->template, bytes, opcidx);
                        kpmd->template[opcidx] = 0xe9;
                        kpmd->dispoff = opcidx + 1;

                        /* Instruction length changes from 2 to 5. */
                        kpmd->tinstlen = 5;
                        kpmd->disp -= 3;
                } else if (bytes[opcidx] >= 0x70 && bytes[opcidx] <= 0x7f)  {
                        memcpy(kpmd->template, bytes, opcidx);
                        kpmd->template[opcidx] = 0x0f;
                        kpmd->template[opcidx + 1] = bytes[opcidx] + 0x10;
                        kpmd->dispoff = opcidx + 2;

                        /* Instruction length changes from 2 to 6. */
                        kpmd->tinstlen = 6;
                        kpmd->disp -= 4;
                } else {
                        panic("unhandled opcode %#x", bytes[opcidx]);
                }
        } else {
                memcpy(kpmd->template, bytes, ilen);
        }

        return (0);
}

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;
        int error, instrsize, n, off;
        uint8_t *instr, *limit, *tmp;
        bool push_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);

        /*
         * Refuse to instrument functions lacking the usual frame pointer
         * manipulations since they might correspond to exception handlers.
         */
        tmp = instr;
        push_found = false;
        while (tmp < limit) {
                /*
                 * Checking for 'pop %rbp' as well makes the filtering too
                 * strict as it would skip functions that never return (e.g.,
                 * vnlru_proc()).
                 */
                if (*tmp == KINST_PUSHL_RBP) {
                        push_found = true;
                        break;
                }
                tmp += dtrace_instr_size(tmp);
        }
        if (!push_found)
                return (0);

        n = 0;
        while (instr < limit) {
                instrsize = dtrace_instr_size(instr);
                off = (int)(instr - (uint8_t *)symval->value);
                if (pd->kpd_off != -1 && off != pd->kpd_off) {
                        instr += instrsize;
                        continue;
                }

                /*
                 * Check for instructions which may enable interrupts.  Such
                 * instructions are tricky to trace since it is unclear whether
                 * to use the per-thread or per-CPU trampolines.  Since they are
                 * rare, we don't bother to implement special handling for them.
                 *
                 * If the caller specified an offset, return an error, otherwise
                 * silently ignore the instruction so that it remains possible
                 * to enable all instructions in a function.
                 */
                if (instrsize == 1 &&
                    (instr[0] == KINST_POPF || instr[0] == KINST_STI)) {
                        if (pd->kpd_off != -1)
                                return (EINVAL);
                        instr += instrsize;
                        continue;
                }

                /*
                 * 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_savedval = *instr;
                kp->kp_patchval = KINST_PATCHVAL;
                kp->kp_patchpoint = instr;

                error = kinst_instr_dissect(kp, &instr);
                if (error != 0)
                        return (error);

                kinst_probe_create(kp, lf);
        }

        return (0);
}

int
kinst_md_init(void)
{
        uint8_t *tramp;
        int cpu;

        CPU_FOREACH(cpu) {
                tramp = kinst_trampoline_alloc(M_WAITOK);
                if (tramp == NULL)
                        return (ENOMEM);
                DPCPU_ID_SET(cpu, intr_tramp, tramp);
        }

        return (0);
}

void
kinst_md_deinit(void)
{
        uint8_t *tramp;
        int cpu;

        CPU_FOREACH(cpu) {
                tramp = DPCPU_ID_GET(cpu, intr_tramp);
                if (tramp != NULL) {
                        kinst_trampoline_dealloc(tramp);
                        DPCPU_ID_SET(cpu, intr_tramp, NULL);
                }
        }
}

/*
 * Exclude machine-dependent functions that are not safe-to-trace.
 */
bool
kinst_md_excluded(const char *name)
{
        return (false);
}