/*
 * DF-2232 verification rig -- DragonFlyBSD kernel module (KLD).
 *
 * Finding: sys/libkern/linux_idr.c idr_grow() growth loop
 *      nf = idp->idr_count;
 *      do { nf = 2 * nf + 1; } while (nf <= want);
 * computes 2*nf+1 in a *signed int*.  For want >= 2^29 the loop exits
 * with nf = 2^30-1 and kmalloc()s (2^30-1)*16 = 17179869168 bytes (~16GB,
 * M_ZERO|M_WAITOK).  For want == INT_MAX the predicate stays true forever:
 * 2*INT_MAX+1 wraps to -1, -1 <= INT_MAX -> infinite spin while the lwkt
 * idr token is held.
 *
 * Rig design (verbatim-code KLD):
 *   The REAL sys/libkern/linux_idr.c is #included byte-identical (the copy
 *   placed next to this file is sha256-verified against the audited tree by
 *   build.sh).  Only the exported symbol NAMES are macro-renamed so the KLD
 *   does not collide with the kernel's own libkern copies -- the algorithm
 *   is untouched.  A kmalloc() observation wrapper (OUTSIDE the verbatim
 *   file) captures the exact byte size requested at the allocator boundary
 *   and then hands the request to the real allocator unchanged.
 *
 *   The kernel's own exported idr_* KPI (the shipped binary) is ALSO
 *   exercised directly for the control and spin phases.
 *
 * Phases (sysctl debug.df2232=<n>, write triggers):
 *   1 = control, verbatim copy:   normal wants 1..1000 round-trip
 *   2 = control, kernel KPI:      same, on the running kernel's own symbols
 *   3 = GIANT KMALLOC:            idr_alloc(start=1<<29) via verbatim copy
 *                                 (watchdog-protected, 15s)
 *   4 = SPIN, verbatim copy:      sid=INT_MAX -> pre_get -> infinite loop;
 *                                 second thread denied the idr token
 *   5 = SPIN, kernel KPI:         same on the shipped binary
 *
 * Watchdog: a systimer armed on a remote CPU fires from hard-interrupt
 * context (immune to kernel-mode spin loops) and panics with a marker
 * string plus the phase state flags, so every phase terminates with
 * decisive output.
 */

#include <sys/param.h>
#include <sys/kernel.h>
#include <sys/systm.h>
#include <sys/libkern.h>
#include <sys/malloc.h>
#include <sys/sysctl.h>
#include <sys/errno.h>
#include <sys/idr.h>
#include <sys/thread.h>
#include <sys/thread2.h>
#include <sys/kthread.h>
#include <sys/globaldata.h>
#include <sys/systimer.h>
#include <machine/cpu.h>
#include <sys/spinlock2.h>
#include <sys/limits.h>

/*
 * Phase state (cross-thread, read by the hard-IRQ watchdog)
 */
static volatile int df2232_a_stage;	/* 0=not run 1=EAGAIN seen 2=in pre_get */
static volatile int df2232_a_done;	/* A returned from pre_get */
static volatile int df2232_b_started;	/* B entered post-spin idr_remove */
static volatile int df2232_b_acquired;	/* B's idr_remove returned */
static volatile unsigned long df2232_last_req;
static int df2232_armed;
static int df2232_threads_live;
static int df2232_danger_ran;	/* any crash-phase ran: stay resident */
static int df2232_use_kernel;

static struct idr df2232_idr_v;		/* verbatim-copy idr (control/giant) */
static struct idr df2232_idr_k;		/* kernel-KPI idr (control) */
static struct idr df2232_idr_c;		/* spin-phase idr */

#define DF2232_WD_SECONDS	15

static struct systimer df2232_wd;
static volatile int df2232_wd_ticks;
static volatile int df2232_wd_cancel;
static volatile int df2232_wd_ctrl;	/* 1 = stop the watchdog thread */
static const char *df2232_wd_marker = "(none)";

/*
 * Watchdog callback -- runs from hard-interrupt context (systimer).
 */
