/*
 * This file and its contents are supplied under the terms of the
 * Common Development and Distribution License ("CDDL"), version 1.0.
 * You may only use this file in accordance with the terms of version
 * 1.0 of the CDDL.
 *
 * A full copy of the text of the CDDL should have accompanied this
 * source.  A copy of the CDDL is also available via the Internet at
 * http://www.illumos.org/license/CDDL.
 */

/*
 * Copyright 2024 Oxide Computer Company
 */

#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <strings.h>
#include <assert.h>
#include <errno.h>

#include <sys/types.h>
#include <sys/segments.h>
#include <sys/psw.h>
#include <sys/controlregs.h>
#include <sys/sysmacros.h>
#include <sys/varargs.h>
#include <sys/debug.h>
#include <sys/mman.h>

#include <sys/vmm.h>
#include <sys/vmm_dev.h>
#include <vmmapi.h>

#include "in_guest.h"


#define	PT_VALID	0x01
#define	PT_WRITABLE	0x02
#define	PT_WRITETHRU	0x08
#define	PT_NOCACHE	0x10
#define	PT_PAGESIZE	0x80

#define	SEG_ACCESS_TYPE_MASK	0x1f
#define	SEG_ACCESS_DPL_MASK	0x60
#define	SEG_ACCESS_P		(1 << 7)
#define	SEG_ACCESS_AVL		(1 << 12)
#define	SEG_ACCESS_L		(1 << 13)
#define	SEG_ACCESS_D		(1 << 14)
#define	SEG_ACCESS_G		(1 << 15)
#define	SEG_ACCESS_UNUSABLE	(1 << 16)


/*
 * Keep the test name and VM context around so the consumer is not required to
 * pass either of them to us for subsequent test-related operations after the
 * initialization has been performed.
 *
 * The test code is not designed to be reentrant at this point.
 */
static struct vmctx *test_vmctx = NULL;
static const char *test_name = NULL;

static uint64_t test_msg_addr = 0;

static int
setup_rom(struct vmctx *ctx)
{
	const size_t seg_sz = 0x1000;
	const uintptr_t seg_addr = MEM_LOC_ROM;
	const int fd = vm_get_device_fd(ctx);
	int err;

	struct vm_memseg memseg = {
		.segid = VM_BOOTROM,
		.len = 0x1000,
	};
	(void) strlcpy(memseg.name, "testrom", sizeof (memseg.name));
	err = ioctl(fd, VM_ALLOC_MEMSEG, &memseg);
	if (err != 0) {
		return (err);
	}
	err = vm_mmap_memseg(ctx, seg_addr, VM_BOOTROM, 0, seg_sz,
	    PROT_READ | PROT_EXEC);
	return (err);
}

static void
populate_identity_table(struct vmctx *ctx)
{
	uint64_t gpa, pte_loc;

	/* Set up 2MiB PTEs for everything up through 0xffffffff */
	for (gpa = 0, pte_loc = MEM_LOC_PAGE_TABLE_2M;
	    gpa < 0x100000000;
	    pte_loc += PAGE_SIZE) {
		uint64_t *ptep = vm_map_gpa(ctx, pte_loc, PAGE_SIZE);

		for (uint_t i = 0; i < 512; i++, ptep++, gpa += 0x200000) {
			*ptep =  gpa | PT_VALID | PT_WRITABLE | PT_PAGESIZE;
			/* Make traditional MMIO space uncachable */
			if (gpa >= 0xc0000000) {
				*ptep |= PT_WRITETHRU | PT_NOCACHE;
			}
		}
	}
	assert(gpa == 0x100000000 && pte_loc == MEM_LOC_PAGE_TABLE_1G);

	uint64_t *pdep = vm_map_gpa(ctx, MEM_LOC_PAGE_TABLE_1G, PAGE_SIZE);
	pdep[0] = MEM_LOC_PAGE_TABLE_2M | PT_VALID | PT_WRITABLE;
	pdep[1] = (MEM_LOC_PAGE_TABLE_2M + PAGE_SIZE) | PT_VALID | PT_WRITABLE;
	pdep[2] =
	    (MEM_LOC_PAGE_TABLE_2M + 2 * PAGE_SIZE) | PT_VALID | PT_WRITABLE;
	pdep[3] =
	    (MEM_LOC_PAGE_TABLE_2M + 3 * PAGE_SIZE) | PT_VALID | PT_WRITABLE;

	pdep = vm_map_gpa(ctx, MEM_LOC_PAGE_TABLE_512G, PAGE_SIZE);
	pdep[0] = MEM_LOC_PAGE_TABLE_1G | PT_VALID | PT_WRITABLE;
}

