/*
 * arc4_harness.c - userland evidence harness for the pass-2 audit of
 * sys/libkern/arc4random.c (DragonFly BSD).
 *
 * The arc4 algorithm below is copied VERBATIM from sys/libkern/arc4random.c
 * (struct arc4_data, arc4_swap, arc4_randomstir, arc4_randbyte, karc4random,
 * karc4random64, karc4random_buf) with only the kernel environment shimmed:
 *   - read_random()  -> filled from getrandom(2), always returns ARC4_KEYBYTES
 *     (post-boot invariant: sys/kern/kern_nrandom.c:718 returns the full
 *     csprng_get_random() count; sys/kern/subr_csprng.c:151-168 always
 *     produces `bytes` for a 32-byte request).
 *   - time_uptime   -> controllable global
 *   - mycpuid       -> 0 (single per-CPU state)
 *
 * Modes:
 *   stats N   generate N bytes through the real API mix (karc4random /
 *             karc4random64 / karc4random_buf) and run chi-square
 *             (byte freq) + serial (digraph) tests.
 *   race      deterministic sub-step interleaving of two contexts sharing
 *             one arc4 state (models LWKT preemption of a syscall thread
 *             by a netisr/TCP thread, sys/kern/lwkt_thread.c:941).
 *             Proves: joint output cannot be a serialization of any RC4
 *             stream from the same initial state, and scans for
 *             cross-consumer output duplication.
 *   bias N    targeted RC4 bias probes on N bytes: Pr[Z==0] excess,
 *             first-word-after-reseed uniformity (drop effectiveness),
 *             digraph-repeat counts at small lags.
 *
 * Build: cc -O2 -o arc4_harness arc4_harness.c -lm
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <math.h>
#include <sys/random.h>

#define	ARC4_MAXRUNS		16384
#define	ARC4_RESEED_SECONDS	300
#define	ARC4_KEYBYTES		32	/* 256 bit key */

struct arc4_data {
	uint8_t			arc4_i;
	uint8_t			arc4_j;
	int			arc4_numruns;
	time_t			arc4_nextreseed;
	uint8_t			arc4_sbox[256];
	/* harness-only: sub-step staging (not part of kernel struct) */
	uint8_t			h_t;
};

static struct arc4_data d0;	/* the single (cpu0) state */
static time_t time_uptime = 1000000;	/* past initial reseed epoch */

/* kernel shim: post-boot read_random(buf, 32, 1) always yields 32 bytes */
static int
read_random_shim(void *buf, unsigned size, int unlimited)
{
	(void)unlimited;
	if (getrandom(buf, size, 0) != (ssize_t)size)
		abort();
	return (int)size;
}

static uint8_t arc4_randbyte(struct arc4_data *);

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

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

/*
 * Stir our S-box.   [verbatim: sys/libkern/arc4random.c:50-82]
 */