static void
df2232_watchdog(struct systimer *st __unused, int inintr __unused,
		struct intrframe *f __unused)
{
	if (df2232_wd_cancel)
		return;
	if (++df2232_wd_ticks < DF2232_WD_SECONDS * 2)
		return;
	kprintf("[DF2232-WD] FIRED after ~%d s\n", DF2232_WD_SECONDS);
	kprintf("[DF2232-WD] a_stage=%d a_done=%d b_started=%d b_acquired=%d "
		"last_kmalloc_req=%lu bytes\n",
		df2232_a_stage, df2232_a_done, df2232_b_started,
		df2232_b_acquired, df2232_last_req);
	panic("DF2232-WD: %s", df2232_wd_marker);
}

/*
 * Watchdog thread: created on a CPU different from the trigger so the
 * systimer is armed on a CPU that no phase ever monopolizes.
 */
static void
df2232_wd_thread(void *arg __unused)
{
	kprintf("[DF2232] watchdog thread on cpu %d\n", mycpu->gd_cpuid);
	systimer_init_periodic(&df2232_wd, df2232_watchdog, NULL, 2);
	for (;;) {
		tsleep(&df2232_wd_ctrl, 0, "df2232wd", hz);
		if (df2232_wd_ctrl) {
			systimer_del(&df2232_wd);
			kprintf("[DF2232] watchdog thread stopped+exiting\n");
			break;
		}
	}
	kthread_exit();
}

static void
df2232_arm_watchdog(const char *marker)
{
	int cpu = (mycpu->gd_cpuid + 3) % ncpus;

	df2232_wd_ticks = 0;
	df2232_wd_cancel = 0;
	df2232_wd_ctrl = 0;
	df2232_wd_marker = marker;
	df2232_armed = 1;
	kprintf("[DF2232] arming %d s watchdog on cpu %d\n",
		DF2232_WD_SECONDS, cpu);
	kthread_create_cpu(df2232_wd_thread, NULL, NULL, cpu,
			   "df2232_wd");
}

static void
df2232_disarm_watchdog(void)
{
	df2232_wd_cancel = 1;
	df2232_wd_ctrl = 1;
	wakeup(&df2232_wd_ctrl);
	/* give the pinned watchdog thread a moment to systimer_del+exit */
	tsleep(&df2232_wd_ctrl, 0, "df2232d", 2);
	df2232_armed = 0;
	kprintf("[DF2232] watchdog disarmed\n");
}

/*=====================================================================
 * Verbatim inclusion of sys/libkern/linux_idr.c.
 *
 * KLD guard: rename the file's exported globals (and its malloc type) so
 * they do not collide with the kernel's own libkern/linux_idr.c symbols.
 * The algorithm itself is compiled byte-identical from the pristine file
 * (sha256-checked by build.sh).  All headers the file includes are already
 * included above, so the kmalloc() observation wrapper defined below
 * survives (malloc.h's own #define of kmalloc is not reprocessed).
 *=====================================================================*/

/* prototypes for the renamed KPI (declared before the renames so the
 * identifiers here are NOT themselves renamed) */
int	df2232_idr_pre_get(struct idr *, unsigned);
int	df2232_idr_get_new_above(struct idr *, void *, int, int *);
int	df2232_idr_get_new(struct idr *, void *, int *);
int	df2232_idr_alloc(struct idr *, void *, int, int, unsigned);
void	*df2232_idr_remove(struct idr *, int);
void	df2232_idr_remove_all(struct idr *);
void	df2232_idr_destroy(struct idr *);
void	*df2232_idr_find(struct idr *, int);
int	df2232_idr_for_each(struct idr *, int (*)(int, void *, void *),
			   void *);
void	*df2232_idr_replace(struct idr *, void *, int);
void	df2232_idr_init(struct idr *);
/* tentative definition: the real initializer comes from MALLOC_DEFINE
 * inside linux_idr.c (renamed to df2232_M_IDR) */
struct malloc_type df2232_M_IDR[1];