static void
populate_desc_tables(struct vmctx *ctx)
{

}

void
test_cleanup(bool is_failure)
{
	if (test_vmctx != NULL) {
		bool keep_on_fail = false;

		const char *keep_var;
		if ((keep_var = getenv("KEEP_ON_FAIL")) != NULL) {
			if (strlen(keep_var) != 0 &&
			    strcmp(keep_var, "0") != 0) {
				keep_on_fail = true;
			}
		}

		/*
		 * Destroy the instance unless the test failed and it was
		 * requested that we keep it around.
		 */
		if (!is_failure || !keep_on_fail) {
			vm_destroy(test_vmctx);
		}
		test_name = NULL;
		test_vmctx = NULL;
	}
}

static void fail_finish(void)
{
	assert(test_name != NULL);
	(void) printf("FAIL %s\n", test_name);

	test_cleanup(true);
	exit(EXIT_FAILURE);
}

void
test_fail(void)
{
	fail_finish();
}

void
test_fail_errno(int err, const char *msg)
{
	const char *err_str = strerror(err);

	(void) fprintf(stderr, "%s: %s\n", msg, err_str);
	fail_finish();
}

void
test_fail_msg(const char *fmt, ...)
{
	va_list ap;

	va_start(ap, fmt);
	(void) vfprintf(stderr, fmt, ap);

	fail_finish();
}

void
test_fail_vmexit(const struct vm_exit *vexit)
{
	const char *hdr_fmt = "Unexpected %s exit:\n\t%%rip: %lx\n";

	switch (vexit->exitcode) {
	case VM_EXITCODE_INOUT:
		(void) fprintf(stderr, hdr_fmt, "IN/OUT", vexit->rip);
		(void) fprintf(stderr,
		    "\teax: %08x\n"
		    "\tport: %04x\n"
		    "\tbytes: %u\n"
		    "\tflags: %x\n",
		    vexit->u.inout.eax,
		    vexit->u.inout.port,
		    vexit->u.inout.bytes,
		    vexit->u.inout.flags);
		break;
	case VM_EXITCODE_RDMSR:
		(void) fprintf(stderr, hdr_fmt, "RDMSR", vexit->rip);
		(void) fprintf(stderr, "\tcode: %08x\n", vexit->u.msr.code);
		break;
	case VM_EXITCODE_WRMSR:
		(void) fprintf(stderr, hdr_fmt, "WRMSR", vexit->rip);
		(void) fprintf(stderr,
		    "\tcode: %08x\n"
		    "\twval: %016lx\n",
		    vexit->u.msr.code, vexit->u.msr.wval);
		break;
	case VM_EXITCODE_MMIO:
		(void) fprintf(stderr, hdr_fmt, "MMIO", vexit->rip);
		(void) fprintf(stderr,
		    "\tbytes: %u\n"
		    "\ttype: %s\n"
		    "\tgpa: %x\n"
		    "\tdata: %016x\n",
		    vexit->u.mmio.bytes,
		    vexit->u.mmio.read == 0 ? "write" : "read",
		    vexit->u.mmio.gpa,
		    vexit->u.mmio.data);
		break;
	case VM_EXITCODE_VMX:
		(void) fprintf(stderr, hdr_fmt, "VMX", vexit->rip);
		(void) fprintf(stderr,
		    "\tstatus: %x\n"
		    "\treason: %x\n"
		    "\tqualification: %lx\n"
		    "\tinst_type: %x\n"
		    "\tinst_error: %x\n",
		    vexit->u.vmx.status,
		    vexit->u.vmx.exit_reason,
		    vexit->u.vmx.exit_qualification,
		    vexit->u.vmx.inst_type,
		    vexit->u.vmx.inst_error);
		break;
	case VM_EXITCODE_SVM:
		(void) fprintf(stderr, hdr_fmt, "SVM", vexit->rip);
		break;
	case VM_EXITCODE_INST_EMUL:
		(void) fprintf(stderr, hdr_fmt, "instruction emulation",
		    vexit->rip);
		const uint_t len = vexit->u.inst_emul.num_valid > 0 ?
		    vexit->u.inst_emul.num_valid : 15;
		(void) fprintf(stderr, "\tinstruction bytes: [");
		for (uint_t i = 0; i < len; i++) {
			(void) fprintf(stderr, "%s%02x",
			    i == 0 ? "" : ", ",
			    vexit->u.inst_emul.inst[i]);
		}
		(void) fprintf(stderr, "]\n");
		break;
	case VM_EXITCODE_SUSPENDED:
		(void) fprintf(stderr, hdr_fmt, "suspend", vexit->rip);
		switch (vexit->u.suspended.how) {
		case VM_SUSPEND_RESET:
			(void) fprintf(stderr, "\thow: reset");
			break;
		case VM_SUSPEND_POWEROFF:
			(void) fprintf(stderr, "\thow: poweroff");
			break;
		case VM_SUSPEND_HALT:
			(void) fprintf(stderr, "\thow: halt");
			break;
		case VM_SUSPEND_TRIPLEFAULT:
			(void) fprintf(stderr, "\thow: triple-fault");
			break;
		default:
			(void) fprintf(stderr, "\thow: unknown - %d",
			    vexit->u.suspended.how);
			break;
		}
		break;
	default:
		(void) fprintf(stderr, "Unexpected code %d exit:\n"
		    "\t%%rip: %lx\n", vexit->exitcode, vexit->rip);
		break;
	}
	fail_finish();
}

