/*
 * arc4x.c -- DF-2243 + DF-2244 verification experiments, compiled INTO the
 * test kernel (runtime kldload is broken on this guest: freshly built
 * modules GPF the kernel linker in strcmp during load, while the Jul-2
 * in-tree modules load fine; root cause not chased -- compiled-in is the
 * reliable path and a kernel build is needed for the boot-order probe and
 * fix validation anyway, since even loader-preloaded modules only get
 * their SYSINITs run at SI_BOOT2_KLD (kern_linker.c:1392), long after the
 * SI_BOOT2_START_APS window under test).
 *
 * Contents:
 *   1. Boot-order probe (SYSINITs around mp_start_aps)      [DF-2243 (a)]
 *   2. Per-boot determinism printer (first 8 karc4random)   [DF-2243 (c)]
 *   3. Stack-residue keying harness (verbatim arc4 copies)  [DF-2243 (b)]
 *   4. karc4random_buf run-accounting timing experiment     [DF-2244]
 *
 * Everything prints with the ARC4X/ARC4PROBE/arc4runs:/arc4stir: prefixes
 * and lands in the console boot log + msgbuffer.
 */
#include <sys/param.h>
#include <sys/types.h>
#include <sys/systm.h>
#include <sys/libkern.h>
#include <sys/time.h>
#include <sys/kernel.h>
#include <sys/random.h>
#include <sys/malloc.h>

#include <vm/vm_extern.h>
#include <machine/cpufunc.h>

static MALLOC_DEFINE(M_ARC4X, "arc4x", "DF-2243/2244 experiments");


/* ============ boot-order probe ============ */


static void
arc4boot_pre(void *arg __unused)
{
	uint8_t t[32];
	u_int r;

	r = read_random(t, 32, 1);
	kprintf("ARC4PROBE pre-aps : read_random r=%u first=%02x%02x%02x%02x "
	    "(expect r=32 => csprng_pcpu ready BEFORE arc4_init_pcpu)\n",
	    r, t[0], t[1], t[2], t[3]);
}

static void
arc4boot_post(void *arg __unused)
{
	uint8_t t[32];
	u_int r;

	r = read_random(t, 32, 1);
	kprintf("ARC4PROBE post-aps: read_random r=%u time_uptime=%ju "
	    "karc4random[0..7]=\n", r, (uintmax_t)time_uptime);
	kprintf("ARC4PROBE   %08x %08x %08x %08x %08x %08x %08x %08x\n",
	    karc4random(), karc4random(), karc4random(), karc4random(),
	    karc4random(), karc4random(), karc4random(), karc4random());
}

SYSINIT(arc4bp0, SI_BOOT2_START_APS, SI_ORDER_EARLIER, arc4boot_pre, NULL);
SYSINIT(arc4bp2, SI_BOOT2_START_APS, SI_ORDER_SECOND, arc4boot_post, NULL);

/* ------------------------------------------------------------------ */
/* Controllable stand-in for read_random()                             */
/* ------------------------------------------------------------------ */
static uint8_t v_inject[256];	/* bytes handed out when r > 0 */
static int v_mode;		/* 0: return 0; 32: inject 32; 256: inject 256 */
static int v_last_r;

static int
v_read_random(void *buf, u_int nbytes, int unlimited __unused)
{
	if (v_mode <= 0) {
		/* csprng_pcpu == NULL: writes NOTHING, returns 0 --
		 * exactly like kern_nrandom.c:693-696 */
		v_last_r = 0;
		return (0);
	}
	memcpy(buf, v_inject, v_mode);
	v_last_r = v_mode;
	return (v_mode);
}

/* ================================================================== */
/* VERBATIM copy of sys/libkern/arc4random.c:21-31, 37-45, 87-99       */
/* (renames only: arc4_data->arc4_data_v, arc4_swap->arc4_swap_v,      */
/*  arc4_randbyte->arc4_randbyte_v)                                    */
/* ================================================================== */
#define ARC4_MAXRUNS_V		16384
#define ARC4_RESEED_SECONDS_V	300
#define ARC4_KEYBYTES_V		32

