/*
 * DF-5321 / DF-5322 verification harness (sys/bus/firewire/firewire.c)
 *
 * DF-5321 - split-transaction completion race: firewire_xfer_timeout()
 *     (firewire.c:320-353) and the fw_rcv() response path
 *     (firewire.c:1757-1799) can BOTH complete the same struct fw_xfer:
 *       * fw_rcv() looks the xfer up with fw_tl2xfer() (firewire.c:1757),
 *         then copies into xfer->recv.* with fw_rcv_copy() (firewire.c:1779)
 *         and only afterwards stores state = FWXF_RCVD (firewire.c:1786)
 *         and calls fw_xfer_done() (firewire.c:1789).  It never removes
 *         the tlabel and holds no reference.
 *       * firewire_xfer_timeout() -- a 10 Hz callout watchdog on another
 *         CPU; its critical section does not exclude other CPUs -- skips
 *         only state == FWXF_START (firewire.c:341), so it removes the
 *         tlabel (firewire.c:348) and calls fw_xfer_done() (firewire.c:349)
 *         on an xfer that fw_rcv() is concurrently completing, or that
 *         fw_rcv() already completed (state FWXF_RCVD) but whose owner
 *         has not freed yet.
 *     Consequences (both demonstrated with the PRODUCTION code):
 *       (a) act.hand == fw_asy_callback_free (production users in this
 *           file: fw_phy_config firewire.c:1078, fw_bmr's fwmem_write_quad
 *           firewire.c:2169-2170): fw_xfer_free() runs TWICE on the same
 *           object -> use-after-free + double kfree -> panic / duplicate
 *           slab pointers.
 *       (b) act.hand == fw_asy_callback with the /dev/fw FW_ASYREQ waiter
 *           (fwdev.c tsleep -> fwdev.c fw_xfer_free_buf): the timeout's
 *           wakeup() releases the waiter, which frees the xfer while the
 *           receiving context is still inside fw_rcv_copy()/storing
 *           xfer->state -> UAF WRITE into a freed (possibly reused)
 *           fw_xfer.
 *
 * DF-5322 - struct tlabel leak on every split-transaction timeout:
 *     firewire_xfer_timeout() unlinks the tlabel with STAILQ_REMOVE_HEAD
 *     (firewire.c:348) but never kfrees it.  The only kfree of a tlabel is
 *     in fw_tl_free() (firewire.c:900), which frees only entries it still
 *     finds in the list -- once the timeout removed it, the owner's
 *     fw_xfer_free() -> fw_xfer_unload() -> fw_tl_free() finds nothing and
 *     the struct tlabel (24 bytes on x86_64) is leaked for EVERY timed-out
 *     transaction.
 *
 * The harness compiles the PRODUCTION sys/bus/firewire/firewire.c into the
 * module (via #include) so the production fw_rcv(), fw_tl2xfer(),
 * firewire_xfer_timeout(), fw_get_tlabel(), fw_xfer_free*() run unchanged.
 * Synthetic state is only the FireWire bus itself (a struct firewire_comm
 * with the tcode table fwohci.c installs at fwohci.c:661, fc->bdev =
 * root_bus) and the "wire" response packet -- no FireWire controller
 * exists on the QEMU guest (hardware-gated finding).  firewire.c's
 * globals are -D-renamed in the Makefile so they do not clash with the
 * copy compiled into the guest kernel (X86_64_GENERIC has
 * `device firewire`).
 *
 * Modes (kern.df5321_mode):
 *   0 = DF-5322  tlabel leak measurement (M_FW inuse/memuse delta)
 *   1 = DF-5321  instrumented race: counting handler; counts double
 *               completions, act.hand calls after free, and stale stores
 *               landing in a re-allocated holder fw_xfer
 *   2 = DF-5321  production handler fw_asy_callback_free: the loser of
 *               the race double-frees the fw_xfer -> expect kernel panic
 *
 * Fire with: sysctl kern.df5321_fire=1
 */
#include <sys/param.h>
#include <sys/systm.h>
#include <sys/types.h>
#include <sys/kernel.h>
#include <sys/module.h>
#include <sys/malloc.h>
#include <sys/sysctl.h>
#include <sys/bus.h>
#include <sys/uio.h>
#include <sys/time.h>
#include <sys/thread.h>
#include <sys/thread2.h>
#include <machine/atomic.h>