void
test_pass(void)
{
	assert(test_name != NULL);
	(void) printf("PASS %s\n", test_name);
	test_cleanup(false);
	exit(EXIT_SUCCESS);
}

const char *
test_msg_get(struct vmctx *ctx)
{
	/* Disregard if the message address is still NULL */
	const uint64_t msg_addr = test_msg_addr;
	if (msg_addr == 0) {
		return (NULL);
	}

	/*
	 * We want to try to map up to one page after the specified message
	 * address, keeping in mind the end of lowmem. (The payload, and
	 * thus message, is assumed to be in lowmem at this time.)
	 */
	const uint64_t lowmem_end = vm_get_lowmem_size(ctx);
	const uint64_t msg_map_end = MIN(msg_addr + PAGE_SIZE, lowmem_end);

	if (msg_map_end >= lowmem_end || msg_map_end <= msg_addr) {
		return (NULL);
	}
	const uint64_t max_msg_len = msg_map_end - msg_addr;

	/*
	 * Get the mapping to that guest memory.  This assumes that the payload
	 * has provided a guest-physical address to us.
	 */
	const char *result = vm_map_gpa(ctx, msg_addr, max_msg_len);
	if (result == NULL) {
		return (NULL);
	}

	/* Demand a NUL-terminated string shorter than the map limit */
	if (strnlen(result, max_msg_len) >= max_msg_len) {
		return (NULL);
	}

	return (result);
}

void
test_msg_print(struct vmctx *ctx)
{
	const char *payload_msg = test_msg_get(ctx);

	if (payload_msg != NULL) {
		(void) fprintf(stderr, "MSG: %s\n", payload_msg);
	}
}

static int
load_payload(struct vmctx *ctx)
{
	extern uint8_t payload_data;
	extern uint32_t payload_size;

	const uint32_t len = payload_size;
	const uint32_t cap = (MEM_TOTAL_SZ - MEM_LOC_PAYLOAD);

	if (len > cap) {
		test_fail_msg("Payload size %u > capacity %u\n", len, cap);
	}

	const size_t map_len = P2ROUNDUP(len, PAGE_SIZE);
	void *outp = vm_map_gpa(ctx, MEM_LOC_PAYLOAD, map_len);
	bcopy(&payload_data, outp, len);

	return (0);
}

static struct vmctx *
test_initialize_opts(const char *tname, uint64_t create_flags, bool is_plain)
{
	char vm_name[VM_MAX_NAMELEN];
	int err;
	struct vmctx *ctx;

	assert(test_vmctx == NULL);
	assert(test_name == NULL);

	test_name = strdup(tname);
	(void) snprintf(vm_name, sizeof (vm_name), "bhyve-test-%s-%d",
	    test_name, getpid());

	err = vm_create(vm_name, create_flags);
	if (err != 0) {
		test_fail_errno(err, "Could not create VM");
	}

	ctx = vm_open(vm_name);
	if (ctx == NULL) {
		test_fail_errno(errno, "Could not open VM");
	}
	test_vmctx = ctx;

	/* No further setup required for a "plain" instance */
	if (is_plain) {
		return (ctx);
	}

	err = vm_setup_memory(ctx, MEM_TOTAL_SZ, VM_MMAP_ALL);
	if (err != 0) {
		test_fail_errno(err, "Could not set up VM memory");
	}

	err = setup_rom(ctx);
	if (err != 0) {
		test_fail_errno(err, "Could not set up VM ROM segment");
	}

	populate_identity_table(ctx);
	populate_desc_tables(ctx);

	err = load_payload(ctx);
	if (err != 0) {
		test_fail_errno(err, "Could not load payload");
	}

	return (ctx);
}