struct arc4_data_v {
	uint8_t		arc4_i;
	uint8_t		arc4_j;
	int		arc4_numruns;
	time_t		arc4_nextreseed;
	uint8_t		arc4_sbox[256];
};

static struct arc4_data_v *d_v;	/* single instance (deviation: not pcpu) */

static __inline void
arc4_swap_v(uint8_t *a, uint8_t *b)
{
	uint8_t c;

	c = *a;
	*a = *b;
	*b = c;
}

/* VERBATIM copy of arc4random.c:87-99 */
static uint8_t
arc4_randbyte_v(struct arc4_data_v *d)
{
	uint8_t arc4_t;

	d->arc4_i = (d->arc4_i + 1) % 256;
	d->arc4_j = (d->arc4_j + d->arc4_sbox[d->arc4_i]) % 256;

	arc4_swap_v(&d->arc4_sbox[d->arc4_i], &d->arc4_sbox[d->arc4_j]);

	arc4_t = (d->arc4_sbox[d->arc4_i] + d->arc4_sbox[d->arc4_j]) % 256;
	return d->arc4_sbox[arc4_t];
}

/* ================================================================== */
/* VERBATIM copy of arc4random.c:50-82 (arc4_randomstir)              */
/* rename: arc4_randomstir->arc4_randomstir_v,                        */
/*         read_random->v_read_random                                 */
/* ================================================================== */
static void
arc4_randomstir_v(struct arc4_data_v *d)
{
	uint8_t key[256];
	int r, n;

	/*
	 * XXX read_random() returns unsafe numbers if the entropy
	 * device is not loaded -- MarkM.
	 */
	r = v_read_random(key, ARC4_KEYBYTES_V, 1);
	/* If r == 0 || -1, just use what was on the stack. */
	if (r > 0) {
		for (n = r; n < sizeof(key); n++)
			key[n] = key[n % r];
	}

	for (n = 0; n < 256; n++) {
		d->arc4_j = (d->arc4_j + d->arc4_sbox[n] + key[n]) % 256;
		arc4_swap_v(&d->arc4_sbox[n], &d->arc4_sbox[d->arc4_j]);
	}

	/*
	 * Discard early keystream, as per recommendations in:
	 * "(Not So) Random Shuffles of RC4" by Ilya Mironov.
	 */
	for (n = 0; n < 768 * 4; n++)
		arc4_randbyte_v(d);

	/* Reset for next reseed cycle. */
	d->arc4_nextreseed = time_uptime + ARC4_RESEED_SECONDS_V;
	d->arc4_numruns = 0;
}

/* ================================================================== */
/* VERBATIM copy of arc4random.c:101-117 (karc4random)                */
/* deviation: per-cpu array replaced by the single harness instance    */
/* ================================================================== */
static uint32_t
karc4random_v(void)
{
	struct arc4_data_v *d = d_v;
	uint32_t ret;

	if (++(d->arc4_numruns) > ARC4_MAXRUNS_V ||
	    time_uptime > d->arc4_nextreseed)
		arc4_randomstir_v(d);

	ret = arc4_randbyte_v(d);
	ret |= arc4_randbyte_v(d) << 8;
	ret |= arc4_randbyte_v(d) << 16;
	ret |= arc4_randbyte_v(d) << 24;

	return ret;
}

/* VERBATIM copy of arc4random.c:167-193 (arc4_init_pcpu)             */
/* deviation: kmem_alloc3 kept verbatim; lifecycle handled by caller   */
/* (free + pointer reset between cases so the KASSERT holds)           */
/* ================================================================== */
static void
arc4_init_pcpu_v(int cpuid)
{
	struct arc4_data_v *d;
	int n;

	KASSERT(d_v == NULL, ("arc4_v was initialized"));

	d = (void *)kmem_alloc3(kernel_map, sizeof(*d), VM_SUBSYS_GD,
	    KM_CPU(cpuid));
	memset(d, 0, sizeof(*d));

	for (n = 0; n < 256; n++)
		d->arc4_sbox[n] = (uint8_t)n;

	arc4_randomstir_v(d);

	/*
	 * Discard early keystream, as per recommendations in:
	 * "(Not So) Random Shuffles of RC4" by Ilya Mironov.
	 */
	for (n = 0; n < 768 * 4; n++)
		arc4_randbyte_v(d);

	d_v = d;
}