/* PRODUCTION source: its statics (fw_get_tlabel, fw_tl_free, fw_tl2xfer,
 * firewire_xfer_timeout) become callable in this translation unit.
 *
 * MODULE_VERSION(firewire, 1) at firewire.c:2208 registers the module
 * interface version metadata "firewire" (the #name stringify ignores the
 * -D renames), which the kernel-resident copy already provides -- EEXIST
 * on kldload.  Neutralize just that registration; everything else in the
 * production file is compiled unchanged. */
#undef MODULE_VERSION
#define MODULE_VERSION(name, version)	struct __hack
#include <bus/firewire/firewire.c>
#undef MODULE_VERSION

struct att_rec {
	volatile int hand_calls;	/* act.hand invocations */
	volatile int freed;		/* fw_xfer_free_buf done (cmpset-claimed) */
	volatile int holder_set;
	int gen;
	int dst;
	volatile int last_hand_seq;
	volatile int free_seq;
	struct fw_xfer *holder;
};

static int	df5321_mode = 1;
static int	df5321_iters = 30000;

SYSCTL_INT(_kern, OID_AUTO, df5321_mode, CTLFLAG_RW, &df5321_mode, 0,
    "harness case: 0=DF-5322 tlabel leak, 1=DF-5321 instrumented race, "
    "2=DF-5321 production double-free");
SYSCTL_INT(_kern, OID_AUTO, df5321_iters, CTLFLAG_RW, &df5321_iters, 0,
    "attempts for the harness case");

/* synthetic bus state */

/* exact copy of fwohci.c:85-102 (tinfo[]), installed as fc->tcode by
 * fwohci.c:661 (sc->fc.tcode = tinfo) */
static struct tcode_info dfw_tinfo[16] = {
/* 0 WREQQ  */ {16,	FWTI_REQ | FWTI_TLABEL},
/* 1 WREQB  */ {16,	FWTI_REQ | FWTI_TLABEL | FWTI_BLOCK_ASY},
/* 2 WRES   */ {12,	FWTI_RES},
/* 3 XXX    */ { 0,	0},
/* 4 RREQQ  */ {12,	FWTI_REQ | FWTI_TLABEL},
/* 5 RREQB  */ {16,	FWTI_REQ | FWTI_TLABEL},
/* 6 RRESQ  */ {16,	FWTI_RES},
/* 7 RRESB  */ {16,	FWTI_RES | FWTI_BLOCK_ASY},
/* 8 CYCS   */ { 0,	0},
/* 9 LREQ   */ {16,	FWTI_REQ | FWTI_TLABEL | FWTI_BLOCK_ASY},
/* a STREAM */ { 4,	FWTI_REQ | FWTI_BLOCK_STR},
/* b LRES   */ {16,	FWTI_RES | FWTI_BLOCK_ASY},
/* c XXX    */ { 0,	0},
/* d XXX    */ { 0, 	0},
/* e PHY    */ {12,	FWTI_REQ},
/* f XXX    */ { 0,	0}
};

static struct firewire_comm *g_fc;

static volatile int g_stop;		/* threads: please exit */
static volatile int g_live;		/* threads still alive */
static volatile int g_newatt;		/* new attempt published */
static struct fw_xfer *g_att_xfer;	/* current attempt's xfer */
static struct att_rec *g_att_rec;
static volatile int g_seq;
static volatile int g_gen;

static volatile int g_attempts;		/* attempts run (race modes) */
static volatile int g_hits;		/* >= 2 completions for one xfer */
static volatile int g_hand_after_free;	/* act.hand called after free */
static volatile int g_uaf_write;	/* stale stores hit holder object */

static u_char pkt[4 * 4 + 65536];

static size_t
mfw_stat(int what)
{
	struct malloc_type *t = M_FW;	/* renamed DFW_M_FW in this TU */
	size_t s = 0;
	int i;

	if (t->ks_use != NULL) {
		for (i = 0; i < SMP_MAXCPU; i++) {
			if (what == 0)
				s += t->ks_use[i].inuse;
			else
				s += t->ks_use[i].memuse;
		}
	}
	return (s);
}

static void
prewarm_zones(void)
{
	void *p[32];
	int i;

	/* populate the tlabel/xfer size classes so the production
	 * fw_get_tlabel() kmalloc(M_WAITOK) inside its critical section
	 * (the known DF-1021 site) does not have to sleep here */
	for (i = 0; i < 32; i++)
		p[i] = kmalloc(sizeof(struct tlabel), M_FW, M_INTWAIT);
	for (i = 0; i < 32; i++)
		kfree(p[i], M_FW);
	for (i = 0; i < 32; i++)
		p[i] = kmalloc(sizeof(struct fw_xfer), M_FWXFER, M_INTWAIT);
	for (i = 0; i < 32; i++)
		kfree(p[i], M_FWXFER);
}

