/*
 * DF-5292 simulation harness - if_rum.c beacon mbuf write-after-free.
 *
 * Models, line-for-line, the lock domains of:
 *
 *  [U] rum_update_beacon()   sys/bus/u4b/wlan/if_rum.c:2792-2831
 *        - called from ieee80211_beacon_notify() with IEEE80211_LOCK(ic) held
 *        - 2800:  m = rvp->bcn_mbuf                      (ic domain only)
 *        - 2803:  RUM_LOCK(sc)
 *        - 2805:  m = ieee80211_beacon_alloc(ni)  (if NULL)
 *        - 2812:  rvp->bcn_mbuf = m
 *        - 2825:  RUM_UNLOCK(sc)          <-- sc protection dropped
 *        - 2828:  ieee80211_beacon_update(ni, m, mcast)  WRITES seq/capinfo/
 *                TIM/ERP element data INTO the mbuf cluster, under the ic
 *                lock only (ieee80211_output.c:3339 ff.)
 *        - 2830:  rum_cmd_sleepable(sc, &vap, ...)       (DF-5293 vector)
 *
 *  [N] rum_newstate()        if_rum.c:916-1025
 *        - 932:   IEEE80211_UNLOCK(ic)                   <-- ic dropped
 *        - 933:   RUM_LOCK(sc)
 *        - 977:   rum_alloc_beacon(sc, vap)
 *            2775-2776:  m_freem(rvp->bcn_mbuf)          (sc domain)
 *            2778:       rvp->bcn_mbuf = new
 *
 *  [D] rum_vap_delete()      if_rum.c:709-721
 *        - 715:   m_freem(rvp->bcn_mbuf)                 (NO lock at all)
 *
 * The ic lock and the sc lock are disjoint.  Thread U writes into the beacon
 * cluster after dropping the sc lock; thread N frees the very same cluster
 * while holding only the sc lock; thread D frees it with no lock.  Any free
 * that lands between U's unlock(sc) (2825) and the end of U's in-mbuf writes
 * (2828) is a kernel write-after-free on a 2 KB cluster.
 *
 * The harness keeps the exact lock protocol and counts (instead of dying on
 * the first corruption) how many beacon_update write phases observed a freed
 * cluster, using an atomic poison flag - the userland equivalent of what the
 * kernel does right before it shreds freed slab memory.
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <pthread.h>
#include <stdatomic.h>
#include <unistd.h>
#include <time.h>

static pthread_mutex_t ic_lock = PTHREAD_MUTEX_INITIALIZER; /* IEEE80211_LOCK */
static pthread_mutex_t sc_lock = PTHREAD_MUTEX_INITIALIZER; /* RUM_LOCK      */

struct mbuf_model {
	atomic_int freed;
	unsigned char data[2048];	/* MCLBYTES cluster */
};

struct rum_vap_model {
	struct mbuf_model *bcn_mbuf;
} rvap;

static atomic_long uaf_writes;	/* writes into a freed cluster */
static atomic_long updates;	/* rum_update_beacon iterations */
static atomic_long newstate_repl;	/* rum_alloc_beacon replacements */
static volatile int stop;

static struct mbuf_model *
beacon_alloc(void)
{
	struct mbuf_model *m = calloc(1, sizeof(*m));
	atomic_store(&m->freed, 0);
	return m;
}

/* [U] rum_update_beacon, ic lock held on entry (net80211 caller) */
static void
update_beacon(unsigned id)
{
	struct mbuf_model *m;

	pthread_mutex_lock(&ic_lock);			/* caller context */
	m = rvap.bcn_mbuf;				/* :2800 */

	pthread_mutex_lock(&sc_lock);			/* :2803 */
	if (m == NULL) {
		m = beacon_alloc();			/* :2805 */
		rvap.bcn_mbuf = m;			/* :2812 */
	}
	pthread_mutex_unlock(&sc_lock);			/* :2825 */

	/*
	 * :2828 ieee80211_beacon_update(ni, m, mcast):
	 * writes seqno, capinfo, TIM bitmap, ERP/HT elements into m->data.
	 * Modelled as a burst of writes across the cluster; a yield models
	 * the scheduler preemption the in-kernel TIM/element recompute path
	 * is exposed to (the real window is microseconds and unbounded under
	 * preemption; the raw yield-less model is the optimistic floor).
	 */
	for (int off = 0; off < 2048; off += 64) {
		m->data[off] = (unsigned char)(id + off);
		if (off == 1024)
			sched_yield();
	}
	if (atomic_load(&m->freed))
		atomic_fetch_add(&uaf_writes, 1);
	atomic_fetch_add(&updates, 1);

	pthread_mutex_unlock(&ic_lock);
}

/* [N] rum_newstate(S_RUN) -> rum_alloc_beacon, sc lock only (ic dropped) */
static void
newstate_run(void)
{
	struct mbuf_model *old, *nm;

	pthread_mutex_lock(&sc_lock);			/* :933 */
	old = rvap.bcn_mbuf;
	nm = beacon_alloc();
	rvap.bcn_mbuf = nm;				/* :2778 */
	if (old != NULL)
		atomic_store(&old->freed, 1);		/* :2776 m_freem(old) */
	pthread_mutex_unlock(&sc_lock);
	atomic_fetch_add(&newstate_repl, 1);
	free(old);
}

/* [D] rum_vap_delete: m_freem(rvp->bcn_mbuf) with NO lock (:715) */
static void
vap_delete(void)
{
	struct mbuf_model *old = rvap.bcn_mbuf;
	rvap.bcn_mbuf = NULL;
	if (old != NULL)
		atomic_store(&old->freed, 1);
	free(old);
}

static void *
thr_update(void *arg)
{
	unsigned id = 0;
	(void)arg;
	while (!stop) {
		update_beacon(id++);
		usleep(200);		/* beacon-update churn (TIM/ERP/CSA) */
	}
	return NULL;
}

static void *
thr_newstate(void *arg)
{
	(void)arg;
	while (!stop) {
		newstate_run();		/* S_RUN re-entry / re-association */
		usleep(300);
	}
	return NULL;
}

static void *
thr_delete(void *arg)
{
	(void)arg;
	while (!stop) {
		pthread_mutex_lock(&sc_lock);
		vap_delete();
		pthread_mutex_unlock(&sc_lock);
		usleep(2000);		/* ifconfig wlanX destroy */
	}
	return NULL;
}

int
main(int argc, char **argv)
{
	int secs = argc > 1 ? atoi(argv[1]) : 10;
	pthread_t tu, tn, td;
	struct timespec a, b;

	clock_gettime(CLOCK_MONOTONIC, &a);
	pthread_create(&tu, NULL, thr_update, NULL);
	pthread_create(&tn, NULL, thr_newstate, NULL);
	pthread_create(&td, NULL, thr_delete, NULL);
	sleep(secs);
	stop = 1;
	pthread_join(tu, NULL);
	pthread_join(tn, NULL);
	pthread_join(td, NULL);
	clock_gettime(CLOCK_MONOTONIC, &b);

	printf("rum beacon-mbuf race model (%d s): update_beacon=%ld "
	    "newstate_replace=%ld freed-in-window UAF WRITES=%ld (%.1f%%)\n",
	    secs, atomic_load(&updates), atomic_load(&newstate_repl),
	    atomic_load(&uaf_writes),
	    atomic_load(&updates) ?
	    100.0 * (double)atomic_load(&uaf_writes) / (double)atomic_load(&updates) : 0.0);
	(void)(b.tv_sec - a.tv_sec);
	return 0;
}