/*
 * DF-5420 — pmap_inval_smp() single-cpu shortcut pmap->pm_active TOCTOU
 *          -> missed TLB invalidation -> persistent stale TLB entry.
 *
 * Race (see findings/DF-5420):
 *   sys/platform/pc64/x86_64/pmap_inval.c:282  shortcut decides the target
 *   set from pmap->pm_active WITHOUT holding CPULOCK_EXCL (and for
 *   non-PMAP_MULTI pmaps without any interlock at all):
 *
 *       if (CPUMASK_CMPMASKEQ(pmap->pm_active, gd->gd_cpumask) ...) {
 *               if (pmap->pm_flags & PMAP_MULTI)
 *                       pmap_inval_init(pmap);     <-- interlock taken only
 *               ...                                     AFTER the sample
 *               opte = atomic_swap_long(ptep, npte);
 *               cpu_invlpg((void *)va);           <-- local cpu only
 *
 *   A cpu concurrently switching INTO the pmap does (swtch.s:408-445):
 *       OR pm_active bit -> test CPULOCK_EXCL -> (clear) -> load_cr3 -> user
 *   If the OR lands after the invalidator's pm_active read and the EXCL
 *   test slips in before pmap_inval_init() sets it, that cpu never receives
 *   the invalidation, and any TLB entry it cached for the va before the
 *   atomic_swap_long() persists forever (until an unrelated full flush).
 *
 * Rig (per instance, process P):
 *   T0 pinned CPU A : munmap(X) ... hold ... mmap(X)+write MAGIC_NEW loop
 *   T1 pinned CPU B : jittered usleep + guarded loads of X
 * Helper process Q (stretcher): one thread spinning on CPU A (keeps Q's
 *   pmap active there), one thread on CPU C hammering munmap/mmap of its
 *   own page -> every invalidation broadcast-Xinvltlb-IPIs CPU A, landing
 *   inside T0's [pm_active read .. pte swap] window and stretching it from
 *   ~50ns to microseconds (Xinvltlb runs even inside critical sections).
 *
 * Detectors (both false-positive-proof on x86 TSO):
 *   A: guarded load of X SUCCEEDS while phase==PH_UNMAPPED on both sides.
 *      PH_UNMAPPED strictly implies *ptep == 0 (set after munmap() returns,
 *      cleared only after the post-remap fault-in store lands), so a
 *      non-faulting user load can only be served by a stale TLB entry.
 *   B: load returns MAGIC_OLD while phase==PH_MAPPED on both sides and
 *      last_written==MAGIC_NEW (current page provably contains NEW; an
 *      OLD result was served from the old physical page via stale TLB).
 *
 * Exit code 42 = stale TLB entry confirmed.
 */
#include <sys/types.h>
#include <sys/usched.h>
#include <sys/mman.h>
#include <sys/wait.h>
#include <sys/time.h>
#include <unistd.h>
#include <signal.h>
#include <setjmp.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <pthread.h>
#include <errno.h>
#include <sys/sysctl.h>

typedef unsigned long cpumask_ul[4];	/* kernel cpumask_t = 4 x u64 */

#define PAGE        4096UL
#define XADDR       ((void *)0x0000200000000000UL)
#define ZADDR       ((void *)0x0000300000000000UL)
#define MAGIC_OLD   0x4f4c444d41474943UL   /* "CIGAMOLD" */
#define MAGIC_NEW   0x4e45474d41474943UL   /* "CIGAMNEW" */

#define PH_UNMAPPED 1   /* *ptep == 0 guaranteed                */
#define PH_WARM     2   /* mmap done, fault-in store in flight  */
#define PH_MAPPED   3   /* PTE present, MAGIC_NEW stored        */

static volatile int  phase = PH_MAPPED;
static volatile unsigned long last_written = MAGIC_OLD;
static volatile unsigned long stop = 0;
static volatile unsigned long hits_a = 0, hits_b = 0;
static volatile int freeze_req = 0, freeze_ack = 0;
static volatile unsigned long t0_iter = 0, t1_iter = 0, reads_ok = 0;
static volatile unsigned long q_iter = 0;
static volatile int pin_ok = 1;