#define idr_pre_get		df2232_idr_pre_get
#define idr_get_new_above	df2232_idr_get_new_above
#define idr_get_new		df2232_idr_get_new
#define idr_alloc		df2232_idr_alloc
#define idr_remove		df2232_idr_remove
#define idr_remove_all		df2232_idr_remove_all
#define idr_destroy		df2232_idr_destroy
#define idr_find		df2232_idr_find
#define idr_for_each		df2232_idr_for_each
#define idr_replace		df2232_idr_replace
#define idr_init		df2232_idr_init
#define M_IDR			df2232_M_IDR

/*
 * kmalloc observation wrapper: capture the exact request size at the
 * allocator boundary, then pass through to the REAL allocator unchanged.
 */
static void *
df2232_kmalloc_obs(unsigned long size, int flags)
{
	df2232_last_req = size;
	kprintf("[DF2232-B] ALLOCATOR BOUNDARY: kmalloc(%lu bytes [0x%lx], "
		"M_IDR, flags 0x%x) -> real allocator\n", size, size, flags);
	return _kmalloc(size, df2232_M_IDR, flags);
}

#undef kmalloc
#define kmalloc(size, type, flags)	df2232_kmalloc_obs((size), (flags))

#include "linux_idr.c"		/* byte-identical, sha256-verified */

/*
 * Restore the plain names: from here on, idr_alloc(...) etc. refer to the
 * RUNNING KERNEL's own exported KPI (the shipped binary).
 */
#undef idr_pre_get
#undef idr_get_new_above
#undef idr_get_new
#undef idr_alloc
#undef idr_remove
#undef idr_remove_all
#undef idr_destroy
#undef idr_find
#undef idr_for_each
#undef idr_replace
#undef idr_init
#undef M_IDR
#undef kmalloc

/*=====================================================================
 * Phase A: normal-path control (verbatim copy)
 *=====================================================================*/
