/*
 * 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 2022 Oxide Computer Company
 */


#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <stropts.h>
#include <strings.h>
#include <signal.h>
#include <setjmp.h>
#include <libgen.h>
#include <sys/debug.h>
#include <sys/fp.h>

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

#include "common.h"

/* Minimal xsave state area (sans any AVX storage) */
struct xsave_min {
	struct fxsave_state	legacy;
	struct xsave_header	header;
};

CTASSERT(sizeof (struct xsave_min) == MIN_XSAVE_SIZE);

struct avx_state {
	/* 16 x 128-bit: high portions of the ymm registers */
	uint64_t	ymm[32];
};

static bool
get_fpu(int fd, struct vm_fpu_state *req)
{
	int res = ioctl(fd, VM_GET_FPU, req);
	if (res != 0) {
		perror("could not read FPU for vCPU");
		return (false);
	}
	return (true);
}

static bool
set_fpu(int fd, struct vm_fpu_state *req)
{
	int res = ioctl(fd, VM_SET_FPU, req);
	if (res != 0) {
		perror("could not write FPU for vCPU");
		return (false);
	}
	return (true);
}

static bool
check_sse(int fd, const struct vm_fpu_desc *desc, void *fpu_area,
    size_t fpu_size)
{
	/* Make sure the x87/MMX/SSE state is described as present */
	bool found_fp = false, found_sse = false;
	for (uint_t i = 0; i < desc->vfd_num_entries; i++) {
		const struct vm_fpu_desc_entry *ent = &desc->vfd_entry_data[i];

		switch (ent->vfde_feature) {
		case XFEATURE_LEGACY_FP:
			found_fp = true;
			if (ent->vfde_off != 0 ||
			    ent->vfde_size != sizeof (struct fxsave_state)) {
				(void) fprintf(stderr,
				    "unexpected entity for %x: "
				    "size=%x off=%x\n", ent->vfde_feature,
				    ent->vfde_size, ent->vfde_off);
				return (false);
			}
			break;
		case XFEATURE_SSE:
			found_sse = true;
			if (ent->vfde_off != 0 ||
			    ent->vfde_size != sizeof (struct fxsave_state)) {
				(void) fprintf(stderr,
				    "unexpected entity for %x: "
				    "size=%x off=%x\n", ent->vfde_feature,
				    ent->vfde_size, ent->vfde_off);
				return (false);
			}
			break;
		}
	}

	if (!found_fp || !found_sse) {
		(void) fprintf(stderr, "did not find x87 and SSE area "
		    "descriptors as expected in initial FPU\n");
		return (false);
	}

	struct vm_fpu_state req = {
		.vcpuid = 0,
		.buf = fpu_area,
		.len = fpu_size,
	};

	if (!get_fpu(fd, &req)) {
		return (false);
	}

	struct xsave_min *xs = fpu_area;
	/*
	 * Executing this test on a freshly-created instance, we expect the FPU
	 * to only have the legacy and SSE features present in its active state.
	 */
	if (xs->header.xsh_xstate_bv != (XFEATURE_LEGACY_FP | XFEATURE_SSE)) {
		(void) fprintf(stderr, "bad xstate_bv %lx, expected %lx",
		    xs->header.xsh_xstate_bv,
		    (XFEATURE_LEGACY_FP | XFEATURE_SSE));
		return (false);
	}

	/* load some SSE values to check for a get/set cycle */
	uint64_t *xmm = (void *)&xs->legacy.fx_xmm[0];
	xmm[0] = UINT64_MAX;
	xmm[2] = 1;

	if (!set_fpu(fd, &req)) {
		return (false);
	}

	/* check that those values made it in/out of the guest FPU */
	bzero(fpu_area, fpu_size);
	if (!get_fpu(fd, &req)) {
		return (false);
	}
	if (xmm[0] != UINT64_MAX || xmm[2] != 1) {
		(void) fprintf(stderr, "SSE test registers not saved\n");
		return (false);
	}

	/* Make sure that a bogus MXCSR value is rejected */
	xs->legacy.fx_mxcsr = UINT32_MAX;
	int res = ioctl(fd, VM_SET_FPU, &req);
	if (res == 0) {
		(void) fprintf(stderr,
		    "write of invalid MXCSR erroneously allowed\n");
		return (false);
	}

	return (true);
}

