/*
 * DF-0712 deterministic code-level harness.
 *
 * The live `ieee80211_alloc_countryie` path requires a WiFi radio (it is
 * called from ieee80211_add_countryie while assembling a beacon/probe-resp on
 * a hostap vap). This audit guest has NO WiFi hardware (only vtnet0/lo0, no
 * wlan/80211 modules loaded), so the live path is unreachable here.
 *
 * This harness faithfully replicates the EXACT allocation + IE-assembly logic
 * of sys/netproto/802_11/wlan/ieee80211_regdomain.c: ieee80211_alloc_countryie
 * (lines 248-326), byte-for-byte, including the struct layouts from
 * sys/netproto/802_11/ieee80211.h:961 and ieee80211_var.h:109. It proves the
 * off-by-3 heap overflow deterministically using a canary-guarded allocation.
 *
 * Build:  cc -O2 -o harness harness.c
 * Run:    ./harness
 */
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>

/* --- Replicate kernel struct/constant definitions exactly --- */

#define IEEE80211_COUNTRY_MAX_BANDS 84		/* ieee80211.h:972 */
struct ieee80211_country_ie {			/* ieee80211.h:961 */
	uint8_t		ie;			/* IEEE80211_ELEMID_COUNTRY */
	uint8_t		len;
	uint8_t		cc[3];			/* ISO CC+(I)ndoor/(O)utdoor */
	struct {
		uint8_t schan;			/* starting channel */
		uint8_t nchan;			/* number channels */
		uint8_t maxtxpwr;		/* tx power cap */
	} __attribute__((packed)) band[1];	/* sub bands (NB: var size) */
} __attribute__((packed));
#define IEEE80211_COUNTRY_MAX_SIZE \
	(sizeof(struct ieee80211_country_ie) + 3*(IEEE80211_COUNTRY_MAX_BANDS-1))  /* ieee80211.h:973 */

struct ieee80211_appie {			/* ieee80211_var.h:109 */
	uint16_t	ie_len;			/* size of ie_data */
	uint8_t		ie_data[];		/* user-specified IE's */
};

/* --- Canary guard region ---
 * The kernel kmalloc(IEEE80211_COUNTRY_MAX_SIZE) returns exactly 257 bytes
 * (from a slab bucket). We allocate a larger region, fill the bytes BEYOND the
 * nominal 257-byte alloc with a canary, run the exact IE-assembly code, then
 * inspect which canary bytes were clobbered = the OOB write.
 */
#define CANARY 0xCD
#define GUARD  16			/* bytes beyond the alloc to monitor */

/* Channel descriptor (only the fields the loop reads). */
struct chan {
	uint8_t ic_ieee;			/* channel number */
	uint8_t ic_maxregpower;		/* tx power cap (attacker-controlled via ioctl) */
};

/*
 * Replicate ieee80211_alloc_countryie body (regdomain.c:248-326), but operate
 * on a caller-provided buffer (the canary-guarded region) instead of a fresh
 * kmalloc. 'allocsz' = the size the kernel passes to kmalloc
 * (IEEE80211_COUNTRY_MAX_SIZE). Returns the number of bytes written past the
 * end of the nominal allocation.
 */
