/*
 * DF-5254 simulation harness: use-after-free WRITE on urtwn beacon mbuf
 * caused by disjoint lock domains between:
 *
 *   thread A: ieee80211_set_tim() [com lock held] -> iv_update_beacon ->
 *             urtwn_update_beacon()  (sys/bus/u4b/wlan/if_urtwn.c:2187)
 *               sc_mtx HELD   : alloc-if-NULL   (lines 2196-2205)
 *               sc_mtx RELEASED: ieee80211_beacon_update(ni, uvp->bcn_mbuf,.)
 *                               writes TIM/CSA fields INTO the mbuf
 *                               (lines 2206-2212)   <-- no sc_mtx, com held
 *               sc_mtx HELD   : urtwn_tx_beacon  (lines 2214-2216)
 *
 *   thread B: ieee80211_ibss_merge() [remote attacker via crafted beacon]
 *             -> ieee80211_sta_join1 -> newstate(S_RUN) -> urtwn_newstate()
 *               IEEE80211_UNLOCK(ic) at line 2551 (com lock DROPPED)
 *               URTWN_LOCK(sc) at 2552, then urtwn_setup_beacon() does:
 *                   m_freem(uvp->bcn_mbuf);          (line 2169)
 *                   uvp->bcn_mbuf = new_m;           (line 2171)
 *             -- com lock NOT held, sc_mtx held: nothing serializes B
 *                against A's unlocked window.
 *
 *   thread C: urtwn_vap_delete() frees uvp->bcn_mbuf with NO lock at all
 *             (lines 840-841).
 *
 * This program models exactly that interleaving with a zone-allocator-like
 * mbuf simulator (generation counters + reuse pool).  A write whose target
 * generation is stale is a write-after-free.  Exit code 3 == detected.
 *
 * Build: cc -O2 -pthread -o race_sim race_sim.c
 * Run:   ./race_sim [seconds]
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <pthread.h>
#include <unistd.h>
#include <time.h>

static pthread_mutex_t sc_mtx  = PTHREAD_MUTEX_INITIALIZER; /* urtwn sc_mtx   */
static pthread_mutex_t com_mtx = PTHREAD_MUTEX_INITIALIZER; /* net80211 com   */

#define MCLBYTES 2048
struct mbuf_sim {
	volatile uint64_t gen;              /* allocation generation           */
	volatile uint64_t dead;             /* 1 once "m_freem" ran            */
	uint8_t         data[MCLBYTES];     /* the "mbuf cluster"              */
};

/* --- zone allocator emulation (like mbuf clusters: freed => reusable) --- */
static struct mbuf_sim *pool[512];
static int pool_n;
static pthread_mutex_t pool_mtx = PTHREAD_MUTEX_INITIALIZER;
static volatile uint64_t cur_gen = 1;

static struct mbuf_sim *mb_alloc(void)
{
	struct mbuf_sim *m;
	pthread_mutex_lock(&pool_mtx);
	if (pool_n > 0) {
		m = pool[--pool_n];            /* reuse a freed cluster           */
	} else {
		m = malloc(sizeof(*m));
		if (m == NULL) { pthread_mutex_unlock(&pool_mtx); exit(9); }
	}
	pthread_mutex_unlock(&pool_mtx);
	memset(m->data, 0, MCLBYTES);
	m->dead = 0;
	m->gen = ++cur_gen;
	return m;
}

/* model of m_freem(): cluster becomes dead immediately, reusable later */
static void mb_freem(struct mbuf_sim *m)
{
	m->dead = 1;
	/* poison the first word, like a real allocator scribbling free pages */
	*(volatile uint64_t *)&m->data[0] = 0xdeadbeefdeadbeefULL;
	pthread_mutex_lock(&pool_mtx);
	if (pool_n < (int)(sizeof(pool)/sizeof(pool[0])))
		pool[pool_n++] = m;
	pthread_mutex_unlock(&pool_mtx);
}

/* ---- the shared driver state: struct urtwn_vap.bcn_mbuf ---- */
static struct mbuf_sim *volatile bcn_mbuf;

/* statistics */
static volatile unsigned long iters_a, iters_b, iters_c;
static volatile unsigned long uaf_writes, uaf_reads, detected;
static volatile unsigned long first_gen;

static double now(void)
{
	struct timespec ts;
	clock_gettime(CLOCK_MONOTONIC, &ts);
	return ts.tv_sec + ts.tv_nsec / 1e9;
}

/*
 * thread A == net80211 task / RX-driven TIM update path:
 *   com_mtx held for the whole call (ieee80211_power.c:310-318),
 *   sc_mtx deliberately released around beacon_update (if_urtwn.c:2206).
 */