static void *
xmap(void *addr)
{
	void *p = mmap(addr, PAGE, PROT_READ | PROT_WRITE,
		       MAP_ANON | MAP_PRIVATE | MAP_FIXED, -1, 0);
	if (p == MAP_FAILED) {
		perror("mmap");
		exit(1);
	}
	return (p);
}

static int
pin_me(int cpu)
{
	cpumask_ul mask;
	int rc;

	memset(mask, 0, sizeof(mask));
	mask[cpu / 64] = 1UL << (cpu % 64);
	rc = usched_set(0, USCHED_SET_CPUMASK, mask, sizeof(mask));
	if (rc != 0)
		fprintf(stderr, "usched_set(cpu %d) failed: %s "
			"(continuing unpinned)\n", cpu, strerror(errno));
	return (rc);
}

/* ------------------------------------------------------------------ */
/* T0: the invalidator (munmap hammer) on CPU A                       */
/* ------------------------------------------------------------------ */
static void *
t0_main(void *arg)
{
	int cpuA = (int)(intptr_t)arg;
	uint64_t v = 0;
	int hold_us;
	unsigned seed = 0x5eed0;

	if (pin_me(cpuA) != 0)
		pin_ok = 0;

	while (!stop) {
		if (munmap(XADDR, PAGE) != 0) {
			perror("munmap");
			exit(1);
		}
		phase = PH_UNMAPPED;          /* PTE provably absent now */
		t0_iter++;

		/*
		 * Hold the unmapped window (us).  Check freeze_req DURING
		 * the hold so the phantom probe can freeze us with the PTE
		 * ABSENT (this is what makes the probe meaningful).
		 */
		hold_us = 20 + rand_r(&seed) % 60;
		while (hold_us > 0) {
			if (freeze_req) {
				freeze_ack = 1;
				while (freeze_req && !stop)
					usleep(1000);
				freeze_ack = 0;
			}
			usleep(5);
			hold_us -= 5;
		}

		xmap(XADDR);
		phase = PH_WARM;
		/* fault-in store installs the new PTE */
		*(volatile uint64_t *)XADDR = MAGIC_NEW + (v++ & 0xff);
		phase = PH_MAPPED;
		last_written = MAGIC_NEW;

		usleep(20 + rand_r(&seed) % 60);
	}
	return (NULL);
}

/* ------------------------------------------------------------------ */
/* T1: the re-activating cpu (sleep/wake + guarded read) on CPU B     */
/* ------------------------------------------------------------------ */
static sigjmp_buf jb;
static volatile sig_atomic_t in_guard = 0;

static void
segv_handler(int sig)
{
	(void)sig;
	if (in_guard)
		siglongjmp(jb, 1);
	_exit(8);
}

static int probe_pipe[2] = { -1, -1 };

static int
guard_read(uint64_t *out)
{
	if (sigsetjmp(jb, 1) == 0) {
		in_guard = 1;
		__asm __volatile("" ::: "memory");
		*out = *(volatile uint64_t *)XADDR;
		__asm __volatile("" ::: "memory");
		in_guard = 0;
		return (1);
	}
	in_guard = 0;
	return (0);
}

/*
 * Frozen phantom probe — run in the detecting thread immediately after
 * a hit, with T0's munmap/remap loop frozen so the PTE state cannot
 * change during the probe:
 *
 *   (1) kernel-mode copyin of the same va (write(2)):
 *         EFAULT  => PTE truly absent in the page tables (kernel view);
 *   (2) user re-read AFTER a syscall round-trip: under meltdown
 *       mitigation every kernel/user transition reloads %cr3 (no PCID),
 *       flushing the non-global TLB.  Survival across the syscall means
 *       the page WALK still finds a mapping (software-level phantom,
 *       e.g. iso-tree desync), NOT a TLB entry;
 *   (3) death at the first syscall means the hit was a stale TLB entry.
 */