static int assemble_countryie(uint8_t *buf, size_t allocsz,
			       const struct chan *chans, int nchan,
			       uint8_t isocc0, uint8_t isocc1, char location,
			       int verbose)
{
	/* Mirror: aie = kmalloc(IEEE80211_COUNTRY_MAX_SIZE, ...)  [regdomain.c:248] */
	struct ieee80211_appie *aie = (struct ieee80211_appie *)buf;
	struct ieee80211_country_ie *ie;
	uint8_t *frm;
	int i, nruns;
	uint8_t nextchan;

	/* Fill the guard zone with canary. */
	memset(buf + allocsz, CANARY, GUARD);

	/* ie = (struct ieee80211_country_ie *) aie->ie_data;  [regdomain.c:260] */
	ie = (struct ieee80211_country_ie *)aie->ie_data;
	ie->ie = 0;	/* IEEE80211_ELEMID_COUNTRY */
	if (isocc0 == 0) {				/* regdomain.c:262 */
		ie->cc[0] = ie->cc[1] = ' ';
	} else {
		ie->cc[0] = isocc0;			/* regdomain.c:267 */
		ie->cc[1] = isocc1;			/* regdomain.c:268 */
	}
	ie->cc[2] = (location == 'I' ? 'I' :		/* regdomain.c:276 */
		     location == 'O' ? 'O' : ' ');

	/* frm = (uint8_t *)&ie->band[0];  [regdomain.c:281] */
	frm = (uint8_t *)&ie->band[0];
	nextchan = 0;					/* regdomain.c:282 */
	nruns = 0;					/* regdomain.c:283 */

	/* Channel loop — regdomain.c:290-317. We omit the skip/band logic
	 * (it only filters channels; to force 84 runs we feed 84 channels that
	 * are guaranteed to each start a new run: non-consecutive ic_ieee AND
	 * distinct ic_maxregpower so the run-break condition at :298 holds).
	 */
	for (i = 0; i < nchan; i++) {
		if (chans[i].ic_ieee != nextchan ||
		    chans[i].ic_maxregpower != frm[-1]) {	/* new run :298 */
			if (nruns == IEEE80211_COUNTRY_MAX_BANDS) {	/* :300 */
				break;
			}
			frm[0] = chans[i].ic_ieee;		/* :307 */
			frm[1] = 1;				/* :308 */
			frm[2] = chans[i].ic_maxregpower;	/* :309 */
			frm += 3;				/* :310 */
			nextchan = chans[i].ic_ieee + 1;	/* :311 */
			nruns++;				/* :312 */
		} else {					/* extend run :313 */
			frm[-2]++;				/* :314 */
			nextchan++;				/* :315 */
		}
	}
	ie->len = frm - ie->cc;				/* :318 */
	if (ie->len & 1) {				/* :319 Zero pad to multiple of 2 */
		ie->len++;				/* :320 */
		*frm++ = 0;				/* :321 */
	}
	aie->ie_len = frm - aie->ie_data;		/* :323 */

	/* Inspect guard zone. */
	int oob = 0;
	if (verbose) {
		printf("alloc size (kernel kmalloc arg) = %zu bytes\n", allocsz);
		printf("ie_data bytes available        = %zu\n", allocsz - sizeof(struct ieee80211_appie));
		printf("ie_data valid indices          = 0..%zu\n",
		       allocsz - sizeof(struct ieee80211_appie) - 1);
		printf("nruns=%d  ie->len=%u  aie->ie_len=%u\n", nruns, ie->len, aie->ie_len);
		printf("highest ie_data index written   = %ld\n", (long)((frm - 1) - aie->ie_data));
	}
	for (i = 0; i < GUARD; i++) {
		if (buf[allocsz + i] != CANARY) {
			if (verbose)
				printf("  OOB+%d: alloc[%zu] = 0x%02x (was 0x%02x)\n",
				       i + 1, allocsz + i, buf[allocsz + i], CANARY);
			if (i + 1 > oob) oob = i + 1;
		}
	}
	return oob;
}

int main(void)
{
	/* Build a channel list that forces the maximum 84 distinct runs:
	 * non-consecutive ic_ieee so nextchan never matches, and distinct
	 * ic_maxregpower so frm[-1] never matches either. Use power values
	 * 0xA0.. so the attacker-controlled byte is recognisable in the OOB
	 * region.
	 */
	struct chan chans[IEEE80211_COUNTRY_MAX_BANDS];
	int i;
	for (i = 0; i < IEEE80211_COUNTRY_MAX_BANDS; i++) {
		chans[i].ic_ieee = (uint8_t)(1 + i * 2);	/* 1,3,5,... never consecutive */
		chans[i].ic_maxregpower = (uint8_t)(0xA0 + i);	/* distinct, attacker-controlled */
	}

	/* The FIXED allocation the kernel should perform:
	 * kmalloc(sizeof(struct ieee80211_appie) + IEEE80211_COUNTRY_MAX_SIZE + 1)
	 * = 2 + 257 + 1 = 260 bytes (regdomain.c:248 after fix.diff) */
	size_t allocsz = sizeof(struct ieee80211_appie) + IEEE80211_COUNTRY_MAX_SIZE + 1;
	uint8_t *buf = calloc(1, allocsz + GUARD);
	if (!buf) { perror("calloc"); return 2; }

	printf("=== DF-0712: ieee80211_alloc_countryie 3-byte heap overflow ===\n");
	printf("IEEE80211_COUNTRY_MAX_SIZE = %zu\n", allocsz);
	printf("sizeof(struct ieee80211_appie) = %zu (ie_len header)\n",
	       sizeof(struct ieee80211_appie));
	printf("Feeding %d channels -> %d runs (== IEEE80211_COUNTRY_MAX_BANDS cap)\n\n",
	       IEEE80211_COUNTRY_MAX_BANDS, IEEE80211_COUNTRY_MAX_BANDS);

	int oob = assemble_countryie(buf, allocsz, chans, IEEE80211_COUNTRY_MAX_BANDS,
				     'U', 'S', ' ', 1 /*verbose*/);

	printf("\n");
	if (oob > 0) {
		printf("RESULT: %d-byte heap OOB write detected past the %zu-byte allocation.\n",
		       oob, allocsz);
		printf("  OOB+1 (alloc[%zu]) = number-of-channels byte = 1 (constant)\n", allocsz);
		printf("  OOB+2 (alloc[%zu]) = ic_maxregpower = ATTACKER-CONTROLLED byte\n", allocsz + 1);
		printf("  OOB+3 (alloc[%zu]) = 0x00 (odd-length IE pad byte)\n", allocsz + 2);
		printf("BUG CONFIRMED: kmalloc(IEEE80211_COUNTRY_MAX_SIZE)=%zu but the IE\n", allocsz);
		printf("  assembly writes up to alloc[%zu] (needs %zu bytes total).\n",
		       allocsz + 2, allocsz + 3);
		free(buf);
		return 0;
	} else {
		printf("RESULT: no overflow detected (bug not present).\n");
		free(buf);
		return 1;
	}
}