// SPDX-License-Identifier: GPL-2.0
#include <test_util.h>
#include <kvm_util.h>
#include <processor.h>
#include <linux/bitfield.h>

#define MDSCR_KDE       (1 << 13)
#define MDSCR_MDE       (1 << 15)
#define MDSCR_SS        (1 << 0)

#define DBGBCR_LEN8     (0xff << 5)
#define DBGBCR_EXEC     (0x0 << 3)
#define DBGBCR_EL1      (0x1 << 1)
#define DBGBCR_E        (0x1 << 0)
#define DBGBCR_LBN_SHIFT        16
#define DBGBCR_BT_SHIFT         20
#define DBGBCR_BT_ADDR_LINK_CTX (0x1 << DBGBCR_BT_SHIFT)
#define DBGBCR_BT_CTX_LINK      (0x3 << DBGBCR_BT_SHIFT)

#define DBGWCR_LEN8     (0xff << 5)
#define DBGWCR_RD       (0x1 << 3)
#define DBGWCR_WR       (0x2 << 3)
#define DBGWCR_EL1      (0x1 << 1)
#define DBGWCR_E        (0x1 << 0)
#define DBGWCR_LBN_SHIFT        16
#define DBGWCR_WT_SHIFT         20
#define DBGWCR_WT_LINK          (0x1 << DBGWCR_WT_SHIFT)

#define SPSR_D          (1 << 9)
#define SPSR_SS         (1 << 21)

extern unsigned char sw_bp, sw_bp2, hw_bp, hw_bp2, bp_svc, bp_brk, hw_wp, ss_start, hw_bp_ctx;
extern unsigned char iter_ss_begin, iter_ss_end;
static volatile uint64_t sw_bp_addr, hw_bp_addr;
static volatile uint64_t wp_addr, wp_data_addr;
static volatile uint64_t svc_addr;
static volatile uint64_t ss_addr[4], ss_idx;
#define  PC(v)  ((uint64_t)&(v))