static void
arc4_fini_v(void)
{
	if (d_v) {
		kmem_free(kernel_map, (vm_offset_t)d_v,
			  sizeof(struct arc4_data_v));
		d_v = NULL;
	}
}

/* ================================================================== */
/* Instrumentation twin of arc4_randomstir_v: IDENTICAL logic, plus a  */
/* dump of the actual key[] bytes right after the r>0 fold.            */
/* (This is the only added line block, clearly marked.)                */
/* ================================================================== */
static void
arc4_randomstir_dump(struct arc4_data_v *d, uint8_t *keyout)
{
	uint8_t key[256];
	int r, n;

	r = v_read_random(key, ARC4_KEYBYTES_V, 1);
	if (r > 0) {
		for (n = r; n < sizeof(key); n++)
			key[n] = key[n % r];
	}

	/* ---- instrumentation only ---- */
	if (keyout)
		memcpy(keyout, key, sizeof(key));
	kprintf("arc4stir: DUMP r=%d key[0..255]:\n", r);
	for (n = 0; n < 256; n += 32) {
		kprintf("arc4stir:   %03x:", n);
		for (int q = 0; q < 32; q++)
			kprintf("%02x", key[n + q]);
		kprintf("\n");
	}
	/* ---- end instrumentation ---- */

	for (n = 0; n < 256; n++) {
		d->arc4_j = (d->arc4_j + d->arc4_sbox[n] + key[n]) % 256;
		arc4_swap_v(&d->arc4_sbox[n], &d->arc4_sbox[d->arc4_j]);
	}
	for (n = 0; n < 768 * 4; n++)
		arc4_randbyte_v(d);
	d->arc4_nextreseed = time_uptime + ARC4_RESEED_SECONDS_V;
	d->arc4_numruns = 0;
}

static void
arc4_init_pcpu_dump(int cpuid, uint8_t *keyout)
{
	struct arc4_data_v *d;
	int n;

	KASSERT(d_v == NULL, ("arc4_v was initialized"));

	d = (void *)kmem_alloc3(kernel_map, sizeof(*d), VM_SUBSYS_GD,
	    KM_CPU(cpuid));
	memset(d, 0, sizeof(*d));

	for (n = 0; n < 256; n++)
		d->arc4_sbox[n] = (uint8_t)n;

	arc4_randomstir_dump(d, keyout);

	for (n = 0; n < 768 * 4; n++)
		arc4_randbyte_v(d);

	d_v = d;
}

/* ================================================================== */
/* FIXED variant (proposed DF-2243 fix) for A/B validation             */
/* ================================================================== */
static void
arc4_randomstir_f(struct arc4_data_v *d)
{
	uint8_t key[256];
	int r, n;

	memset(key, 0, sizeof(key));
	r = v_read_random(key, ARC4_KEYBYTES_V, 1);
	if (r > 0) {
		for (n = r; n < sizeof(key); n++)
			key[n] = key[n % r];
	} else {
		struct timespec ts;
		static volatile int arc4_stircnt;
		uint64_t mix;

		nanouptime(&ts);
		mix = (uint64_t)ts.tv_nsec ^ (uint64_t)(uintptr_t)d ^
		    ((uint64_t)(unsigned)arc4_stircnt++ *
		     0x9E3779B97F4A7C15ULL);
		for (n = 0; n < (int)sizeof(key); n++) {
			key[n] ^= (uint8_t)(mix >> ((n & 7) * 8));
			if ((n & 7) == 7)
				mix = mix * 6364136223846793005ULL +
				      1442695040888963407ULL;
		}
	}

	for (n = 0; n < 256; n++) {
		d->arc4_j = (d->arc4_j + d->arc4_sbox[n] + key[n]) % 256;
		arc4_swap_v(&d->arc4_sbox[n], &d->arc4_sbox[d->arc4_j]);
	}
	for (n = 0; n < 768 * 4; n++)
		arc4_randbyte_v(d);
	d->arc4_nextreseed = time_uptime + ARC4_RESEED_SECONDS_V;
	d->arc4_numruns = 0;
}

