/*
 * 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 <fcntl.h>
#include <libgen.h>
#include <sys/stat.h>
#include <errno.h>
#include <err.h>
#include <assert.h>
#include <sys/sysmacros.h>
#include <stdbool.h>

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

#include "common.h"

#define	PAGESZ		4096
#define	TEST_PAGE_COUNT	256
#define	TEST_MEM_SZ	(PAGESZ * 256)

static struct vmctx *
check_vmm_capability(const char *tname)
{
	char vmname[VM_MAX_NAMELEN];

	name_test_vm(tname, vmname);
	int res = vm_create(vmname, VCF_TRACK_DIRTY);

	if (res != 0) {
		if (errno == ENOTSUP) {
			(void) fprintf(stderr,
			    "VMM lacks dirty page tracking capability");
			(void) printf("%s\tSKIP\n", tname);
			exit(EXIT_SUCCESS);
		}
		err(EXIT_FAILURE, "could not create VM");
	}
	struct vmctx *ctx = vm_open(vmname);
	if (ctx == NULL) {
		err(EXIT_FAILURE, "could not open test VM");
	}

	return (ctx);
}

static void
expect_errno(int expected)
{
	if (errno != expected) {
		errx(EXIT_FAILURE, "unexpected errno %d != %d",
		    errno, expected);
	}
}

static uint8_t
popc8(uint8_t val)
{
	uint8_t cnt;

	for (cnt = 0; val != 0; val &= (val - 1)) {
		cnt++;
	}
	return (cnt);
}

static uint_t
legacy_clear_dirty(struct vmctx *ctx)
{
	uint8_t bitmap[TEST_PAGE_COUNT / 8] = { 0 };
	struct vmm_dirty_tracker req = {
		.vdt_start_gpa = 0,
		.vdt_len = TEST_MEM_SZ,
		.vdt_pfns = bitmap,
	};

	if (ioctl(vm_get_device_fd(ctx), VM_TRACK_DIRTY_PAGES, &req) != 0) {
		err(EXIT_FAILURE, "VM_TRACK_DIRTY_PAGES failed");
	}

	uint_t bits_set = 0;
	for (uint_t i = 0; i < (TEST_PAGE_COUNT / 8); i++) {
		bits_set += popc8(bitmap[i]);
	}
	return (bits_set);
}

static void
do_npt_op(int vmfd, struct vm_npt_operation *vno)
{
	if (ioctl(vmfd, VM_NPT_OPERATION, vno) != 0) {
		err(EXIT_FAILURE, "VM_NPT_OPERATION failed");
	}
}

static void
test_legacy(struct vmctx *ctx)
{
	const int vmfd = vm_get_device_fd(ctx);
	uint8_t *datap = vm_map_gpa(ctx, 0, PAGESZ);

	/* dirty the first page */
	*datap = 0xff;

	uint8_t bitmap[TEST_PAGE_COUNT / 8] = { 0 };
	struct vmm_dirty_tracker req = {
		.vdt_start_gpa = 0,
		.vdt_len = TEST_MEM_SZ,
		.vdt_pfns = bitmap,
	};

	if (ioctl(vmfd, VM_TRACK_DIRTY_PAGES, &req) != 0) {
		err(EXIT_FAILURE, "VM_TRACK_DIRTY_PAGES failed");
	}

	if (bitmap[0] != 1) {
		errx(EXIT_FAILURE, "first page not marked dirty");
	}
	for (uint_t i = 1; i < (TEST_PAGE_COUNT / 8); i++) {
		if (bitmap[i] != 0) {
			errx(EXIT_FAILURE,
			    "unexpected non-zero entry: bitmap[%u] = %x\n",
			    i, bitmap[i]);
		}
	}
}

static void
test_toggle_tracking(struct vmctx *ctx)
{
	const int vmfd = vm_get_device_fd(ctx);
	struct vm_npt_operation vno = {
		.vno_operation = VNO_OP_GET_TRACK_DIRTY,
		.vno_gpa = 0,
		.vno_len = 0,
	};

	/*
	 * Since the VM was created with VCF_TRACK_DIRTY set, dirty tracking
	 * should already be active.
	 */
	if (ioctl(vmfd, VM_NPT_OPERATION, &vno) != 1) {
		errx(EXIT_FAILURE, "expected dirty tracking to be active");
	}

	vno.vno_operation = VNO_OP_DIS_TRACK_DIRTY;
	if (ioctl(vmfd, VM_NPT_OPERATION, &vno) != 0) {
		err(EXIT_FAILURE, "VM_NPT_OPERATION failed");
	}

	vno.vno_operation = VNO_OP_GET_TRACK_DIRTY;
	if (ioctl(vmfd, VM_NPT_OPERATION, &vno) != 0) {
		errx(EXIT_FAILURE, "expected dirty tracking to be inactive");
	}

	vno.vno_operation = VNO_OP_EN_TRACK_DIRTY;
	if (ioctl(vmfd, VM_NPT_OPERATION, &vno) != 0) {
		err(EXIT_FAILURE, "VM_NPT_OPERATION failed");
	}

	vno.vno_operation = VNO_OP_GET_TRACK_DIRTY;
	if (ioctl(vmfd, VM_NPT_OPERATION, &vno) != 1) {
		errx(EXIT_FAILURE,
		    "expected dirty tracking to be active again");
	}
}