#define GEN_DEBUG_WRITE_REG(reg_name)                   \
static void write_##reg_name(int num, uint64_t val)     \
{                                                       \
        switch (num) {                                  \
        case 0:                                         \
                write_sysreg(val, reg_name##0_el1);     \
                break;                                  \
        case 1:                                         \
                write_sysreg(val, reg_name##1_el1);     \
                break;                                  \
        case 2:                                         \
                write_sysreg(val, reg_name##2_el1);     \
                break;                                  \
        case 3:                                         \
                write_sysreg(val, reg_name##3_el1);     \
                break;                                  \
        case 4:                                         \
                write_sysreg(val, reg_name##4_el1);     \
                break;                                  \
        case 5:                                         \
                write_sysreg(val, reg_name##5_el1);     \
                break;                                  \
        case 6:                                         \
                write_sysreg(val, reg_name##6_el1);     \
                break;                                  \
        case 7:                                         \
                write_sysreg(val, reg_name##7_el1);     \
                break;                                  \
        case 8:                                         \
                write_sysreg(val, reg_name##8_el1);     \
                break;                                  \
        case 9:                                         \
                write_sysreg(val, reg_name##9_el1);     \
                break;                                  \
        case 10:                                        \
                write_sysreg(val, reg_name##10_el1);    \
                break;                                  \
        case 11:                                        \
                write_sysreg(val, reg_name##11_el1);    \
                break;                                  \
        case 12:                                        \
                write_sysreg(val, reg_name##12_el1);    \
                break;                                  \
        case 13:                                        \
                write_sysreg(val, reg_name##13_el1);    \
                break;                                  \
        case 14:                                        \
                write_sysreg(val, reg_name##14_el1);    \
                break;                                  \
        case 15:                                        \
                write_sysreg(val, reg_name##15_el1);    \
                break;                                  \
        default:                                        \
                GUEST_ASSERT(0);                        \
        }                                               \
}

/* Define write_dbgbcr()/write_dbgbvr()/write_dbgwcr()/write_dbgwvr() */
GEN_DEBUG_WRITE_REG(dbgbcr)
GEN_DEBUG_WRITE_REG(dbgbvr)
GEN_DEBUG_WRITE_REG(dbgwcr)
GEN_DEBUG_WRITE_REG(dbgwvr)

static void reset_debug_state(void)
{
        uint8_t brps, wrps, i;
        uint64_t dfr0;

        asm volatile("msr daifset, #8");

        write_sysreg(0, osdlr_el1);
        write_sysreg(0, oslar_el1);
        isb();

        write_sysreg(0, mdscr_el1);
        write_sysreg(0, contextidr_el1);

        /* Reset all bcr/bvr/wcr/wvr registers */
        dfr0 = read_sysreg(id_aa64dfr0_el1);
        brps = FIELD_GET(ID_AA64DFR0_EL1_BRPs, dfr0);
        for (i = 0; i <= brps; i++) {
                write_dbgbcr(i, 0);
                write_dbgbvr(i, 0);
        }
        wrps = FIELD_GET(ID_AA64DFR0_EL1_WRPs, dfr0);
        for (i = 0; i <= wrps; i++) {
                write_dbgwcr(i, 0);
                write_dbgwvr(i, 0);
        }

        isb();
}

static void enable_os_lock(void)
{
        write_sysreg(1, oslar_el1);
        isb();

        GUEST_ASSERT(read_sysreg(oslsr_el1) & 2);
}

static void enable_monitor_debug_exceptions(void)
{
        uint64_t mdscr;

        asm volatile("msr daifclr, #8");

        mdscr = read_sysreg(mdscr_el1) | MDSCR_KDE | MDSCR_MDE;
        write_sysreg(mdscr, mdscr_el1);
        isb();
}

static void install_wp(uint8_t wpn, uint64_t addr)
{
        uint32_t wcr;

        wcr = DBGWCR_LEN8 | DBGWCR_RD | DBGWCR_WR | DBGWCR_EL1 | DBGWCR_E;
        write_dbgwcr(wpn, wcr);
        write_dbgwvr(wpn, addr);

        isb();

        enable_monitor_debug_exceptions();
}

static void install_hw_bp(uint8_t bpn, uint64_t addr)
{
        uint32_t bcr;

        bcr = DBGBCR_LEN8 | DBGBCR_EXEC | DBGBCR_EL1 | DBGBCR_E;
        write_dbgbcr(bpn, bcr);
        write_dbgbvr(bpn, addr);
        isb();

        enable_monitor_debug_exceptions();
}

static void install_wp_ctx(uint8_t addr_wp, uint8_t ctx_bp, uint64_t addr,
                           uint64_t ctx)
{
        uint32_t wcr;
        uint64_t ctx_bcr;

        /* Setup a context-aware breakpoint for Linked Context ID Match */
        ctx_bcr = DBGBCR_LEN8 | DBGBCR_EXEC | DBGBCR_EL1 | DBGBCR_E |
                  DBGBCR_BT_CTX_LINK;
        write_dbgbcr(ctx_bp, ctx_bcr);
        write_dbgbvr(ctx_bp, ctx);

        /* Setup a linked watchpoint (linked to the context-aware breakpoint) */
        wcr = DBGWCR_LEN8 | DBGWCR_RD | DBGWCR_WR | DBGWCR_EL1 | DBGWCR_E |
              DBGWCR_WT_LINK | ((uint32_t)ctx_bp << DBGWCR_LBN_SHIFT);
        write_dbgwcr(addr_wp, wcr);
        write_dbgwvr(addr_wp, addr);
        isb();

        enable_monitor_debug_exceptions();
}

void install_hw_bp_ctx(uint8_t addr_bp, uint8_t ctx_bp, uint64_t addr,
                       uint64_t ctx)
{
        uint32_t addr_bcr, ctx_bcr;

        /* Setup a context-aware breakpoint for Linked Context ID Match */
        ctx_bcr = DBGBCR_LEN8 | DBGBCR_EXEC | DBGBCR_EL1 | DBGBCR_E |
                  DBGBCR_BT_CTX_LINK;
        write_dbgbcr(ctx_bp, ctx_bcr);
        write_dbgbvr(ctx_bp, ctx);

        /*
         * Setup a normal breakpoint for Linked Address Match, and link it
         * to the context-aware breakpoint.
         */
        addr_bcr = DBGBCR_LEN8 | DBGBCR_EXEC | DBGBCR_EL1 | DBGBCR_E |
                   DBGBCR_BT_ADDR_LINK_CTX |
                   ((uint32_t)ctx_bp << DBGBCR_LBN_SHIFT);
        write_dbgbcr(addr_bp, addr_bcr);
        write_dbgbvr(addr_bp, addr);
        isb();

        enable_monitor_debug_exceptions();
}

static void install_ss(void)
{
        uint64_t mdscr;

        asm volatile("msr daifclr, #8");

        mdscr = read_sysreg(mdscr_el1) | MDSCR_KDE | MDSCR_SS;
        write_sysreg(mdscr, mdscr_el1);
        isb();
}

static volatile char write_data;

static void guest_code(uint8_t bpn, uint8_t wpn, uint8_t ctx_bpn)
{
        uint64_t ctx = 0xabcdef;        /* a random context number */

        /* Software-breakpoint */
        reset_debug_state();
        asm volatile("sw_bp: brk #0");
        GUEST_ASSERT_EQ(sw_bp_addr, PC(sw_bp));

        /* Hardware-breakpoint */
        reset_debug_state();
        install_hw_bp(bpn, PC(hw_bp));
        asm volatile("hw_bp: nop");
        GUEST_ASSERT_EQ(hw_bp_addr, PC(hw_bp));

        /* Hardware-breakpoint + svc */
        reset_debug_state();
        install_hw_bp(bpn, PC(bp_svc));
        asm volatile("bp_svc: svc #0");
        GUEST_ASSERT_EQ(hw_bp_addr, PC(bp_svc));
        GUEST_ASSERT_EQ(svc_addr, PC(bp_svc) + 4);

        /* Hardware-breakpoint + software-breakpoint */
        reset_debug_state();
        install_hw_bp(bpn, PC(bp_brk));
        asm volatile("bp_brk: brk #0");
        GUEST_ASSERT_EQ(sw_bp_addr, PC(bp_brk));
        GUEST_ASSERT_EQ(hw_bp_addr, PC(bp_brk));

        /* Watchpoint */
        reset_debug_state();
        install_wp(wpn, PC(write_data));
        write_data = 'x';
        GUEST_ASSERT_EQ(write_data, 'x');
        GUEST_ASSERT_EQ(wp_data_addr, PC(write_data));

        /* Single-step */
        reset_debug_state();
        install_ss();
        ss_idx = 0;
        asm volatile("ss_start:\n"
                     "mrs x0, esr_el1\n"
                     "add x0, x0, #1\n"
                     "msr daifset, #8\n"
                     : : : "x0");
        GUEST_ASSERT_EQ(ss_addr[0], PC(ss_start));
        GUEST_ASSERT_EQ(ss_addr[1], PC(ss_start) + 4);
        GUEST_ASSERT_EQ(ss_addr[2], PC(ss_start) + 8);

        /* OS Lock does not block software-breakpoint */
        reset_debug_state();
        enable_os_lock();
        sw_bp_addr = 0;
        asm volatile("sw_bp2: brk #0");
        GUEST_ASSERT_EQ(sw_bp_addr, PC(sw_bp2));

        /* OS Lock blocking hardware-breakpoint */
        reset_debug_state();
        enable_os_lock();
        install_hw_bp(bpn, PC(hw_bp2));
        hw_bp_addr = 0;
        asm volatile("hw_bp2: nop");
        GUEST_ASSERT_EQ(hw_bp_addr, 0);

        /* OS Lock blocking watchpoint */
        reset_debug_state();
        enable_os_lock();
        write_data = '\0';
        wp_data_addr = 0;
        install_wp(wpn, PC(write_data));
        write_data = 'x';
        GUEST_ASSERT_EQ(write_data, 'x');
        GUEST_ASSERT_EQ(wp_data_addr, 0);

        /* OS Lock blocking single-step */
        reset_debug_state();
        enable_os_lock();
        ss_addr[0] = 0;
        install_ss();
        ss_idx = 0;
        asm volatile("mrs x0, esr_el1\n\t"
                     "add x0, x0, #1\n\t"
                     "msr daifset, #8\n\t"
                     : : : "x0");
        GUEST_ASSERT_EQ(ss_addr[0], 0);

        /* Linked hardware-breakpoint */
        hw_bp_addr = 0;
        reset_debug_state();
        install_hw_bp_ctx(bpn, ctx_bpn, PC(hw_bp_ctx), ctx);
        /* Set context id */
        write_sysreg(ctx, contextidr_el1);
        isb();
        asm volatile("hw_bp_ctx: nop");
        write_sysreg(0, contextidr_el1);
        GUEST_ASSERT_EQ(hw_bp_addr, PC(hw_bp_ctx));

        /* Linked watchpoint */
        reset_debug_state();
        install_wp_ctx(wpn, ctx_bpn, PC(write_data), ctx);
        /* Set context id */
        write_sysreg(ctx, contextidr_el1);
        isb();
        write_data = 'x';
        GUEST_ASSERT_EQ(write_data, 'x');
        GUEST_ASSERT_EQ(wp_data_addr, PC(write_data));

        GUEST_DONE();
}

static void guest_sw_bp_handler(struct ex_regs *regs)
{
        sw_bp_addr = regs->pc;
        regs->pc += 4;
}

static void guest_hw_bp_handler(struct ex_regs *regs)
{
        hw_bp_addr = regs->pc;
        regs->pstate |= SPSR_D;
}

static void guest_wp_handler(struct ex_regs *regs)
{
        wp_data_addr = read_sysreg(far_el1);
        wp_addr = regs->pc;
        regs->pstate |= SPSR_D;
}

static void guest_ss_handler(struct ex_regs *regs)
{
        __GUEST_ASSERT(ss_idx < 4, "Expected index < 4, got '%lu'", ss_idx);
        ss_addr[ss_idx++] = regs->pc;
        regs->pstate |= SPSR_SS;
}

static void guest_svc_handler(struct ex_regs *regs)
{
        svc_addr = regs->pc;
}

static void guest_code_ss(int test_cnt)
{
        uint64_t i;
        uint64_t bvr, wvr, w_bvr, w_wvr;

        for (i = 0; i < test_cnt; i++) {
                /* Bits [1:0] of dbg{b,w}vr are RES0 */
                w_bvr = i << 2;
                w_wvr = i << 2;

                /*
                 * Enable Single Step execution.  Note!  This _must_ be a bare
                 * ucall as the ucall() path uses atomic operations to manage
                 * the ucall structures, and the built-in "atomics" are usually
                 * implemented via exclusive access instructions.  The exlusive
                 * monitor is cleared on ERET, and so taking debug exceptions
                 * during a LDREX=>STREX sequence will prevent forward progress
                 * and hang the guest/test.
                 */
                GUEST_UCALL_NONE();

                /*
                 * The userspace will verify that the pc is as expected during
                 * single step execution between iter_ss_begin and iter_ss_end.
                 */
                asm volatile("iter_ss_begin:nop\n");

                write_sysreg(w_bvr, dbgbvr0_el1);
                write_sysreg(w_wvr, dbgwvr0_el1);
                bvr = read_sysreg(dbgbvr0_el1);
                wvr = read_sysreg(dbgwvr0_el1);

                /* Userspace disables Single Step when the end is nigh. */
                asm volatile("iter_ss_end:\n");

                GUEST_ASSERT_EQ(bvr, w_bvr);
                GUEST_ASSERT_EQ(wvr, w_wvr);
        }
        GUEST_DONE();
}

static int debug_version(uint64_t id_aa64dfr0)
{
        return FIELD_GET(ID_AA64DFR0_EL1_DebugVer, id_aa64dfr0);
}

static void test_guest_debug_exceptions(uint8_t bpn, uint8_t wpn, uint8_t ctx_bpn)
{
        struct kvm_vcpu *vcpu;
        struct kvm_vm *vm;
        struct ucall uc;

        vm = vm_create_with_one_vcpu(&vcpu, guest_code);

        vm_init_descriptor_tables(vm);
        vcpu_init_descriptor_tables(vcpu);

        vm_install_sync_handler(vm, VECTOR_SYNC_CURRENT,
                                ESR_ELx_EC_BRK64, guest_sw_bp_handler);
        vm_install_sync_handler(vm, VECTOR_SYNC_CURRENT,
                                ESR_ELx_EC_BREAKPT_CUR, guest_hw_bp_handler);
        vm_install_sync_handler(vm, VECTOR_SYNC_CURRENT,
                                ESR_ELx_EC_WATCHPT_CUR, guest_wp_handler);
        vm_install_sync_handler(vm, VECTOR_SYNC_CURRENT,
                                ESR_ELx_EC_SOFTSTP_CUR, guest_ss_handler);
        vm_install_sync_handler(vm, VECTOR_SYNC_CURRENT,
                                ESR_ELx_EC_SVC64, guest_svc_handler);

        /* Specify bpn/wpn/ctx_bpn to be tested */
        vcpu_args_set(vcpu, 3, bpn, wpn, ctx_bpn);
        pr_debug("Use bpn#%d, wpn#%d and ctx_bpn#%d\n", bpn, wpn, ctx_bpn);

        vcpu_run(vcpu);
        switch (get_ucall(vcpu, &uc)) {
        case UCALL_ABORT:
                REPORT_GUEST_ASSERT(uc);
                break;
        case UCALL_DONE:
                goto done;
        default:
                TEST_FAIL("Unknown ucall %lu", uc.cmd);
        }

done:
        kvm_vm_free(vm);
}

void test_single_step_from_userspace(int test_cnt)
{
        struct kvm_vcpu *vcpu;
        struct kvm_vm *vm;
        struct ucall uc;
        struct kvm_run *run;
        uint64_t pc, cmd;
        uint64_t test_pc = 0;
        bool ss_enable = false;
        struct kvm_guest_debug debug = {};

        vm = vm_create_with_one_vcpu(&vcpu, guest_code_ss);
        run = vcpu->run;
        vcpu_args_set(vcpu, 1, test_cnt);

        while (1) {
                vcpu_run(vcpu);
                if (run->exit_reason != KVM_EXIT_DEBUG) {
                        cmd = get_ucall(vcpu, &uc);
                        if (cmd == UCALL_ABORT) {
                                REPORT_GUEST_ASSERT(uc);
                                /* NOT REACHED */
                        } else if (cmd == UCALL_DONE) {
                                break;
                        }

                        TEST_ASSERT(cmd == UCALL_NONE,
                                    "Unexpected ucall cmd 0x%lx", cmd);

                        debug.control = KVM_GUESTDBG_ENABLE |
                                        KVM_GUESTDBG_SINGLESTEP;
                        ss_enable = true;
                        vcpu_guest_debug_set(vcpu, &debug);
                        continue;
                }

                TEST_ASSERT(ss_enable, "Unexpected KVM_EXIT_DEBUG");

                /* Check if the current pc is expected. */
                pc = vcpu_get_reg(vcpu, ARM64_CORE_REG(regs.pc));
                TEST_ASSERT(!test_pc || pc == test_pc,
                            "Unexpected pc 0x%lx (expected 0x%lx)",
                            pc, test_pc);

                if ((pc + 4) == (uint64_t)&iter_ss_end) {
                        test_pc = 0;
                        debug.control = KVM_GUESTDBG_ENABLE;
                        ss_enable = false;
                        vcpu_guest_debug_set(vcpu, &debug);
                        continue;
                }

                /*
                 * If the current pc is between iter_ss_bgin and
                 * iter_ss_end, the pc for the next KVM_EXIT_DEBUG should
                 * be the current pc + 4.
                 */
                if ((pc >= (uint64_t)&iter_ss_begin) &&
                    (pc < (uint64_t)&iter_ss_end))
                        test_pc = pc + 4;
                else
                        test_pc = 0;
        }

        kvm_vm_free(vm);
}

/*
 * Run debug testing using the various breakpoint#, watchpoint# and
 * context-aware breakpoint# with the given ID_AA64DFR0_EL1 configuration.
 */
void test_guest_debug_exceptions_all(uint64_t aa64dfr0)
{
        uint8_t brp_num, wrp_num, ctx_brp_num, normal_brp_num, ctx_brp_base;
        int b, w, c;

        /* Number of breakpoints */
        brp_num = FIELD_GET(ID_AA64DFR0_EL1_BRPs, aa64dfr0) + 1;
        __TEST_REQUIRE(brp_num >= 2, "At least two breakpoints are required");

        /* Number of watchpoints */
        wrp_num = FIELD_GET(ID_AA64DFR0_EL1_WRPs, aa64dfr0) + 1;

        /* Number of context aware breakpoints */
        ctx_brp_num = FIELD_GET(ID_AA64DFR0_EL1_CTX_CMPs, aa64dfr0) + 1;

        pr_debug("%s brp_num:%d, wrp_num:%d, ctx_brp_num:%d\n", __func__,
                 brp_num, wrp_num, ctx_brp_num);

        /* Number of normal (non-context aware) breakpoints */
        normal_brp_num = brp_num - ctx_brp_num;

        /* Lowest context aware breakpoint number */
        ctx_brp_base = normal_brp_num;

        /* Run tests with all supported breakpoints/watchpoints */
        for (c = ctx_brp_base; c < ctx_brp_base + ctx_brp_num; c++) {
                for (b = 0; b < normal_brp_num; b++) {
                        for (w = 0; w < wrp_num; w++)
                                test_guest_debug_exceptions(b, w, c);
                }
        }
}

static void help(char *name)
{
        puts("");
        printf("Usage: %s [-h] [-i iterations of the single step test]\n", name);
        puts("");
        exit(0);
}

int main(int argc, char *argv[])
{
        struct kvm_vcpu *vcpu;
        struct kvm_vm *vm;
        int opt;
        int ss_iteration = 10000;
        uint64_t aa64dfr0;

        vm = vm_create_with_one_vcpu(&vcpu, guest_code);
        aa64dfr0 = vcpu_get_reg(vcpu, KVM_ARM64_SYS_REG(SYS_ID_AA64DFR0_EL1));
        __TEST_REQUIRE(debug_version(aa64dfr0) >= 6,
                       "Armv8 debug architecture not supported.");
        kvm_vm_free(vm);

        while ((opt = getopt(argc, argv, "i:")) != -1) {
                switch (opt) {
                case 'i':
                        ss_iteration = atoi_positive("Number of iterations", optarg);
                        break;
                case 'h':
                default:
                        help(argv[0]);
                        break;
                }
        }

        test_guest_debug_exceptions_all(aa64dfr0);
        test_single_step_from_userspace(ss_iteration);

        return 0;
}