static void
arc4_init_pcpu_f(int cpuid)
{
	struct arc4_data_v *d;
	int n;

	KASSERT(d_v == NULL, ("arc4_v was initialized"));

	d = (void *)kmem_alloc3(kernel_map, sizeof(*d), VM_SUBSYS_GD,
	    KM_CPU(cpuid));
	memset(d, 0, sizeof(*d));

	for (n = 0; n < 256; n++)
		d->arc4_sbox[n] = (uint8_t)n;

	arc4_randomstir_f(d);

	for (n = 0; n < 768 * 4; n++)
		arc4_randbyte_v(d);

	d_v = d;
}

/* ================================================================== */
/* Stack-poisoning drivers                                             */
/* ================================================================== */
static int v_pattern;	/* poison byte */

static void __noinline
run_case_v(int mode, uint32_t *ks, int nks)
{
	v_mode = mode;
	arc4_fini_v();
	arc4_init_pcpu_v(0);
	for (int i = 0; i < nks; i++)
		ks[i] = karc4random_v();
}

static void __noinline
run_case_dump(int mode, uint32_t *ks, int nks, uint8_t *keyout)
{
	v_mode = mode;
	arc4_fini_v();
	arc4_init_pcpu_dump(0, keyout);
	for (int i = 0; i < nks; i++)
		ks[i] = karc4random_v();
}

static void __noinline
run_case_f(int mode, uint32_t *ks, int nks)
{
	v_mode = mode;
	arc4_fini_v();
	arc4_init_pcpu_f(0);
	for (int i = 0; i < nks; i++)
		ks[i] = karc4random_v();
}

/*
 * Poison 1024 bytes of stack with the pattern, THEN descend into the
 * noinline runner whose frame (and the stir's key[256] below it) overlaps
 * the poisoned region.
 */
static void __noinline
poison_and_run_dump(int pattern, int mode, uint32_t *ks, int nks,
    uint8_t *keyout)
{
	volatile char pad[1024];

	for (int q = 0; q < (int)sizeof(pad); q++)
		pad[q] = (char)pattern;
	__asm __volatile("" ::: "memory");
	run_case_dump(mode, ks, nks, keyout);
}

static void __noinline
poison_and_run_f(int pattern, int mode, uint32_t *ks, int nks)
{
	volatile char pad[1024];

	for (int q = 0; q < (int)sizeof(pad); q++)
		pad[q] = (char)pattern;
	__asm __volatile("" ::: "memory");
	run_case_f(mode, ks, nks);
}

static void
print_ks(const char *tag, uint32_t *ks, int n)
{
	kprintf("arc4stir: %s:", tag);
	for (int i = 0; i < n; i++)
		kprintf(" %08x", ks[i]);
	kprintf("\n");
}

static int
ks_equal(uint32_t *a, uint32_t *b, int n)
{
	for (int i = 0; i < n; i++)
		if (a[i] != b[i])
			return (0);
	return (1);
}

