#define _GNU_SOURCE
#include <err.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <sys/ucontext.h>

static volatile unsigned char *buf, *xsave_addr;
static volatile int nfailures = 0;

static void sethandler(int sig, void (*handler)(int, siginfo_t *, void *),
		       int flags)
{
	struct sigaction sa;
	memset(&sa, 0, sizeof(sa));
	sa.sa_sigaction = handler;
	sa.sa_flags = SA_SIGINFO | flags;
	sigemptyset(&sa.sa_mask);
	if (sigaction(sig, &sa, 0))
		err(1, "sigaction");
}

static void clearhandler(int sig)
{
	struct sigaction sa;
	memset(&sa, 0, sizeof(sa));
	sa.sa_handler = SIG_DFL;
	sigemptyset(&sa.sa_mask);
	if (sigaction(sig, &sa, 0))
		err(1, "sigaction");
}

static void sigsegv(int sig, siginfo_t *si, void *ctx_void)
{
	ucontext_t *ctx = (ucontext_t*)ctx_void;

	unsigned long cr2 = (unsigned long)ctx->uc_mcontext.gregs[REG_CR2];
	unsigned long start = (unsigned long)buf;

	extern unsigned char xsave_insn[], after_xsave_insn[];
	
	if (ctx->uc_mcontext.gregs[REG_RIP] != (unsigned long)xsave_insn) {
		printf("Uncorrectable segfault\n");
		clearhandler(SIGSEGV);
		return;
	}

	if (si->si_code != SEGV_ACCERR) {
		printf("Segfault was %d (trap %d), not SEGV_ACCERR\n",
		       si->si_code, ctx->uc_mcontext.gregs[REG_TRAPNO]);
		clearhandler(SIGSEGV);
		return;
	}

	if (cr2 != (unsigned long)si->si_addr) {
		printf("CR2 (0x%lx) != si_addr (0x%lx)\n",
		       cr2, (unsigned long)si->si_addr);
		clearhandler(SIGSEGV);
		return;
	}

	if (cr2 >= start && cr2 <= (start + 4095)) {
		printf("[OK]\txsave offset = %d, cr2 offset = %d\n",
		       (int)(xsave_addr - buf), (int)(cr2 - start));
	} else if (cr2 >= start + 4096 && cr2 <= start + 8191) {
		printf("[FAIL]\txsave offset = %d, cr2 offset = %d\n",
		       (int)(xsave_addr - buf), (int)(cr2 - start));

		nfailures++;
	} else if (cr2 >= start + 8192 && cr2 <= start + 12287) {
		printf("[OK]\txsave offset = %d, cr2 offset = %d\n",
		       (int)(xsave_addr - buf), (int)(cr2 - start));
	} else {
		printf("[FAIL]\tcr2 is completely out of range\n");
		abort();
	}

	ctx->uc_mcontext.gregs[REG_RIP] = (unsigned long)after_xsave_insn;
}

int main()
{
	int i;

	buf = mmap(NULL, 4096*3, PROT_NONE,
		   MAP_PRIVATE | MAP_ANONYMOUS | MAP_NORESERVE, -1, 0);
	if (buf == MAP_FAILED)
		err(1, "mmap");

	if (mmap((unsigned char *)buf + 4096, 4096, PROT_READ | PROT_WRITE,
		 MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED, -1, 0) == MAP_FAILED)
		err(1, "mmap");

	sethandler(SIGSEGV, sigsegv, 0);

	for (i = 0; i < 8193; i += 64) {
		xsave_addr = buf + i;
		printf("XSAVE to offset %d\n", i);
		asm volatile ("xsave_insn: xsaveq %0 ; after_xsave_insn:"
			      : "=m" (*xsave_addr)
			      : "a" (0xffffffff), "d" (0xffffffff));
	}

	if (nfailures)
		printf("%d failures\n", nfailures);
	else
		printf("PASS!\n");

	return 0;
}