struct vmctx *
test_initialize(const char *tname)
{
	return (test_initialize_opts(tname, 0, false));
}

struct vmctx *
test_initialize_plain(const char *tname)
{
	return (test_initialize_opts(tname, 0, true));
}

struct vmctx *
test_initialize_flags(const char *tname, uint64_t create_flags)
{
	return (test_initialize_opts(tname, create_flags, false));
}

void
test_reinitialize(struct vmctx *ctx, uint64_t flags)
{
	int err;

	if ((err = vm_reinit(ctx, flags)) != 0) {
		test_fail_errno(err, "Could not reinit VM");
	}

	/* Reload tables and payload in case they were altered */

	populate_identity_table(ctx);
	populate_desc_tables(ctx);

	err = load_payload(ctx);
	if (err != 0) {
		test_fail_errno(err, "Could not load payload");
	}
}

int
test_setup_vcpu(struct vcpu *vcpu, uint64_t rip, uint64_t rsp)
{
	int err;

	err = vm_activate_cpu(vcpu);
	if (err != 0 && err != EBUSY) {
		return (err);
	}

	/*
	 * Granularity bit important here for VMX validity:
	 * "If any bit in the limit field in the range 31:20 is 1, G must be 1"
	 */
	err = vm_set_desc(vcpu, VM_REG_GUEST_CS, 0, UINT32_MAX,
	    SDT_MEMERA | SEG_ACCESS_P | SEG_ACCESS_L | SEG_ACCESS_G);
	if (err != 0) {
		return (err);
	}

	err = vm_set_desc(vcpu, VM_REG_GUEST_SS, 0, UINT32_MAX,
	    SDT_MEMRWA | SEG_ACCESS_P | SEG_ACCESS_L |
	    SEG_ACCESS_D | SEG_ACCESS_G);
	if (err != 0) {
		return (err);
	}

	err = vm_set_desc(vcpu, VM_REG_GUEST_DS, 0, UINT32_MAX,
	    SDT_MEMRWA | SEG_ACCESS_P | SEG_ACCESS_D | SEG_ACCESS_G);
	if (err != 0) {
		return (err);
	}

	/*
	 * While SVM will happilly run with an otherwise unusable TR, VMX
	 * includes it among its entry checks.
	 */
	err = vm_set_desc(vcpu, VM_REG_GUEST_TR, MEM_LOC_TSS, 0xff,
	    SDT_SYSTSSBSY | SEG_ACCESS_P);
	if (err != 0) {
		return (err);
	}
	err = vm_set_desc(vcpu, VM_REG_GUEST_GDTR, MEM_LOC_GDT, 0x1ff, 0);
	if (err != 0) {
		return (err);
	}
	err = vm_set_desc(vcpu, VM_REG_GUEST_IDTR, MEM_LOC_IDT, 0xfff, 0);
	if (err != 0) {
		return (err);
	}

	/* Mark unused segments as explicitly unusable (for VMX) */
	const int unsable_segs[] = {
		VM_REG_GUEST_ES,
		VM_REG_GUEST_FS,
		VM_REG_GUEST_GS,
		VM_REG_GUEST_LDTR,
	};
	for (uint_t i = 0; i < ARRAY_SIZE(unsable_segs); i++) {
		err = vm_set_desc(vcpu, unsable_segs[i], 0, 0,
		    SEG_ACCESS_UNUSABLE);
		if (err != 0) {
			return (err);
		}
	}

	/* Place CPU directly in long mode */
	const int regnums[] = {
		VM_REG_GUEST_CR0,
		VM_REG_GUEST_CR3,
		VM_REG_GUEST_CR4,
		VM_REG_GUEST_EFER,
		VM_REG_GUEST_RFLAGS,
		VM_REG_GUEST_RIP,
		VM_REG_GUEST_RSP,
		VM_REG_GUEST_CS,
		VM_REG_GUEST_SS,
		VM_REG_GUEST_DS,
		VM_REG_GUEST_TR,
	};
	uint64_t regvals[] = {
		CR0_PG | CR0_AM | CR0_WP | CR0_NE | CR0_ET | CR0_TS |
		    CR0_MP | CR0_PE,
		MEM_LOC_PAGE_TABLE_512G,
		CR4_DE | CR4_PSE | CR4_PAE | CR4_MCE | CR4_PGE | CR4_FSGSBASE,
		AMD_EFER_SCE | AMD_EFER_LME | AMD_EFER_LMA | AMD_EFER_NXE,
		/* start with interrupts disabled */
		PS_MB1,
		rip,
		rsp,
		(GDT_KCODE << 3),
		(GDT_KDATA << 3),
		(GDT_KDATA << 3),
		(GDT_KTSS << 3),
	};
	assert(ARRAY_SIZE(regnums) == ARRAY_SIZE(regvals));

	err = vm_set_register_set(vcpu, ARRAY_SIZE(regnums), regnums,
	    regvals);
	if (err != 0) {
		return (err);
	}

	err = vm_set_run_state(vcpu, VRS_RUN, 0);
	if (err != 0) {
		return (err);
	}

	return (0);
}