static void
init_fc(void)
{
	int i;

	g_fc = kmalloc(sizeof(*g_fc), M_TEMP, M_INTWAIT | M_ZERO);
	for (i = 0; i < 0x40; i++)
		STAILQ_INIT(&g_fc->tlabels[i]);
	STAILQ_INIT(&g_fc->binds);
	STAILQ_INIT(&g_fc->devices);
	g_fc->tcode = dfw_tinfo;	/* as fwohci.c:661 installs it */
	g_fc->bdev = root_bus;		/* device_printf(fc->bdev,..) target */
	g_fc->nodeid = 1;
	g_fc->status = FWBUSEXPLORE;
	g_fc->max_asyretry = FW_MAXASYRTY;
}

/* ------------------------------------------------------------------ */
/* DF-5322: tlabel leak                                                */

static void
run_leak(int n)
{
	struct fw_xfer *xfer;
	size_t in0, in1, mu0, mu1;
	int i, tl, leaked;

	prewarm_zones();
	init_fc();

	in0 = mfw_stat(0);
	mu0 = mfw_stat(1);

	for (i = 0; i < n; i++) {
		xfer = fw_xfer_alloc(M_FWXFER);		/* production */
		if (xfer == NULL)
			continue;
		xfer->fc = g_fc;
		/* production fwdev FW_ASYREQ style completion handler */
		xfer->act.hand = fw_asy_callback;	/* production */
		xfer->send.hdr.mode.hdr.dst = 0x111;
		xfer->state = FWXF_SENT;		/* fwohci_txd() DMA'd it */
		microtime(&xfer->tv);
		xfer->tv.tv_sec -= 2;			/* past the 200ms split */
		tl = fw_get_tlabel(g_fc, xfer);		/* production: inserts */
		if (tl < 0) {
			fw_xfer_free(xfer);
			continue;
		}
		xfer->tl = tl;				/* as fw_asyreq:243 */
		/* production 10 Hz watchdog: removes the tlabel WITHOUT
		 * kfree (firewire.c:348) and completes the xfer */
		firewire_xfer_timeout(g_fc);
		/* production user-side free (fwdev.c out: label): the
		 * fw_tl_free() inside finds nothing -> tlabel leaked */
		fw_xfer_free(xfer);
	}

	in1 = mfw_stat(0);
	mu1 = mfw_stat(1);
	leaked = (int)(in1 - in0);

	kprintf("DF5322: %d split-transaction timeouts\n", n);
	kprintf("DF5322: M_FW inuse delta = %d objects, memuse delta = %zd"
	    " bytes (each leaked struct tlabel = %zu bytes)\n",
	    leaked, (ssize_t)(mu1 - mu0), sizeof(struct tlabel));
	if (leaked >= n / 2)
		kprintf("DF5322: LEAK REPRODUCED: ~1 tlabel leaked per"
		    " timeout (firewire.c:348 unlinks without kfree)\n");
	else
		kprintf("DF5322: no significant leak (fixed kernel?)\n");

	kfree(g_fc, M_TEMP);
	g_fc = NULL;
}

/* ------------------------------------------------------------------ */
/* DF-5321: the completion race                                        */

static void
count_hand(struct fw_xfer *x)
{
	struct att_rec *rec = (struct att_rec *)x->sc;

	if (rec != NULL) {
		atomic_add_int(&rec->hand_calls, 1);
		rec->last_hand_seq = atomic_fetchadd_int(&g_seq, 1) + 1;
	}
	/* production fw_asy_callback behaviour (firewire.c:255-259):
	 * wakeup the FW_ASYREQ waiter (fwdev.c tsleep) */
	wakeup(x);
}

/*
 * Receiving context: fwohci_arcv() -> fw_rcv() on an interrupt/serializing
 * thread; a response arrives for an outstanding xfer whose 200 ms split
 * timeout is already due.
 */