static void
probe_phantom(const char *kind, uint64_t firstval)
{
	struct timeval tv;
	uint64_t v = 0;
	int r, spin, ok;

	freeze_req = 1;
	for (spin = 0; spin < 2000 && !freeze_ack; spin++)
		usleep(1000);			/* up to 2s for T0 to freeze */

	/* On the instrumented kernel (root): dump the kernel-side trace
	 * ring NOW, frozen at hit time, so the last ring entries are the
	 * munmap command(s) surrounding the stale window. */
	if (geteuid() == 0) {
		int one = 1;
		sysctlbyname("machdep.df5420_dump", NULL, NULL, &one,
			     sizeof(one));
	}

	r = write(probe_pipe[1], (void *)XADDR, 8);   /* kernel copyin */
	gettimeofday(&tv, NULL);
	printf("  probe[%s]: frozen=%d phase=%d (1=UNMAPPED) "
	       "kernel-copyin write()=%d errno=%d "
	       "(first read val=0x%016lx) at %ld.%06ld\n",
	       kind, freeze_ack, phase, r, r < 0 ? errno : 0, firstval,
	       tv.tv_sec, tv.tv_usec);
	fflush(stdout);

	if (phase != PH_UNMAPPED) {
		printf("  probe[%s]: freeze did NOT land in unmapped hold — "
		       "probe invalid, skipping reads\n", kind);
		fflush(stdout);
		freeze_req = 0;
		while (freeze_ack)
			usleep(1000);
		return;
	}

	ok = guard_read(&v);
	printf("  probe[%s]: post-syscall user read #1: %s val=0x%016lx\n",
	       kind, ok ? "OK" : "FAULT", v);
	fflush(stdout);
	if (ok) {
		gettimeofday(&tv, NULL);	/* another syscall */
		ok = guard_read(&v);
		printf("  probe[%s]: post-syscall user read #2: %s val=0x%016lx "
		       "(t=%ld.%06ld)\n",
		       kind, ok ? "OK" : "FAULT", v, tv.tv_sec, tv.tv_usec);
		fflush(stdout);
	}
	freeze_req = 0;
	while (freeze_ack)
		usleep(1000);
}

static void
report(const char *kind, uint64_t v, int s1, int s2,
       unsigned long it, unsigned long rk)
{
	struct timeval tv;

	gettimeofday(&tv, NULL);
	printf("[%ld.%06ld] STALE-TLB-HIT type=%s val=0x%016lx "
	       "(s1=%d s2=%d phase=%d last_written=0x%lx) "
	       "t1_iter=%lu reads_ok=%lu t0_iter=%lu\n",
	       tv.tv_sec, tv.tv_usec, kind, v, s1, s2, phase,
	       last_written, it, rk, t0_iter);
	fflush(stdout);
	if (probe_pipe[1] >= 0)
		probe_phantom(kind, v);
}

static void *
t1_main(void *arg)
{
	int cpuB = (int)(intptr_t)arg;
	unsigned seed = 0x5eed1;
	uint64_t v;
	int ok, s1, s2, burst;

	if (pin_me(cpuB) != 0)
		pin_ok = 0;

	while (!stop) {
		usleep(2 + rand_r(&seed) % 78);   /* block -> pm_active bit B cleared */

		for (burst = 0; burst < 4 && !stop; burst++) {
			s1 = phase;
			ok = 0;
			v = 0;
			if (sigsetjmp(jb, 1) == 0) {
				in_guard = 1;
				__asm __volatile("" ::: "memory");
				v = *(volatile uint64_t *)XADDR;
				__asm __volatile("" ::: "memory");
				in_guard = 0;
				ok = 1;
			} else {
				in_guard = 0;
			}
			s2 = phase;
			t1_iter++;
			if (ok) {
				reads_ok++;
				if (s1 == PH_UNMAPPED && s2 == PH_UNMAPPED) {
					hits_a++;
					report("A(unmapped-load-ok)", v, s1, s2,
					       t1_iter, reads_ok);
				} else if (s1 == PH_MAPPED && s2 == PH_MAPPED &&
					   v == MAGIC_OLD &&
					   last_written == MAGIC_NEW) {
					hits_b++;
					report("B(old-magic)", v, s1, s2,
					       t1_iter, reads_ok);
				}
			}
		}
	}
	return (NULL);
}