static int
arc4stir_load(struct module *m, int what, void *arg)
{
	uint32_t ksA[8], ksA2[8], ksB[8], ksP[8], ksD1[8], ksD2[8];
	uint32_t ksF[8], ksF32[8], ksV32[8];
	uint8_t observed[256];
	int pat_hits, i;

	switch (what) {
	case MOD_LOAD:
		kprintf("arc4stir: === DF-2243 mechanism harness (cpu%d) ===\n",
		    mycpuid);

		/* ---- Case A: r=0, poisoned stack, dump actual key ---- */
		v_pattern = 0xA5;
		poison_and_run_dump(0xA5, 0, ksA, 8, observed);
		pat_hits = 0;
		for (i = 0; i < 256; i++)
			if (observed[i] == 0xA5)
				pat_hits++;
		kprintf("arc4stir: A(r=0,poison=A5): key bytes matching "
		    "poison pattern: %d/256 (non-pattern bytes = frame "
		    "junk: ret addrs, saved regs)\n", pat_hits);
		print_ks("ksA", ksA, 8);

		/* ---- Case A2: repeat -> determinism within boot ---- */
		poison_and_run_dump(0xA5, 0, ksA2, 8, NULL);
		print_ks("ksA2", ksA2, 8);
		kprintf("arc4stir: ksA2 == ksA : %s\n",
		    ks_equal(ksA, ksA2, 8) ? "YES (deterministic)" : "NO");

		/* ---- Case P: different poison -> different stream ---- */
		poison_and_run_dump(0x3C, 0, ksP, 8, NULL);
		print_ks("ksP(0x3C)", ksP, 8);
		kprintf("arc4stir: ksP == ksA : %s (expected NO)\n",
		    ks_equal(ksA, ksP, 8) ? "YES" : "NO");

		/* ---- Case B: r=256 injecting the observed residue ---- */
		memcpy(v_inject, observed, 256);
		run_case_v(256, ksB, 8);
		print_ks("ksB(RC4(observed))", ksB, 8);
		kprintf("arc4stir: *** ksB == ksA : %s  <== keystream is "
		    "RC4(stack residue) ***\n",
		    ks_equal(ksA, ksB, 8) ? "YES" : "NO");

		/* ---- Case D: normal path r=32, fixed entropy ---- */
		for (i = 0; i < 32; i++)
			v_inject[i] = (uint8_t)i;
		run_case_v(32, ksD1, 8);
		run_case_v(32, ksD2, 8);
		print_ks("ksD1(r=32)", ksD1, 8);
		kprintf("arc4stir: ksD2 == ksD1 : %s (normal path "
		    "deterministic given entropy)\n",
		    ks_equal(ksD1, ksD2, 8) ? "YES" : "NO");

		/* ---- Fix A/B: fixed stir, r=0 + same poison ---- */
		poison_and_run_f(0xA5, 0, ksF, 8);
		print_ks("ksF(fixed,r=0,poison=A5)", ksF, 8);
		kprintf("arc4stir: FIX: ksF == ksA : %s (expected NO -- "
		    "residue no longer the key)\n",
		    ks_equal(ksA, ksF, 8) ? "YES" : "NO");

		/* ---- Fix A/B: fixed stir, r=32 -> bit-identical ---- */
		for (i = 0; i < 32; i++)
			v_inject[i] = (uint8_t)i;
		run_case_f(32, ksF32, 8);
		run_case_v(32, ksV32, 8);
		kprintf("arc4stir: FIX: ksF32 == ksV32 : %s (expected YES -- "
		    "normal post-init path bit-identical)\n",
		    ks_equal(ksF32, ksV32, 8) ? "YES" : "NO");

		arc4_fini_v();
		kprintf("arc4stir: === done ===\n");
		break;
	}
	return (0);
}



#define NPH 20000

static uint64_t deltas[NPH];

static MALLOC_DEFINE(M_ARC4RUNS, "arc4runs", "DF-2244 test");

static void
report(const char *tag, int n)
{
	uint64_t mean = 0, maxv = 0, thresh;
	int i, nspike = 0;
	int top_i[5];
	uint64_t top_v[5];
	int t, j;

	for (i = 0; i < n; i++) {
		mean += deltas[i];
		if (deltas[i] > maxv)
			maxv = deltas[i];
	}
	mean /= n;
	thresh = 30 * mean;
	if (thresh < 20000)
		thresh = 20000;

	for (t = 0; t < 5; t++) {
		top_v[t] = 0;
		top_i[t] = -1;
	}
	for (i = 0; i < n; i++) {
		for (t = 0; t < 5; t++) {
			if (deltas[i] > top_v[t]) {
				for (j = 4; j > t; j--) {
					top_v[j] = top_v[j-1];
					top_i[j] = top_i[j-1];
				}
				top_v[t] = deltas[i];
				top_i[t] = i;
				break;
			}
		}
	}

	kprintf("arc4runs: %s n=%d mean=%llu max=%llu thresh=%llu\n",
	    tag, n, (unsigned long long)mean,
	    (unsigned long long)maxv, (unsigned long long)thresh);
	kprintf("arc4runs: %s top5:", tag);
	for (t = 0; t < 5; t++)
		kprintf(" [%d]=%llu", top_i[t], (unsigned long long)top_v[t]);
	kprintf("\n");
	kprintf("arc4runs: %s spikes(i>thresh):", tag);
	for (i = 0; i < n; i++) {
		if (deltas[i] > thresh) {
			kprintf(" %d", i);
			nspike++;
		}
	}
	kprintf("  (count=%d)\n", nspike);
}