static void
trcv_thread(void *arg)
{
	struct fw_rcv_buf rb;
	struct iovec vec[2];
	struct att_rec *rec;
	struct fw_xfer *xfer, *holder;
	u_int32_t *ld;
	int n = df5321_iters, tl, k, mode = df5321_mode;
	struct timespec deadline;
	int printed_deadline = 0;

	/* hard wall-clock bound so the harness can never wedge the guest */
	nanouptime(&deadline);
	deadline.tv_sec += 90;

	while (n-- > 0 && !g_stop) {
		if ((n & 0xff) == 0) {
			struct timespec now;
			nanouptime(&now);
			if (timespeccmp(&now, &deadline, >=)) {
				if (!printed_deadline) {
					kprintf("DF5321: trcv wall-clock"
					    " deadline hit\n");
					printed_deadline = 1;
				}
				break;
			}
		}
		xfer = fw_xfer_alloc(M_FWXFER);		/* production */
		if (xfer == NULL)
			continue;
		rec = kmalloc(sizeof(*rec), M_TEMP, M_INTWAIT | M_ZERO);
		rec->gen = atomic_fetchadd_int(&g_gen, 1) + 1;
		rec->dst = 0x100 + (rec->gen % 0x400);	/* unique per attempt */

		xfer->fc = g_fc;
		xfer->send.hdr.mode.hdr.dst = rec->dst;
		xfer->state = FWXF_SENT;		/* DMA'd, awaiting resp */
		microtime(&xfer->tv);
		xfer->tv.tv_sec -= 2;			/* split timeout due */
		xfer->recv.payload = kmalloc(65536, M_FWXFER, M_INTWAIT);
		xfer->recv.pay_len = 65535;
		xfer->recv.spd = 3;
		xfer->sc = (caddr_t)rec;	/* count_hand must find it even
						 * if the watchdog completes
						 * the xfer immediately */
		if (mode == 2)
			xfer->act.hand = fw_asy_callback_free; /* production */
		else
			xfer->act.hand = count_hand;

		tl = fw_get_tlabel(g_fc, xfer);		/* production: inserts */
		if (tl < 0) {
			kfree(xfer->recv.payload, M_FWXFER);
			fw_xfer_free(xfer);
			kfree(rec, M_TEMP);
			continue;
		}
		xfer->tl = tl;

		/*
		 * Build the RRESB response "from the wire" for this xfer.
		 * Little-endian bitfield layout (firewire.h COMMON_HDR /
		 * BIT16x2):
		 *   ld[0] = dst | tlrt<<8 | tcode(7)<<4 | pri
		 *   ld[1] = src << 16
		 *   ld[2] = 0		(rtcode = RESP_CMP)
		 *   ld[3] = 65535 << 16	(rresb.len = plen)
		 */
		ld = (u_int32_t *)pkt;
		ld[0] = (0xffc0u << 16) | ((u_int32_t)(tl << 2) << 8) |
		    (FWTCODE_RRESB << 4);
		ld[1] = (u_int32_t)rec->dst << 16;
		ld[2] = 0;
		ld[3] = 65535u << 16;

		vec[0].iov_base = pkt;
		vec[0].iov_len = 16 + 32768;
		vec[1].iov_base = pkt + 16 + 32768;
		vec[1].iov_len = 32768;
		rb.fc = g_fc;
		rb.xfer = NULL;
		rb.vec = vec;
		rb.nvec = 2;
		rb.spd = 3;

		xfer->sc = (caddr_t)rec;	/* (set above, before tlabel
						 * insertion) */
		if (mode == 1) {
			/* publish for the FW_ASYREQ waiter thread */
			g_att_rec = rec;
			g_att_xfer = xfer;
			atomic_set_int(&g_newatt, 1);
		}

		fw_rcv(&rb);				/* production */

		if (mode == 1) {
			/* wait for the waiter thread's role (bounded) */
			for (k = 0; k < 200 &&
			    (!rec->freed || !rec->holder_set); k++)
				tsleep(rec, 0, "dfwtr", 1);
			/* if nobody freed it, the owner does (production) */
			if (atomic_cmpset_int(&rec->freed, 0, 1))
				fw_xfer_free_buf(xfer);
			holder = rec->holder;
			if (holder != NULL) {
				/*
				 * holder is M_ZERO; the receiving context's
				 * stale stores (xfer->state = FWXF_RCVD at
				 * firewire.c:1786, xfer->resp at :1781-1783,
				 * recv.pay_len at :1723) landing AFTER the
				 * free show up as non-zero fields here.
				 */
				if (holder->state != FWXF_INIT ||
				    holder->resp != 0 ||
				    holder->recv.pay_len != 0) {
					atomic_add_int(&g_uaf_write, 1);
				}
				fw_xfer_free(holder);
			}
			if (rec->hand_calls >= 2)
				atomic_add_int(&g_hits, 1);
			if (rec->free_seq != 0 &&
			    rec->last_hand_seq > rec->free_seq)
				atomic_add_int(&g_hand_after_free, 1);
			atomic_add_int(&g_attempts, 1);
			kfree(rec, M_TEMP);
		} else {
			/*
			 * mode 2: the xfer was already freed by whichever
			 * completion ran first (production
			 * fw_asy_callback_free); nothing to clean up here.
			 */
			atomic_add_int(&g_attempts, 1);
		}
	}

	g_stop = 1;	/* receiver done: release the other threads */
	atomic_subtract_int(&g_live, 1);
	wakeup(&g_live);
	lwkt_exit();
}