static void
test_inval_args(struct vmctx *ctx)
{
	const int vmfd = vm_get_device_fd(ctx);
	struct vm_npt_operation vno = { 0 };

	/* invalid vno_operation */
	vno.vno_operation = ~0;
	if (ioctl(vmfd, VM_NPT_OPERATION, &vno) == 0) {
		err(EXIT_FAILURE, "unexpected VM_NPT_OPERATION success");
	}
	expect_errno(EINVAL);

	/* valid operation, but gpa which is not page-aligned */
	vno.vno_operation = VNO_OP_GET_DIRTY | VNO_FLAG_BITMAP_IN;
	vno.vno_gpa = 0x100;
	vno.vno_len = PAGESZ;

	if (ioctl(vmfd, VM_NPT_OPERATION, &vno) == 0) {
		err(EXIT_FAILURE, "unexpected VM_NPT_OPERATION success");
	}
	expect_errno(EINVAL);

	/* gpa is page-aligned, but len isn't */
	vno.vno_gpa = 0;
	vno.vno_len = PAGESZ + 0x100;

	if (ioctl(vmfd, VM_NPT_OPERATION, &vno) == 0) {
		err(EXIT_FAILURE, "unexpected VM_NPT_OPERATION success");
	}
	expect_errno(EINVAL);

	/* overflowing region */
	vno.vno_gpa = 0xffffffffffffe000;
	vno.vno_len = 512 * PAGESZ;

	if (ioctl(vmfd, VM_NPT_OPERATION, &vno) == 0) {
		err(EXIT_FAILURE, "unexpected VM_NPT_OPERATION success");
	}
	expect_errno(EOVERFLOW);
}

static void
test_op_get_dirty(struct vmctx *ctx)
{
	const int vmfd = vm_get_device_fd(ctx);
	uint8_t *datap = vm_map_gpa(ctx, 0, TEST_MEM_SZ);

	/* Use legacy mechanism to ensure dirty bits are clear to start */
	(void) legacy_clear_dirty(ctx);

	/* Dirty the first page out of every 8 */
	for (uint_t i = 0; i < TEST_MEM_SZ; i += (PAGESZ * 8)) {
		datap[i] = 0xff;
	}

	uint8_t bits[TEST_PAGE_COUNT / 8] = { 0 };
	struct vm_npt_operation vno = {
		.vno_gpa = 0,
		.vno_len = TEST_MEM_SZ,
		.vno_operation = VNO_OP_GET_DIRTY | VNO_FLAG_BITMAP_OUT,
		.vno_bitmap = bits,
	};
	do_npt_op(vmfd, &vno);

	for (uint_t i = 0; i < TEST_PAGE_COUNT / 8; i++) {
		if (bits[i] != 0x01) {
			errx(EXIT_FAILURE,
			    "unexpected dirty bits %02x at base gpa %08x",
			    bits[i], i * PAGESZ * 8);
		}
	}

	/* Clear those bits again */
	(void) legacy_clear_dirty(ctx);

	/* And check that they are zeroed now */
	do_npt_op(vmfd, &vno);
	for (uint_t i = 0; i < TEST_PAGE_COUNT / 8; i++) {
		if (bits[i] != 0) {
			errx(EXIT_FAILURE,
			    "unexpected dirty bits %02x at base gpa %08x",
			    bits[i], i * PAGESZ * 8);
		}
	}
}