/* ------------------------------------------------------------------ */
/* Q: stretcher process — keeps its pmap active on CPU A while its    */
/* other thread on CPU C hammers invalidations -> Xinvltlb storm to A */
/* ------------------------------------------------------------------ */
static void *
q_hammer(void *arg)
{
	int cpuC = (int)(intptr_t)arg;
	pid_t ppid = getppid();

	if (pin_me(cpuC) != 0)
		pin_ok = 0;

	/* DF-5420: map Z once before the hammer loop (munmap of an
	 * unmapped range returns EINVAL and would kill the stretcher). */
	*(volatile uint64_t *)xmap(ZADDR) = 1;

	while (!stop && getppid() == ppid) {
		if (munmap(ZADDR, PAGE) != 0) {
			perror("Q munmap");
			exit(1);
		}
		xmap(ZADDR);
		*(volatile uint64_t *)ZADDR = 1;
		q_iter++;
	}
	return (NULL);
}

static void
run_q(int cpuA, int cpuC)
{
	pthread_t th;
	pid_t ppid = getppid();

	if (pin_me(cpuA) != 0)      /* Q's main thread lives on CPU A */
		pin_ok = 0;
	if (pthread_create(&th, NULL, q_hammer, (void *)(intptr_t)cpuC) != 0) {
		perror("pthread_create q");
		exit(1);
	}
	while (!stop && getppid() == ppid)
		__asm __volatile("pause");   /* stay runnable on CPU A */
	pthread_join(th, NULL);
	_exit(0);
}

int
main(int argc, char **argv)
{
	int secs = (argc > 1) ? atoi(argv[1]) : 120;
	int storm = (argc > 2) ? atoi(argv[2]) : 1;
	int cpuA = (argc > 3) ? atoi(argv[3]) : 0;
	int cpuB = (argc > 4) ? atoi(argv[4]) : 1;
	int cpuC = (argc > 5) ? atoi(argv[5]) : 2;
	pthread_t t0, t1;
	struct sigaction sa;

	memset(&sa, 0, sizeof(sa));
	sa.sa_handler = segv_handler;
	sigemptyset(&sa.sa_mask);
	sigaction(SIGSEGV, &sa, NULL);
	if (pipe(probe_pipe) != 0) {
		perror("pipe");
		exit(1);
	}

	*(volatile uint64_t *)xmap(XADDR) = MAGIC_OLD;
	phase = PH_MAPPED;
	last_written = MAGIC_OLD;

	if (pin_me(cpuA) != 0)
		pin_ok = 0;

	if (storm) {
		pid_t q = fork();
		if (q == 0)
			run_q(cpuA, cpuC);    /* noreturn */
	}

	pthread_create(&t0, NULL, t0_main, (void *)(intptr_t)cpuA);
	pthread_create(&t1, NULL, t1_main, (void *)(intptr_t)cpuB);

	printf("DF-5420 rig: secs=%d storm=%d cpus T0=%d T1=%d Q=%d pin_ok=%d "
	       "pid=%d\n", secs, storm, cpuA, cpuB, cpuC, pin_ok, getpid());
	fflush(stdout);

	sleep(secs);
	stop = 1;
	pthread_join(t0, NULL);
	pthread_join(t1, NULL);
	if (storm)
		wait(NULL);

	printf("DONE: hits_a=%lu hits_b=%lu t0_iter=%lu t1_iter=%lu "
	       "reads_ok=%lu pin_ok=%d\n",
	       hits_a, hits_b, t0_iter, t1_iter, reads_ok, pin_ok);
	fflush(stdout);

	return ((hits_a || hits_b) ? 42 : 0);
}