static enum vm_exit_kind
which_exit_kind(struct vm_entry *ventry, const struct vm_exit *vexit)
{
	const struct vm_inout *inout = &vexit->u.inout;

	switch (vexit->exitcode) {
	case VM_EXITCODE_BOGUS:
		bzero(ventry, sizeof (ventry));
		return (VEK_REENTR);
	case VM_EXITCODE_INOUT:
		if (inout->port == IOP_TEST_RESULT &&
		    (inout->flags & INOUT_IN) == 0) {
			if (inout->eax == TEST_RESULT_PASS) {
				return (VEK_TEST_PASS);
			} else {
				return (VEK_TEST_FAIL);
			}
		}
		if (inout->port == IOP_TEST_MSG &&
		    (inout->flags & INOUT_IN) == 0 &&
		    inout->bytes == 4) {
			test_msg_addr = inout->eax;
			ventry_fulfill_inout(vexit, ventry, 0);
			return (VEK_TEST_MSG);
		}
		break;
	default:
		break;
	}
	return (VEK_UNHANDLED);
}

enum vm_exit_kind
test_run_vcpu(struct vcpu *vcpu, struct vm_entry *ventry, struct vm_exit *vexit)
{
	int err;

	err = vm_run(vcpu, ventry, vexit);
	if (err != 0) {
		test_fail_errno(err, "Failure during vcpu entry");
	}

	return (which_exit_kind(ventry, vexit));
}

void
ventry_fulfill_inout(const struct vm_exit *vexit, struct vm_entry *ventry,
    uint32_t data)
{
	VERIFY3U(vexit->exitcode, ==, VM_EXITCODE_INOUT);

	ventry->cmd = VEC_FULFILL_INOUT;
	bcopy(&vexit->u.inout, &ventry->u.inout, sizeof (struct vm_inout));
	if ((ventry->u.inout.flags & INOUT_IN) != 0) {
		ventry->u.inout.eax = data;
	}
}

void
ventry_fulfill_mmio(const struct vm_exit *vexit, struct vm_entry *ventry,
    uint64_t data)
{
	VERIFY3U(vexit->exitcode, ==, VM_EXITCODE_MMIO);

	ventry->cmd = VEC_FULFILL_MMIO;
	bcopy(&vexit->u.mmio, &ventry->u.mmio, sizeof (struct vm_mmio));
	if (ventry->u.mmio.read != 0) {
		ventry->u.mmio.data = data;
	}
}

bool
vexit_match_inout(const struct vm_exit *vexit, bool is_read, uint16_t port,
    uint_t len, uint32_t *valp)
{
	if (vexit->exitcode != VM_EXITCODE_INOUT) {
		return (false);
	}

	const uint_t flag = is_read ? INOUT_IN : 0;
	if (vexit->u.inout.port != port ||
	    vexit->u.inout.bytes != len ||
	    (vexit->u.inout.flags & INOUT_IN) != flag) {
		return (false);
	}

	if (!is_read && valp != NULL) {
		*valp = vexit->u.inout.eax;
	}
	return (true);
}

bool
vexit_match_mmio(const struct vm_exit *vexit, bool is_read, uint64_t addr,
    uint_t len, uint64_t *valp)
{
	if (vexit->exitcode != VM_EXITCODE_MMIO) {
		return (false);
	}

	if (vexit->u.mmio.gpa != addr ||
	    vexit->u.mmio.bytes != len ||
	    (vexit->u.mmio.read != 0) != is_read) {
		return (false);
	}

	if (!is_read && valp != NULL) {
		*valp = vexit->u.mmio.data;
	}
	return (true);
}