static int
df2232_phase_a(int use_kernel)
{
	struct idr *idp = use_kernel ? &df2232_idr_k : &df2232_idr_v;
	void *p;
	int id = -1, r, fails = 0;
	const char *tag = use_kernel ? "[DF2232-A2]" : "[DF2232-A1]";

	kprintf("%s === control phase (%s) ===\n", tag,
		use_kernel ? "kernel KPI (shipped binary)"
			   : "verbatim KLD copy");
	kprintf("%s sizeof(struct idr_node)=%lu\n", tag,
		(unsigned long)sizeof(struct idr_node));

#define DF2232_CALL(ret, expr_kernel, expr_verbatim)			\
	do {								\
		if (use_kernel)						\
			ret = (expr_kernel);				\
		else							\
			ret = (expr_verbatim);				\
	} while (0)

	kprintf("%s idr_init()\n", tag);
	if (use_kernel)
		idr_init(idp);
	else
		df2232_idr_init(idp);
	kprintf("%s idr_count=%d (expect 255)\n", tag, idp->idr_count);
	if (idp->idr_count != 255)
		fails++;

	/* grow to want=1 -> no growth needed, id should be 1 */
	DF2232_CALL(r, idr_alloc(idp, (void *)0xAA, 1, 0, 0),
		     df2232_idr_alloc(idp, (void *)0xAA, 1, 0, 0));
	kprintf("%s idr_alloc(start=1) -> %d (expect 1), idr_count=%d "
		"(expect 255)\n", tag, r, idp->idr_count);
	if (r != 1 || idp->idr_count != 255)
		fails++;

	/* grow to want=1000 -> count 255->1023 via idr_grow loop
	 * (nf stops at the first 2^n-1 EXCEEDING want: 1023 > 1000) */
	DF2232_CALL(r, idr_alloc(idp, (void *)0xBB, 1000, 0, 0),
		     df2232_idr_alloc(idp, (void *)0xBB, 1000, 0, 0));
	kprintf("%s idr_alloc(start=1000) -> %d (expect 1000), "
		"idr_count=%d (expect 1023)\n", tag, r, idp->idr_count);
	if (r != 1000 || idp->idr_count != 1023)
		fails++;

	DF2232_CALL(p, idr_find(idp, 1), df2232_idr_find(idp, 1));
	kprintf("%s idr_find(1) -> %p (expect 0xaa)\n", tag, p);
	DF2232_CALL(p, idr_find(idp, 1000), df2232_idr_find(idp, 1000));
	kprintf("%s idr_find(1000) -> %p (expect 0xbb)\n", tag, p);
	DF2232_CALL(p, idr_find(idp, 999), df2232_idr_find(idp, 999));
	kprintf("%s idr_find(999) -> %p (expect NULL)\n", tag, p);
	if (p != NULL)
		fails++;

	/* the documented pre_get + get_new_above(sid) protocol: the FIRST
	 * get_new_above with a fresh sid >= idr_maxwant returns -EAGAIN and
	 * poisons idr_maxwant -- the exact state the finding's INT_MAX
	 * attack uses -- so loop like the in-tree callers do */
	{
		int iter;

		for (iter = 1; iter <= 5; iter++) {
			DF2232_CALL(r, idr_pre_get(idp, 0),
				     df2232_idr_pre_get(idp, 0));
			kprintf("%s iter %d: idr_pre_get -> %d "
				"(expect 1)\n", tag, iter, r);
			if (r != 1)
				fails++;
			DF2232_CALL(r,
				     idr_get_new_above(idp, (void *)0xCC,
						       500, &id),
				     df2232_idr_get_new_above(idp,
							      (void *)0xCC,
							      500, &id));
			kprintf("%s iter %d: idr_get_new_above(sid=500) -> %d "
				"(EAGAIN=%d expected on iter 1), id=%d\n",
				tag, iter, r, -EAGAIN, id);
			if (r == -EAGAIN)
				continue;
			break;
		}
		kprintf("%s get_new_above settled: ret=%d id=%d "
			"(expect 0, 500) after %d iteration(s), "
			"idr_maxwant=%d (expect 500)\n",
			tag, r, id, iter, idp->idr_maxwant);
		if (r != 0 || id != 500)
			fails++;
	}

	/* remove + re-alloc reuses the slot */
	DF2232_CALL(p, idr_remove(idp, 1000), df2232_idr_remove(idp, 1000));
	kprintf("%s idr_remove(1000) -> %p (expect 0xbb)\n", tag, p);
	if (p != (void *)0xBB)
		fails++;
	DF2232_CALL(r, idr_alloc(idp, (void *)0xDD, 1000, 0, 0),
		     df2232_idr_alloc(idp, (void *)0xDD, 1000, 0, 0));
	kprintf("%s idr_alloc(start=1000) again -> %d (expect 1000)\n",
		tag, r);
	if (r != 1000)
		fails++;

	kprintf("%s idr_nexpands=%d\n", tag, idp->idr_nexpands);
	if (use_kernel)
		idr_destroy(idp);
	else
		df2232_idr_destroy(idp);
	kprintf("%s idr_destroy done\n", tag);
	kprintf("%s === control phase %s (fails=%d) ===\n", tag,
		fails ? "FAILED" : "PASSED", fails);
	/* always return 0: the console log is the evidence */
	return (0);
}

/*=====================================================================
 * Phase B: giant kmalloc (want = 1<<29) via the verbatim copy
 *=====================================================================*/
static int
df2232_phase_b(void)
{
	int r;

	kprintf("[DF2232-B] === giant-kmalloc phase (verbatim copy) ===\n");
	kprintf("[DF2232-B] plan: idr_alloc(start=%d, end=0)\n", 1 << 29);
	kprintf("[DF2232-B] expected idr_grow loop result: nf = 2^30-1 = %d\n",
		(1 << 30) - 1);
	kprintf("[DF2232-B] expected kmalloc request: %d * %lu = %lu bytes "
		"(~16 GiB, M_ZERO|M_WAITOK)\n", (1 << 30) - 1,
		(unsigned long)sizeof(struct idr_node),
		(unsigned long)((1 << 30) - 1) *
		sizeof(struct idr_node));
	df2232_arm_watchdog("B: idr_grow kmalloc(~16GB, M_WAITOK) never "
			    "returned (vm_wait wedge / OOM)");
	kprintf("[DF2232-B] idr_init\n");
	df2232_idr_init(&df2232_idr_v);
	kprintf("[DF2232-B] entering idr_alloc ... (wedge expected on "
		"unfixed kernel; watchdog %d s)\n", DF2232_WD_SECONDS);
	df2232_danger_ran = 1;
	r = df2232_idr_alloc(&df2232_idr_v, (void *)0xB0B, 1 << 29, 0, 0);
	kprintf("[DF2232-B] idr_alloc RETURNED %d (unexpected on unfixed "
		"kernel)\n", r);
	if (r < 0) {
		kprintf("[DF2232-B] NOTE: negative errno return, no giant "
			"kmalloc reached the allocator (last request was "
			"%lu bytes)\n", df2232_last_req);
		df2232_idr_destroy(&df2232_idr_v);
	}
	df2232_disarm_watchdog();
	return (0);
}