static void *thread_update_beacon(void *arg)
{
	(void)arg;
	for (;;) {
		pthread_mutex_lock(&com_mtx);               /* set_tim entry */

		pthread_mutex_lock(&sc_mtx);                /* urtwn_update_beacon */
		if (bcn_mbuf == NULL)
			bcn_mbuf = mb_alloc();              /* 2197-2205 */
		pthread_mutex_unlock(&sc_mtx);              /* 2206 *** WINDOW *** */

		/* ieee80211_beacon_update(ni, uvp->bcn_mbuf, mcast) -- line 2212:
		 * the mbuf pointer is loaded NOW and then written through,
		 * without sc_mtx. */
		{
			struct mbuf_sim *m = bcn_mbuf;
			uint64_t g0 = m->gen;
			/* beacon_update writes TIM bitmap, seq, csa count... */
			for (int i = 16; i < 512; i += 8) {
				if (m->dead && !detected) {
					detected = 1;
					first_gen = g0;
					fprintf(stderr,
"[UAF-WRITE] thread A wrote into a freed beacon mbuf (gen %llu) -- "
"freed by concurrent setup_beacon/vap_delete while A was inside "
"ieee80211_beacon_update() with sc_mtx released (if_urtwn.c:2206-2212)\n",
					    (unsigned long long)g0);
				}
				if (m->dead)
					uaf_writes++;
				*(volatile uint64_t *)&m->data[i] = 0xA5A5A5A5A5A5A5A5ULL;
			}
		}

		pthread_mutex_lock(&sc_mtx);                /* 2214 */
		if (bcn_mbuf && bcn_mbuf->dead)		/* tx_beacon reads it too */
			uaf_reads++;
		pthread_mutex_unlock(&sc_mtx);

		pthread_mutex_unlock(&com_mtx);             /* set_tim exit */
		iters_a++;
	}
	return NULL;
}

/*
 * thread B == IBSS-merge-driven newstate(S_RUN) re-entry:
 *   urtwn_newstate drops com_mtx (2551), takes sc_mtx (2552) and
 *   urtwn_setup_beacon frees + swaps bcn_mbuf (2168-2171).
 */
static void *thread_newstate(void *arg)
{
	(void)arg;
	for (;;) {
		struct mbuf_sim *m = mb_alloc();

		/* emulate scheduling jitter so both interleavings get exercised */
		if (iters_b % 64 == 0) usleep(50);

		pthread_mutex_lock(&com_mtx);   /* outer state machine */
		pthread_mutex_unlock(&com_mtx); /* IEEE80211_UNLOCK(ic):2551 */
		pthread_mutex_lock(&sc_mtx);    /* URTWN_LOCK(sc):2552 */
		if (bcn_mbuf)
			mb_freem(bcn_mbuf);     /* m_freem:2169 */
		bcn_mbuf = m;                  /* 2171 */
		pthread_mutex_unlock(&sc_mtx);
		iters_b++;
	}
	return NULL;
}

/*
 * thread C == urtwn_vap_delete(): frees bcn_mbuf holding NO lock (840-841).
 * The kernel leaves uvp->bcn_mbuf dangling (vap is being freed), so the
 * model frees the cluster but keeps the field pointing at it until the
 * next setup_beacon swap.  Models ~5 vap create/destroy cycles.
 */
static void *thread_vap_delete(void *arg)
{
	(void)arg;
	int cycles = 5;
	while (cycles-- > 0) {
		usleep(1000 + (rand() & 1023));
		if (bcn_mbuf && !bcn_mbuf->dead) {
			mb_freem(bcn_mbuf);     /* NO sc_mtx, NO com_mtx; field dangles */
		}
		iters_c++;
	}
	return NULL;
}

int main(int argc, char **argv)
{
	double budget = argc > 1 ? atof(argv[1]) : 5.0;
	double t0 = now();
	pthread_t ta, tb, tc;

	srand(1234);
	pthread_create(&ta, NULL, thread_update_beacon, NULL);
	pthread_create(&tb, NULL, thread_newstate, NULL);
	pthread_create(&tc, NULL, thread_vap_delete, NULL);

	while (now() - t0 < budget) {
		if (detected)
			break;
		usleep(1000);
	}

	printf("DF-5254 disjoint-lock-domain beacon mbuf race simulation\n");
	printf("  update_beacon iterations (A) : %lu\n", iters_a);
	printf("  newstate re-entries     (B)   : %lu\n", iters_b);
	printf("  vap_delete frees        (C)   : %lu\n", iters_c);
	printf("  writes into freed mbuf        : %lu\n", uaf_writes);
	printf("  reads  of  freed mbuf (tx)    : %lu\n", uaf_reads);
	if (detected) {
		printf("RESULT: USE-AFTER-FREE WRITE reproduced (freed gen %llu)\n",
		    (unsigned long long)first_gen);
		printf("The interleaving urtwn_update_beacon(if_urtwn.c:2206-2212) "
		       "x {urtwn_setup_beacon(2168-2171) | urtwn_vap_delete(840-841)} "
		       "is realizable: no lock orders A against B/C.\n");
		exit(3);
	}
	printf("RESULT: not observed within budget (race still permitted by code)\n");
	return 1;
}