static void
arc4_randomstir(struct arc4_data *d)
{
	uint8_t key[256];
	int r, n;

	r = read_random_shim(key, ARC4_KEYBYTES, 1);
	/* If r == 0 || -1, just use what was on the stack. */
	if (r > 0) {
		for (n = r; n < (int)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(&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(d);

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

/*
 * Generate a random byte.   [verbatim: sys/libkern/arc4random.c:87-99]
 */
static uint8_t
arc4_randbyte(struct arc4_data *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(&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: sys/libkern/arc4random.c:101-117 */
uint32_t
karc4random(void)
{
	struct arc4_data *d = &d0;
	uint32_t ret;

	if (++(d->arc4_numruns) > ARC4_MAXRUNS ||
	    time_uptime > d->arc4_nextreseed)
		arc4_randomstir(d);

	ret = (uint64_t)arc4_randbyte(d);
	ret |= (uint64_t)arc4_randbyte(d) << 8;
	ret |= (uint64_t)arc4_randbyte(d) << 16;
	ret |= (uint64_t)arc4_randbyte(d) << 24;

	return ret;
}

/* verbatim: sys/libkern/arc4random.c:119-139 */
uint64_t
karc4random64(void)
{
	struct arc4_data *d = &d0;
	uint64_t ret;

	if (++(d->arc4_numruns) > ARC4_MAXRUNS ||
	    time_uptime > d->arc4_nextreseed)
		arc4_randomstir(d);

	ret = (uint64_t)arc4_randbyte(d);
	ret |= (uint64_t)arc4_randbyte(d) << 8;
	ret |= (uint64_t)arc4_randbyte(d) << 16;
	ret |= (uint64_t)arc4_randbyte(d) << 24;
	ret |= (uint64_t)arc4_randbyte(d) << 32;
	ret |= (uint64_t)arc4_randbyte(d) << 40;
	ret |= (uint64_t)arc4_randbyte(d) << 48;
	ret |= (uint64_t)arc4_randbyte(d) << 56;

	return ret;
}

/* verbatim: sys/libkern/arc4random.c:141-162 (crit_enter/exit are #if 0) */
void
karc4random_buf(void *ptr, size_t len)
{
	struct arc4_data *d = &d0;
	uint8_t *p = ptr;

	if (++(d->arc4_numruns) > ARC4_MAXRUNS ||
	    time_uptime > d->arc4_nextreseed)
		arc4_randomstir(d);

	while (len--)
		*p++ = arc4_randbyte(d);
}

/* verbatim: sys/libkern/arc4random.c:167-193 */
static void
arc4_init_pcpu(struct arc4_data *d)
{
	int n;

	memset(d, 0, sizeof(*d));

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

	arc4_randomstir(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(d);
}

/* ------------------------------------------------------------------ */
/* stats mode                                                          */
/* ------------------------------------------------------------------ */
static void
run_stats(unsigned long long nbytes)
{
	unsigned long long freq[256] = {0};
	unsigned long long digraph[65536] = {0};
	uint8_t *buf = malloc(1 << 20);
	unsigned long long done = 0;
	unsigned prev = 0;
	int alt = 0;
	double x2, e, diff;

	if (!buf) abort();
	while (done < nbytes) {
		unsigned long long chunk =
		    (nbytes - done) > (1u << 20) ? (1u << 20) : nbytes - done;
		unsigned long long i;
		if (alt % 2 == 0) {
			for (i = 0; i + 4 <= chunk; i += 4) {
				uint32_t v = karc4random();
				memcpy(buf + i, &v, 4);
			}
			if (i < chunk) karc4random_buf(buf + i, chunk - i);
		} else {
			/* karc4random64 deliberately excluded: it is biased
			 * by the sign-extension bug (see bias64 mode) and
			 * would contaminate the distribution battery. */
			karc4random_buf(buf, chunk);
		}
		alt++;
		for (i = 0; i < chunk; i++) {
			freq[buf[i]]++;
			digraph[(prev << 8) | buf[i]]++;
			prev = buf[i];
		}
		done += chunk;
	}
	free(buf);

	e = (double)done / 256.0;
	x2 = 0;
	for (int i = 0; i < 256; i++) {
		diff = (double)freq[i] - e;
		x2 += diff * diff / e;
	}
	printf("bytes=%llu chi2_bytefreq=%.1f (df=255, mean=255, sd=%.1f, z=%+.2f)\n",
	       done, x2, sqrt(2 * 255), (x2 - 255) / sqrt(2 * 255));

	e = (double)(done - 1) / 65536.0;
	x2 = 0;
	for (int i = 0; i < 65536; i++) {
		diff = (double)digraph[i] - e;
		x2 += diff * diff / e;
	}
	printf("chi2_serial=%.1f (df=65535, mean=65535, sd=%.1f, z=%+.2f)\n",
	       x2, sqrt(2 * 65535), (x2 - 65535) / sqrt(2 * 65535));
}

/* ------------------------------------------------------------------ */
/* race mode: sub-step interleaving of two contexts on one state       */
/*                                                                     */
/* arc4_randbyte decomposed into its memory-visible sub-steps:         */
/*   S1: i = i+1                                                       */
/*   S2: j = j + sbox[i]                                               */
/*   S3: swap(sbox[i], sbox[j])                                        */
/*   S4: t = sbox[i] + sbox[j]   (t staged)                            */
/*   S5: out = sbox[t]                                                 */
/* Preemption at ANY sub-step boundary is instruction-reachable        */
/* (the kernel code runs with no crit section:                        */
/* sys/libkern/arc4random.c:147-161).                                  */
/*                                                                     */
/* Victim: syscall thread doing karc4random() (e.g. kern_exec.c:1173,  */
/*         vm_map.c:4408, pf.c:6131 outbound).                         */
/* Preemptor: runs complete karc4random() calls (netisr/TCP thread:    */
/*         pf.c:4459, tcp_syncache.c:1070, in_pcb.c:885).              */
/*                                                                     */
/* Check 1 (invariant break): emit bytes in true time order; compare   */
/*   with the serialized reference stream from the same initial state. */
/*   A synchronized implementation must be equivalent to SOME          */
/*   serialization; divergence proves the joint output is not any RC4  */
/*   stream (bytes consumed by neither consumer's contract).           */
/* Check 2 (duplication): does the victim's 4-byte word ever equal a   */
/*   preemptor word (cross-consumer value reuse)?                      */
/* ------------------------------------------------------------------ */
static uint8_t
rb_step(struct arc4_data *d, int step)	/* returns 1 if byte emitted */
{
	switch (step) {
	case 1: d->arc4_i = (d->arc4_i + 1) % 256; return 0;
	case 2: d->arc4_j = (d->arc4_j + d->arc4_sbox[d->arc4_i]) % 256; return 0;
	case 3: arc4_swap(&d->arc4_sbox[d->arc4_i], &d->arc4_sbox[d->arc4_j]); return 0;
	case 4: d->h_t = (d->arc4_sbox[d->arc4_i] + d->arc4_sbox[d->arc4_j]) % 256; return 0;
	case 5: return d->arc4_sbox[d->h_t];
	}
	abort();
}

static void
run_race(unsigned trials)
{
	struct arc4_data race, ref;
	unsigned long long tot_bytes = 0, divergent = 0;
	unsigned long long dup_words = 0;
	unsigned trial;

	for (trial = 0; trial < trials; trial++) {
		/* fresh shared state, deterministic per trial via getrandom */
		uint8_t seedkey[32];
		struct arc4_data base;
		int n, b;

		memset(&base, 0, sizeof(base));
		for (n = 0; n < 256; n++)
			base.arc4_sbox[n] = (uint8_t)n;
		if (getrandom(seedkey, sizeof(seedkey), 0) != (int)sizeof(seedkey))
			abort();
		for (n = 0; n < 256; n++) {
			base.arc4_j = (base.arc4_j + base.arc4_sbox[n] +
			    seedkey[n % 32]) % 256;
			arc4_swap(&base.arc4_sbox[n], &base.arc4_sbox[base.arc4_j]);
		}
		for (n = 0; n < 768 * 4; n++) {
			base.arc4_i = (base.arc4_i + 1) % 256;
			base.arc4_j = (base.arc4_j + base.arc4_sbox[base.arc4_i]) % 256;
			arc4_swap(&base.arc4_sbox[base.arc4_i],
			    &base.arc4_sbox[base.arc4_j]);
		}

		race = ref = base;

		/* time-ordered emission log for the interleaved run */
		uint8_t emitted[8 * 64];
		unsigned nem = 0;
		/* preemptor words, for duplication scan */
		uint32_t preempt_words[64];
		unsigned npw = 0;

		/*
		 * Victim executes karc4random() = 4 randbytes = 20 sub-steps.
		 * Preemption points: victim_byte in {0..3}, substep in {1..5}
		 * (preemption before S1 == between-byte boundary: included
		 * as substep==1 with the convention "run preemptor first").
		 * Deterministic schedule: rotate through all 4x5 points.
		 */
		unsigned npts = 20;
		unsigned pt = trial % npts;	/* chosen preemption point index */
		unsigned preempt_at_byte = pt / 5;
		unsigned preempt_at_step = (pt % 5);  /* 0 => between-byte */

		uint32_t victim_word = 0;
		int vpos = 0;

		for (b = 0; b < 4; b++) {
			for (int s = 1; s <= 5; s++) {
				if (b == (int)preempt_at_byte && s == (int)preempt_at_step + 1 &&
				    npw < 64) {
					/* LWKT preemption: another context
					 * runs 1 complete karc4random() */
					uint32_t w = 0;
					for (int k = 0; k < 4; k++) {
						uint8_t bb = arc4_randbyte(&race);
						emitted[nem++] = bb;
						w |= (uint32_t)bb << (8 * k);
					}
					preempt_words[npw++] = w;
				}
				if (rb_step(&race, s)) {
					uint8_t bb = race.arc4_sbox[race.h_t];
					emitted[nem++] = bb;
					victim_word |= (uint32_t)bb << (8 * vpos);
					vpos++;
				}
			}
		}

		/* Check 2: cross-consumer duplication */
		for (unsigned i = 0; i < npw; i++)
			if (preempt_words[i] == victim_word)
				dup_words++;

		/* Check 1: serialized reference from identical state */
		for (unsigned i = 0; i < nem; i++) {
			uint8_t rb = arc4_randbyte(&ref);
			tot_bytes++;
			if (rb != emitted[i])
				divergent++;
		}
	}
	printf("trials=%u timed_ordered_bytes=%llu divergent_from_serialized=%llu (%.1f%%)\n",
	       trials, tot_bytes, divergent,
	       tot_bytes ? 100.0 * (double)divergent / (double)tot_bytes : 0.0);
	printf("cross_consumer_duplicate_words=%llu\n", dup_words);
}

/* ------------------------------------------------------------------ */
/* bias mode                                                           */
/* ------------------------------------------------------------------ */
static void
run_bias(unsigned long long nbytes)
{
	unsigned long long z0 = 0, done = 0;
	unsigned long long fw[256] = {0};
	unsigned long long lagrep[9] = {0};
	uint8_t *buf = malloc(1 << 20);
	uint8_t ring[8] = {0};
	unsigned long long ringfill = 0;
	int nlag;

	if (!buf) abort();

	while (done < nbytes) {
		unsigned long long chunk =
		    (nbytes - done) > (1u << 20) ? (1u << 20) : nbytes - done;
		karc4random_buf(buf, chunk);
		for (unsigned long long i = 0; i < chunk; i++) {
			if (buf[i] == 0) z0++;
			if (done + i < 4096 * 4)
				fw[buf[i]]++;	/* first words after init drop */
			/* digraph-repeat: byte equals the one `lag' bytes back */
			for (nlag = 1; nlag <= 8; nlag++) {
				if ((unsigned long long)nlag <= ringfill &&
				    ring[(ringfill - nlag) & 7] == buf[i])
					lagrep[nlag]++;
			}
			ring[ringfill & 7] = buf[i];
			ringfill++;
		}
		done += chunk;
	}
	free(buf);

	double p = 1.0 / 256.0;
	double sd = sqrt((double)done * p * (1 - p));
	printf("bytes=%llu\n", done);
	printf("Pr[Z==0]: exp=%.0f got=%llu z=%+.2f\n",
	       (double)done * p, z0, ((double)z0 - (double)done * p) / sd);
	/* byte equals the byte `lag' positions back: P = 1/256 */
	double e2 = (double)done * (1.0 / 256.0);
	double sd2 = sqrt((double)done * (1.0 / 256.0) * (255.0 / 256.0));
	for (nlag = 1; nlag <= 8; nlag++)
		printf("lag%d_repeat: exp=%.0f got=%llu z=%+.2f\n", nlag, e2,
		    lagrep[nlag], ((double)lagrep[nlag] - e2) / sd2);
	/* first-word uniformity (Mironov drop effectiveness spot check) */
	double x2 = 0, e = (4096 * 4) / 256.0;
	for (int i = 0; i < 256; i++) {
		double diff = (double)fw[i] - e;
		x2 += diff * diff / e;
	}
	printf("first16k_word_freq chi2=%.1f (df=255, z=%+.2f)\n", x2,
	    (x2 - 255) / sqrt(2 * 255));
}

/* ------------------------------------------------------------------ */
/* bias64 mode: PoC for the karc4random64 sign-extension bug.          */
/*                                                                     */
/* arc4random.c:132:  ret |= (uint64_t)arc4_randbyte(d) << 24;                   */
/* arc4_randbyte() returns uint8_t -> promoted to int.  For bytes      */
/* >= 0x80 the shifted result is not representable in int (C11         */
/* 6.5.7p4 UB); real compilers emit shl (wrap) leaving bit 31 set,     */
/* i.e. a negative int.  The |= into uint64_t ret converts that        */
/* negative int modulo 2^64 -> SIGN EXTENSION -> bits 32..63 all set.  */
/* Later `|=' of bytes 4..7 cannot clear them.  Net effect: whenever   */
/* the 4th byte of the word is >= 0x80 (p = 1/2), the upper 32 bits    */
/* of karc4random64() are exactly 0xffffffff.                          */
/* ------------------------------------------------------------------ */
static void
run_bias64(unsigned long long ndraws)
{
	unsigned long long top32_ones = 0, b3_high = 0, both = 0;
	unsigned long long bucket[8] = {0};	/* vm_map.c addend buckets */
	unsigned long long vmhits = 0;

	for (unsigned long long k = 0; k < ndraws; k++) {
		uint64_t v = karc4random64();
		uint8_t b3 = (uint8_t)(v >> 24);
		int topset = ((v >> 32) & 0xffffffffULL) == 0xffffffffULL;

		if (topset) top32_ones++;
		if (b3 >= 0x80) b3_high++;
		if (topset && b3 >= 0x80) both++;

		/* replicate vm_map.c:4408 with default dsiz = 32GB */
		unsigned long long addend =
		    (v & 0x7FFFFFFFFFFFFFFFULL) % (32ULL * 1024 * 1024 * 1024);
		bucket[addend >> 32]++;		/* 8 buckets of 4GB */
		vmhits++;
	}
	printf("draws=%llu\n", ndraws);
	printf("P(top32==0xffffffff) = %.6f  (bug predicts ~0.500000)\n",
	    (double)top32_ones / (double)ndraws);
	printf("P(byte3>=0x80)       = %.6f  (must be ~0.500000)\n",
	    (double)b3_high / (double)ndraws);
	printf("P(both)              = %.6f\n", (double)both / (double)ndraws);
	printf("vm_map mmap-ASLR addend distribution over 8x4GB buckets "
	    "(uniform would be 0.125000 each):\n");
	for (int i = 0; i < 8; i++)
		printf("  bucket[%d] (top 4GB=%s): %.6f\n", i,
		    i == 7 ? "yes" : "no ",
		    (double)bucket[i] / (double)vmhits);
	/* min-entropy of the observed addend distribution */
	double hmin = 0.0, pmax = 0.0;
	for (int i = 0; i < 8; i++) {
		double p = (double)bucket[i] / (double)vmhits;
		if (p > pmax) pmax = p;
	}
	hmin = -log2(pmax) + 32;	/* per-byte within bucket ~uniform */
	printf("addend min-entropy ~ %.2f bits (ideal 35.00)\n", hmin);
}

int
main(int argc, char **argv)
{
	if (argc < 2) {
		fprintf(stderr, "usage: %s stats N | race [trials] | bias N |"
		    " bias64 N\n", argv[0]);
		return 2;
	}
	if (!strcmp(argv[1], "stats")) {
		if (argc < 3) return 2;
		arc4_init_pcpu(&d0);
		run_stats(strtoull(argv[2], NULL, 0));
		return 0;
	}
	if (!strcmp(argv[1], "race")) {
		unsigned trials = (argc > 2) ? (unsigned)strtoul(argv[2], NULL, 0) : 200;
		run_race(trials);
		return 0;
	}
	if (!strcmp(argv[1], "bias")) {
		if (argc < 3) return 2;
		arc4_init_pcpu(&d0);
		run_bias(strtoull(argv[2], NULL, 0));
		return 0;
	}
	if (!strcmp(argv[1], "bias64")) {
		if (argc < 3) return 2;
		arc4_init_pcpu(&d0);
		run_bias64(strtoull(argv[2], NULL, 0));
		return 0;
	}
	return 2;
}