/*=====================================================================
 * Phases C1/C2: want == INT_MAX infinite spin holding the idr token
 *=====================================================================*/
static void
df2232_thread_a(void *arg __unused)
{
	int error, id = -1;

	kprintf("[DF2232-C] A: running on cpu %d (%s)\n", mycpu->gd_cpuid,
		df2232_use_kernel ? "kernel KPI" : "verbatim copy");
	kprintf("[DF2232-C] A: calling idr_get_new_above(sid=%d) -- "
		"expect -EAGAIN and idr_maxwant poisoned to INT_MAX\n",
		INT_MAX);
	if (df2232_use_kernel)
		error = idr_get_new_above(&df2232_idr_c, (void *)0xC1,
					  INT_MAX, &id);
	else
		error = df2232_idr_get_new_above(&df2232_idr_c,
						 (void *)0xC1, INT_MAX, &id);
	df2232_a_stage = 1;
	kprintf("[DF2232-C] A: get_new_above returned %d (expect %d), "
		"id=%d, idr_maxwant=%d (expect %d)\n",
		error, -EAGAIN, id, df2232_idr_c.idr_maxwant, INT_MAX);
	kprintf("[DF2232-C] A: entering idr_pre_get -- expect INFINITE "
		"SPIN in idr_grow: nf=2*nf+1 wraps 2*INT_MAX+1 to -1, "
		"-1 <= INT_MAX forever, idr token held\n");
	df2232_a_stage = 2;
	if (df2232_use_kernel)
		error = idr_pre_get(&df2232_idr_c, 0);
	else
		error = df2232_idr_pre_get(&df2232_idr_c, 0);
	df2232_a_done = 1;
	kprintf("[DF2232-C] A: pre_get RETURNED %d (unexpected on unfixed "
		"kernel)\n", error);
	kthread_exit();
}

static void
df2232_thread_b(void *arg __unused)
{
	int waited;

	kprintf("[DF2232-C] B: running on cpu %d\n", mycpu->gd_cpuid);
	/* prove the token is acquirable BEFORE A enters pre_get */
	kprintf("[DF2232-C] B: pre-spin token check: idr_remove(0) ...\n");
	if (df2232_use_kernel)
		idr_remove(&df2232_idr_c, 0);
	else
		df2232_idr_remove(&df2232_idr_c, 0);
	kprintf("[DF2232-C] B: pre-spin token check returned (token was "
		"free)\n");

	for (waited = 0; waited < 100 && df2232_a_stage != 2; waited++)
		tsleep(&df2232_a_stage, 0, "df2232b", 1);
	kprintf("[DF2232-C] B: A entered pre_get (a_stage=%d); calling "
		"idr_remove(0) on the SAME idr -- expect NEVER to return "
		"(lwkt token held forever by spinning A)\n",
		df2232_a_stage);
	df2232_b_started = 1;
	if (df2232_use_kernel)
		idr_remove(&df2232_idr_c, 0);
	else
		df2232_idr_remove(&df2232_idr_c, 0);
	df2232_b_acquired = 1;
	kprintf("[DF2232-C] B: idr_remove RETURNED (unexpected on unfixed "
		"kernel)\n");
	kthread_exit();
}