/*
 * Watchdog context: firewire_watchdog() -> firewire_xfer_timeout() at
 * 10 Hz on its own CPU.
 */
static void
ttmo_thread(void *arg)
{
	int i = 0;

	while (!g_stop) {
		firewire_xfer_timeout(g_fc);	/* production */
		if ((++i & 63) == 0)
			tsleep(&g_stop, 0, "dfwtmo", 1);
		else
			lwkt_yield();
	}
	atomic_subtract_int(&g_live, 1);
	wakeup(&g_live);
	lwkt_exit();
}

/*
 * FW_ASYREQ waiter context: fwdev.c's ioctl thread -- tsleep(xfer) until
 * act.hand wakes it, then fw_xfer_free_buf(xfer).
 */
static void
tuser_thread(void *arg)
{
	struct fw_xfer *x;
	struct att_rec *r;
	int k;

	while (!g_stop) {
		if (atomic_swap_int(&g_newatt, 0) == 1) {
			x = g_att_xfer;
			r = g_att_rec;
			if (x != NULL && r != NULL) {
				/* fwdev.c: tsleep(xfer, FWPRI, "asyreq", hz) */
				for (k = 0; k < 10 && r->hand_calls == 0; k++)
					tsleep(x, 0, "dfwtu", 1);
				/* fwdev.c out: fw_xfer_free_buf(xfer) */
				if (atomic_cmpset_int(&r->freed, 0, 1))
					fw_xfer_free_buf(x);
				r->free_seq =
				    atomic_fetchadd_int(&g_seq, 1) + 1;
				/* keep the freed slot stable & zeroed so
				 * stale stores by the receiving context are
				 * visible */
				r->holder = fw_xfer_alloc(M_FWXFER);
				r->holder_set = 1;
				wakeup(r);
			}
		} else {
			tsleep(&g_newatt, 0, "dfwtui", 1);
		}
	}
	atomic_subtract_int(&g_live, 1);
	wakeup(&g_live);
	lwkt_exit();
}

static void
run_case(void)
{

	g_stop = 0;
	g_live = 0;
	g_newatt = 0;
	g_seq = 0;
	g_attempts = 0;
	g_hits = 0;
	g_hand_after_free = 0;
	g_uaf_write = 0;

	if (df5321_mode == 0) {
		run_leak(df5321_iters);
		return;
	}

	prewarm_zones();
	init_fc();

	kprintf("DF5321: mode %d starting (%d attempts, PRODUCTION firewire.c"
	    " completion paths)\n", df5321_mode, df5321_iters);

	g_live = 2;
	lwkt_create(trcv_thread, NULL, NULL, NULL, 0, 1, "dfw_rcv");
	lwkt_create(ttmo_thread, NULL, NULL, NULL, 0, 2, "dfw_tmo");
	if (df5321_mode == 1) {
		g_live = 3;
		lwkt_create(tuser_thread, NULL, NULL, NULL, 0, 3, "dfw_usr");
	}

	while (g_live)
		tsleep(&g_live, 0, "dfwwait", hz / 4);

	if (df5321_mode == 1) {
		kprintf("DF5321: mode1 done: attempts=%d double-completions=%d"
		    " act.hand-after-free=%d"
		    " stale-stores-into-reused-fw_xfer=%d\n",
		    g_attempts, g_hits, g_hand_after_free, g_uaf_write);
	} else {
		kprintf("DF5321: mode2 SURVIVED %d attempts (baseline expects"
		    " a double-free panic; a fixed kernel expects this)\n",
		    g_attempts);
	}

	kfree(g_fc, M_TEMP);
	g_fc = NULL;
}

static int
df5321_fire_sysctl(SYSCTL_HANDLER_ARGS)
{
	int error, v = 0;

	error = SYSCTL_OUT(req, &v, sizeof(v));
	if (error != 0 || req->newptr == NULL)
		return (error);
	error = SYSCTL_IN(req, &v, sizeof(v));
	if (error != 0)
		return (error);
	if (v == 1)
		run_case();
	return (0);
}

SYSCTL_PROC(_kern, OID_AUTO, df5321_fire, CTLTYPE_INT | CTLFLAG_RW,
    0, 0, df5321_fire_sysctl, "I",
    "write 1 to fire the current harness case");