static int
phase(const char *tag, int n)
{
	uint64_t t0;
	int i, cpu0 = mycpuid;

	for (i = 0; i < n; i++) {
		t0 = rdtsc();
		karc4random();
		deltas[i] = rdtsc() - t0;
	}
	if (mycpuid != cpu0)
		kprintf("arc4runs: %s WARNING migrated cpu%d->cpu%d, "
		    "results unreliable\n", tag, cpu0, mycpuid);
	report(tag, n);
	return (mycpuid == cpu0 ? 0 : 1);
}

static int
spikes_in_first(const char *tag, int n, int first)
{
	uint64_t mean = 0, thresh;
	int i, cnt = 0;

	for (i = 0; i < n; i++)
		mean += deltas[i];
	mean /= n;
	thresh = 30 * mean;
	if (thresh < 20000)
		thresh = 20000;
	for (i = 0; i < first && i < n; i++)
		if (deltas[i] > thresh)
			cnt++;
	kprintf("arc4runs: %s spikes in FIRST %d calls: %d\n", tag, first, cnt);
	return (cnt);
}

static int
arc4runs_load(struct module *m, int what, void *arg)
{
	uint8_t *mb;
	uint64_t t0, t1;
	int rc = 0;

	switch (what) {
	case MOD_LOAD:
		kprintf("arc4runs: === DF-2244 karc4random_buf run-accounting "
		    "experiment (cpu%d) ===\n", mycpuid);

		/* settle: absorbs any pending time-based reseed */
		t0 = rdtsc();
		karc4random();
		t1 = rdtsc() - t0;
		kprintf("arc4runs: settle delta=%llu%s\n",
		    (unsigned long long)t1,
		    t1 > 20000 ? "  <== time-based reseed spike" : "");

		mb = kmalloc(1 << 20, M_ARC4RUNS, M_WAITOK | M_ZERO);
		if (mb == NULL) {
			kprintf("arc4runs: no mem\n");
			return (ENOMEM);
		}

		rc |= phase("phaseA(20000x karc4random)", NPH);

		/* 5 x 1MiB through karc4random_buf */
		for (int k = 0; k < 5; k++) {
			t0 = rdtsc();
			karc4random_buf(mb, 1 << 20);
			t1 = rdtsc() - t0;
			kprintf("arc4runs: phaseB[%d] karc4random_buf(1MiB) "
			    "delta=%llu first16=%02x%02x%02x%02x%02x%02x%02x%02x"
			    "%02x%02x%02x%02x%02x%02x%02x%02x\n", k,
			    (unsigned long long)t1,
			    mb[0],mb[1],mb[2],mb[3],mb[4],mb[5],mb[6],mb[7],
			    mb[8],mb[9],mb[10],mb[11],mb[12],mb[13],mb[14],mb[15]);
		}

		rc |= phase("phaseC(20000x karc4random)", NPH);
		spikes_in_first("phaseC", NPH, 100);

		/*
		 * Expected on STOCK kernel: phaseA 1 spike near i=16384;
		 * phaseC 0 spikes in first 100 calls (5MiB charged only 5
		 * runs), next spike near i~12700-12800.
		 * Expected on FIXED kernel: phaseC spikes in first 100 != 0
		 * (262144*5 pending runs force immediate reseed).
		 */
		kfree(mb, M_ARC4RUNS);
		kprintf("arc4runs: === done (cpu%d) ===\n", mycpuid);
		break;
	}
	return (rc);
}



/* ================================================================== */
/* Master experiment runner: late, quiet, single-threaded boot point    */
/* ================================================================== */
static void
arc4x_run(void *arg __unused)
{
	kprintf("arc4x: === experiments start (cpu%d) ===\n", mycpuid);
	arc4stir_load(NULL, MOD_LOAD, NULL);
	arc4runs_load(NULL, MOD_LOAD, NULL);
	kprintf("arc4x: === experiments end ===\n");
}
SYSINIT(arc4x, SI_BOOT2_KLD, SI_ORDER_ANY, arc4x_run, NULL);