static bool
check_avx(int fd, const struct vm_fpu_desc *desc, void *fpu_area,
    size_t fpu_size)
{
	bool found_avx = false;
	size_t avx_size, avx_off;
	for (uint_t i = 0; i < desc->vfd_num_entries; i++) {
		const struct vm_fpu_desc_entry *ent = &desc->vfd_entry_data[i];

		if (ent->vfde_feature == XFEATURE_AVX) {
			found_avx = true;
			avx_size = ent->vfde_size;
			avx_off = ent->vfde_off;
			break;
		}
	}

	if (!found_avx) {
		(void) printf("AVX capability not found on host CPU, "
		    "skipping related tests\n");
		return (true);
	}

	if (avx_size != sizeof (struct avx_state)) {
		(void) fprintf(stderr, "unexpected AVX state size: %x, "
		    "expected %x\n", avx_size, sizeof (struct avx_state));
		return (false);
	}
	if ((avx_off + avx_size) > fpu_size) {
		(void) fprintf(stderr, "AVX data falls outside fpu size: "
		    "%x > %x\n", avx_off + avx_size, fpu_size);
		return (false);
	}

	struct xsave_min *xs = fpu_area;
	struct avx_state *avx = fpu_area + avx_off;

	/* do a simple data round-trip */
	struct vm_fpu_state req = {
		.vcpuid = 0,
		.buf = fpu_area,
		.len = fpu_size,
	};
	if (!get_fpu(fd, &req)) {
		return (false);
	}

	/* With AVX unused so far, we expect it to be absent from the BV */
	if (xs->header.xsh_xstate_bv != (XFEATURE_LEGACY_FP | XFEATURE_SSE)) {
		(void) fprintf(stderr, "bad xstate_bv %lx, expected %lx\n",
		    xs->header.xsh_xstate_bv,
		    (XFEATURE_LEGACY_FP | XFEATURE_SSE));
		return (false);
	}

	avx->ymm[0] = UINT64_MAX;
	avx->ymm[2] = 2;

	/* first write without asserting AVX in BV */
	if (!set_fpu(fd, &req)) {
		return (false);
	}

	/* And check that the AVX state stays empty */
	bzero(fpu_area, fpu_size);
	if (!get_fpu(fd, &req)) {
		return (false);
	}
	if (xs->header.xsh_xstate_bv != (XFEATURE_LEGACY_FP | XFEATURE_SSE)) {
		(void) fprintf(stderr, "xstate_bv changed unexpectedly %lx\n",
		    xs->header.xsh_xstate_bv);
		return (false);
	}
	if (avx->ymm[0] != 0 || avx->ymm[2] != 0) {
		(void) fprintf(stderr, "YMM state changed unexpectedly "
		    "%lx %lx\n", avx->ymm[0], avx->ymm[2]);
		return (false);
	}

	/* Now write YMM and set the appropriate AVX BV state */
	avx->ymm[0] = UINT64_MAX;
	avx->ymm[2] = 2;
	xs->header.xsh_xstate_bv |= XFEATURE_AVX;
	if (!set_fpu(fd, &req)) {
		return (false);
	}

	/* ... and now check that it stuck */
	bzero(fpu_area, fpu_size);
	if (!get_fpu(fd, &req)) {
		return (false);
	}
	if ((xs->header.xsh_xstate_bv & XFEATURE_AVX) == 0) {
		(void) fprintf(stderr, "AVX missing from xstate_bv %lx\n",
		    xs->header.xsh_xstate_bv);
		return (false);
	}
	if (avx->ymm[0] != UINT64_MAX || avx->ymm[2] != 2) {
		(void) fprintf(stderr, "YMM state not preserved "
		    "%lx != %lx | %lx != %lx\n",
		    avx->ymm[0], UINT64_MAX, avx->ymm[2], 2);
		return (false);
	}


	return (true);
}

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

	ctx = create_test_vm(suite_name);
	if (ctx == NULL) {
		perror("could not open test VM");
		return (EXIT_FAILURE);
	}
	fd = vm_get_device_fd(ctx);

	struct vm_fpu_desc_entry entries[64];
	struct vm_fpu_desc desc = {
		.vfd_entry_data = entries,
		.vfd_num_entries = 64,
	};

	res = ioctl(fd, VM_DESC_FPU_AREA, &desc);
	if (res != 0) {
		perror("could not query fpu area description");
		goto bail;
	}

	/* Make sure the XSAVE area described for this machine is reasonable */
	if (desc.vfd_num_entries == 0) {
		(void) fprintf(stderr, "no FPU description entries found\n");
		goto bail;
	}
	if (desc.vfd_req_size < MIN_XSAVE_SIZE) {
		(void) fprintf(stderr, "required XSAVE size %lu < "
		    "expected %lu\n", desc.vfd_req_size, MIN_XSAVE_SIZE);
		goto bail;
	}

	const size_t fpu_size = desc.vfd_req_size;
	void *fpu_area = malloc(fpu_size);
	if (fpu_area == NULL) {
		perror("could not allocate fpu area");
		goto bail;
	}
	bzero(fpu_area, fpu_size);

	if (!check_sse(fd, &desc, fpu_area, fpu_size)) {
		goto bail;
	}
	if (!check_avx(fd, &desc, fpu_area, fpu_size)) {
		goto bail;
	}

	/* mission accomplished */
	vm_destroy(ctx);
	(void) printf("%s\tPASS\n", suite_name);
	return (EXIT_SUCCESS);

bail:
	vm_destroy(ctx);
	(void) printf("%s\tFAIL\n", suite_name);
	return (EXIT_FAILURE);
}