static int
df2232_phase_c(int use_kernel)
{
	int waited;

	df2232_use_kernel = use_kernel;
	df2232_danger_ran = 1;
	kprintf("[DF2232-C] === spin phase (%s) ===\n",
		use_kernel ? "kernel KPI (shipped binary)" :
			     "verbatim KLD copy");
	kprintf("[DF2232-C] want=INT_MAX=%d: growth loop 2*INT_MAX+1 "
		"overflows signed int to -1; -1 <= INT_MAX -> infinite "
		"loop holding lwkt idr token\n", INT_MAX);
	df2232_arm_watchdog("C: idr_grow growth-loop infinite spin "
			    "(want=INT_MAX), idr token held forever; "
			    "thread B denied token (livelock)");
	if (use_kernel)
		idr_init(&df2232_idr_c);
	else
		df2232_idr_init(&df2232_idr_c);
	kprintf("[DF2232-C] idr_count=%d\n", df2232_idr_c.idr_count);

	kthread_create_cpu(df2232_thread_a, NULL, NULL,
			   (mycpu->gd_cpuid + 1) % ncpus, "df2232_a");
	for (waited = 0; waited < 100 && df2232_a_stage == 0; waited++)
		tsleep(&df2232_a_stage, 0, "df2232c", 1);
	kthread_create_cpu(df2232_thread_b, NULL, NULL,
			   (mycpu->gd_cpuid + 2) % ncpus, "df2232_b");

	/* wait for both threads to finish (fixed kernel) or WD panic */
	for (waited = 0; waited < 30 * 10; waited++) {
		if (df2232_a_done && df2232_b_acquired)
			break;
		tsleep(&df2232_a_done, 0, "df2232w", hz / 10);
	}
	kprintf("[DF2232-C] phase finished: a_done=%d b_acquired=%d\n",
		df2232_a_done, df2232_b_acquired);
	if (df2232_a_done && df2232_b_acquired) {
		kprintf("[DF2232-C] no spin, no livelock (fixed kernel "
			"behavior)\n");
		if (use_kernel)
			idr_destroy(&df2232_idr_c);
		else
			df2232_idr_destroy(&df2232_idr_c);
		df2232_disarm_watchdog();
	}
	return (0);
}

/*=====================================================================
 * sysctl trigger
 *=====================================================================*/
static int	df2232_last_phase;

static int
sysctl_df2232_trigger(SYSCTL_HANDLER_ARGS)
{
	int error, phase = 0;

	error = sysctl_handle_int(oidp, &phase, 0, req);
	if (error != 0 || req->newptr == NULL)
		return (error);

	switch (phase) {
	case 1:
	case 2:
		error = df2232_phase_a(phase - 1);
		break;
	case 3:
		error = df2232_phase_b();
		break;
	case 4:
	case 5:
		error = df2232_phase_c(phase - 4);
		break;
	default:
		return (EINVAL);
	}
	df2232_last_phase = phase;
	return (error);
}

SYSCTL_PROC(_debug, OID_AUTO, df2232, CTLTYPE_INT | CTLFLAG_RW,
	    0, 0, sysctl_df2232_trigger, "I",
	    "DF-2232: 1=control(verbatim) 2=control(kernel) 3=giant-kmalloc "
	    "4=spin(verbatim) 5=spin(kernel)");

static int
df2232_modevent(module_t mod __unused, int type, void *data __unused)
{
	switch (type) {
	case MOD_LOAD:
		kprintf("[DF2232] rig loaded: verbatim linux_idr.c + kernel "
			"KPI, watchdog %ds\n", DF2232_WD_SECONDS);
		return (0);
	case MOD_UNLOAD:
		/* Stay resident once a crash-phase ran: tearing down the
		 * armed systimers and threads racing kldunload is the
		 * rig's own hazard, not the finding's.  EBUSY is clean. */
		if (df2232_armed || df2232_threads_live || df2232_danger_ran)
			return (EBUSY);
		kprintf("[DF2232] rig unloaded\n");
		return (0);
	default:
		return (EOPNOTSUPP);
	}
}

static moduledata_t df2232_mod = {
	"df2232",
	df2232_modevent,
	NULL
};

DECLARE_MODULE(df2232, df2232_mod, SI_SUB_DRIVERS, SI_ORDER_ANY);
MODULE_VERSION(df2232, 1);