static void
test_op_set_dirty(struct vmctx *ctx)
{
	const int vmfd = vm_get_device_fd(ctx);

	/* Use legacy mechanism to ensure dirty bits are clear to start */
	(void) legacy_clear_dirty(ctx);

	/* Mark first 17 pages as dirty */
	uint8_t bits[TEST_PAGE_COUNT / 8] = { 0xff, 0xff, 0x80 };
	struct vm_npt_operation vno = {
		.vno_gpa = 0,
		.vno_len = TEST_MEM_SZ,
		.vno_operation = VNO_OP_SET_DIRTY | VNO_FLAG_BITMAP_IN,
		.vno_bitmap = bits,
	};
	do_npt_op(vmfd, &vno);

	uint_t legacy_dirty = legacy_clear_dirty(ctx);
	if (legacy_dirty != 17) {
		errx(EXIT_FAILURE, "unexpected dirty count after OP_SET_DIRTY");
	}
}

#define	BMAP_IDX(gpa)	((gpa) / (PAGESZ * 8))
#define	BMAP_BIT(gpa)	(((gpa) / PAGESZ) % 8)

static void
test_op_reset_dirty(struct vmctx *ctx)
{
	const int vmfd = vm_get_device_fd(ctx);
	uint8_t *datap = vm_map_gpa(ctx, 0, TEST_MEM_SZ);

	/* Use legacy mechanism to ensure dirty bits are clear to start */
	(void) legacy_clear_dirty(ctx);

	/* Dirty the front half of memory */
	for (uintptr_t gpa = 0; gpa < (TEST_MEM_SZ / 2); gpa += PAGESZ) {
		datap[gpa] = 0xff;
	}

	uint8_t bits[TEST_PAGE_COUNT / 8] = { 0 };
	/* Mark bitmap for every other page, starting at 0 */
	for (uintptr_t gpa = 0; gpa < TEST_MEM_SZ; gpa += (2 * PAGESZ)) {
		bits[BMAP_IDX(gpa)] |= (1 << BMAP_BIT(gpa));
	}

	struct vm_npt_operation vno = {
		.vno_gpa = 0,
		.vno_len = TEST_MEM_SZ,
		.vno_operation = VNO_OP_RESET_DIRTY |
		    VNO_FLAG_BITMAP_IN | VNO_FLAG_BITMAP_OUT,
		.vno_bitmap = bits,
	};
	do_npt_op(vmfd, &vno);

	/* Check that pages marked dirty were reported back as such */
	for (uintptr_t gpa = 0; gpa < TEST_MEM_SZ; gpa += PAGESZ) {
		const bool is_even_page = (BMAP_BIT(gpa) % 2) == 0;
		const bool is_dirty =
		    (bits[BMAP_IDX(gpa)] & (1 << BMAP_BIT(gpa))) != 0;

		/* Even pages in the first half should be set */
		if (is_even_page && gpa < (TEST_MEM_SZ / 2) && !is_dirty) {
			errx(EXIT_FAILURE,
			    "missing dirty bit set at gpa %08lx", gpa);
		}

		/* Odd pages and even pages in second half should be unset */
		if (is_dirty && (!is_even_page || gpa >= (TEST_MEM_SZ / 2))) {
			errx(EXIT_FAILURE,
			    "unexpected dirty bit set at gpa %08lx", gpa);
		}
	}

	/*
	 * With half of the pages dirtied at first, and then half of those reset
	 * from dirty in the NPT operation, we expect 1/4 to be remaining.
	 */
	uint_t remaining_dirty = legacy_clear_dirty(ctx);
	if (remaining_dirty != (TEST_PAGE_COUNT / 4)) {
		errx(EXIT_FAILURE,
		    "expected %u pages remaining dirty, found %u",
		    TEST_PAGE_COUNT / 2, remaining_dirty);
	}
}

int
main(int argc, char *argv[])
{
	const char *suite_name = basename(argv[0]);
	struct vmctx *ctx;

	ctx = check_vmm_capability(suite_name);

	if (vm_setup_memory(ctx, TEST_MEM_SZ, VM_MMAP_ALL) != 0) {
		err(EXIT_FAILURE, "could not setup VM memory");
	}

	/* Test "legacy" VM_TRACK_DIRTY_PAGES mechanism first */
	test_legacy(ctx);

	/* Confirm that dirty tracking can be queried and toggled on/off */
	test_toggle_tracking(ctx);

	/* Check some invalid argument conditions */
	test_inval_args(ctx);

	/* Can dirty bits be queried with VNO_OP_GET_DIRTY */
	test_op_get_dirty(ctx);

	/* Can dirty bits be set with VNO_OP_SET_DIRTY */
	test_op_set_dirty(ctx);

	/*
	 * Can dirty bits be reset (simultaneously queried and cleared )
	 * with VNO_OP_RESET_DIRTY
	 */
	test_op_reset_dirty(ctx);

	vm_destroy(ctx);
	(void) printf("%s\tPASS\n", suite_name);
	return (EXIT_SUCCESS);
}