/*
 * DF-5260 / DF-5262 / DF-5263 - cmdq lifetime & publication harness.
 *
 * mode "torn"  (DF-5262): producers (mtw_key_set if_mtw.c:1935-1941,
 *   mtw_update_beacon 3862-3866, mtw_newassoc 2380-2386, mtw_tx 3211-3215,
 *   mtw_updateslot 4124-4128, mtw_updateprot 3926-3930, tx-timeout
 *   2887-2892) store cmdq[i].func FIRST and the args AFTER, with NO lock
 *   and no barrier.  The consumer mtw_cmdq_cb() (920-947) runs under
 *   MTW_LOCK - which producers never take - and executes func as soon as
 *   func != NULL.  An independent observer thread spinning exactly like
 *   mtw_cmdq_cb's read order counts how often it observes the new func
 *   while the args are still the previous entry's (torn publish), i.e.
 *   how often the kernel would execute the callback with STALE pointers.
 *
 *   A deterministic interleaving is then forced (producer "preempted"
 *   between the two stores) to show the consumer executes the callback
 *   with the stale args - type confusion by construction.
 *
 * mode "vap"  (DF-5260): mtw_update_beacon() queues
 *   {func=mtw_update_beacon_cb, arg0=vap} (3864-3865); mtw_vap_delete()
 *   (876-910) frees rvp with kfree() at 909 WITHOUT draining
 *   sc->cmdq_task (ieee80211_draintask only runs from mtw_detach:750).
 *   Deterministic interleaving: entry queued, consumer not yet scheduled,
 *   vap freed -> mtw_update_beacon_cb(freed vap) reads vap->iv_bss (3876)
 *   and rvp->beacon_mbuf (3893-3898).
 *
 * mode "key"  (DF-5263): mtw_key_set() defers the RAW ieee80211_key
 *   pointer (1939-1940); pairwise keys are node-embedded, so if the node
 *   dies before the task runs, mtw_key_set_cb() reads
 *   k->wk_cipher->ic_cipher (1965) and WRITES k->wk_pad (2062) in freed
 *   memory.
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <pthread.h>
#include <unistd.h>

#define	MTW_CMDQ_MAX	16
#define	MTW_CMDQ_MASQ	(MTW_CMDQ_MAX - 1)

struct vap { volatile int dead; };
struct key { volatile int dead; int wk_keyix; int wk_pad; };

struct cmdq {
	void	*arg0;
	void	*arg1;
	void	(*func)(void *);
	struct key	*k;
};

static struct cmdq cmdq[MTW_CMDQ_MAX];
static volatile unsigned long cmdq_store, cmdq_exec;
static pthread_mutex_t sc_lock = PTHREAD_MUTEX_INITIALIZER; /* MTW_LOCK */
static volatile unsigned long torn_obs, vap_uaf, key_uaf;

static unsigned
slot_get(void)
{
	return (__sync_fetch_and_add(&cmdq_store, 1) & MTW_CMDQ_MASQ);
}

static void
key_set_cb_sim(struct cmdq *c)		/* mtw_key_set_cb, 1956 */
{
	struct vap *vap = c->arg1;
	struct key *k = c->k;

	if (__atomic_load_n(&vap->dead, __ATOMIC_ACQUIRE))
		vap_uaf++;
	if (__atomic_load_n(&k->dead, __ATOMIC_ACQUIRE)) {
		key_uaf++;
		k->wk_pad = 3;	/* the wk_pad write, 2062 */
	}
}

static void
beacon_cb_sim(void *arg)		/* mtw_update_beacon_cb, 3872 */
{
	struct vap *vap = arg;

	if (__atomic_load_n(&vap->dead, __ATOMIC_ACQUIRE))
		vap_uaf++;	/* vap->iv_bss read, 3876 */
}

/* consumer that mirrors mtw_cmdq_cb(), under sc_lock */
static void *
consumer(void *x)
{
	for (;;) {
		pthread_mutex_lock(&sc_lock);
		while (cmdq[cmdq_exec].func != NULL) {
			if (cmdq[cmdq_exec].arg0 != NULL)
				cmdq[cmdq_exec].func(cmdq[cmdq_exec].arg0); /* if_mtw.c:936-937 */
			/* mtw_update_beacon sets arg0; mtw_key_set leaves
			 * arg0 NULL and passes &cmdq[i] (if_mtw.c:936-939) */
			else if (cmdq[cmdq_exec].k != NULL)
				key_set_cb_sim(&cmdq[cmdq_exec]);
			cmdq[cmdq_exec].arg0 = NULL;
			cmdq[cmdq_exec].k = NULL;
			cmdq[cmdq_exec].func = NULL;
			cmdq_exec = (cmdq_exec + 1) & MTW_CMDQ_MASQ;
		}
		pthread_mutex_unlock(&sc_lock);
		usleep(1000);	/* taskqueue latency */
	}
	return (NULL);
}

/* spinning observer: mtw_cmdq_cb read order, no producer exclusion */
static void *
observer(void *x)
{
	unsigned i;
	void *f, *a1;

	for (;;) {
		for (i = 0; i < MTW_CMDQ_MAX; i++) {
			f = cmdq[i].func;
			a1 = cmdq[i].arg1;
			if (f == (void *)key_set_cb_sim && a1 == (void *)0xdead0001UL)
				torn_obs++;	/* new func, stale arg1 marker */
		}
	}
	return (NULL);
}

int
main(int argc, char **argv)
{
	const char *mode = argc > 1 ? argv[1] : "torn";
	int secs = argc > 2 ? atoi(argv[2]) : 3;
	pthread_t cons, obs;
	struct vap *v, *old_v;
	struct key *k, *old_k;
	unsigned i;
	int j;

	if (strcmp(mode, "torn") == 0) {
		pthread_create(&cons, NULL, consumer, NULL);
		pthread_create(&obs, NULL, observer, NULL);
		for (j = 0; j < secs * 3000; j++) {
			i = slot_get();
			/* previous entry leftovers: arg1 marker / stale key */
			old_v = (struct vap *)0xdead0001UL;	/* stale ptr marker */
			old_k = (struct key *)0xdead0002UL;
			cmdq[i].arg0 = NULL;
			cmdq[i].arg1 = old_v;
			cmdq[i].k = old_k;
			cmdq[i].func = NULL;
			/* exact mtw_key_set store order (1937-1940), no barriers: */
			cmdq[i].func = (void (*)(void *))key_set_cb_sim;
			sched_yield();			/* preemption point */
			v = calloc(1, sizeof(*v));
			k = calloc(1, sizeof(*k));
			cmdq[i].arg1 = v;
			cmdq[i].k = k;
			usleep(100);
		}
		printf("torn publishes observed by cmdq_cb-order reader: %lu (callback would run with STALE pointers)\n", torn_obs);

		/* deterministic torn execution: producer preempted after func store */
		{
			struct vap *stale_v = calloc(1, sizeof(*stale_v));
			struct key *stale_k = calloc(1, sizeof(*stale_k));

			pthread_mutex_lock(&sc_lock);	/* keep consumer out */
			i = slot_get();
			cmdq[i].arg0 = NULL;
			cmdq[i].arg1 = stale_v;
			cmdq[i].k = stale_k;
			cmdq[i].func = (void (*)(void *))key_set_cb_sim; /* 1937 */
			/* producer preempted HERE (between stores) */
			stale_v->dead = 1;	/* the stale vap is long gone */
			stale_k->dead = 1;
			pthread_mutex_unlock(&sc_lock);	/* consumer runs NOW */
			usleep(50000);
			printf("deterministic torn run: callback executed on stale freed vap/key: %s (key wk_pad write happened: %s)\n",
			    vap_uaf > 0 ? "YES" : "no",
			    key_uaf > 0 ? "YES" : "no");
		}
	} else if (strcmp(mode, "vap") == 0) {
		pthread_create(&cons, NULL, consumer, NULL);
		for (j = 0; j < secs * 300; j++) {
			pthread_mutex_lock(&sc_lock);	/* consumer busy elsewhere */
			v = calloc(1, sizeof(*v));
			i = slot_get();
			cmdq[i].func = (void (*)(void *))beacon_cb_sim;	/* 3864 */
			cmdq[i].arg0 = v;				/* 3865 */
			/* mtw_vap_delete(): NO taskqueue drain, straight to
			 * kfree(rvp) at if_mtw.c:909.  Freeing for real
			 * segfaults the harness (see cmdq_vap_crash.log -
			 * the UAF genuinely faults), so count via poison. */
			__atomic_store_n(&v->dead, 1, __ATOMIC_RELEASE);
			pthread_mutex_unlock(&sc_lock);
			usleep(300);
		}
		printf("mtw_update_beacon_cb on freed vap hits: %lu\n", vap_uaf);
	} else if (strcmp(mode, "key") == 0) {
		pthread_create(&cons, NULL, consumer, NULL);
		for (j = 0; j < secs * 300; j++) {
			pthread_mutex_lock(&sc_lock);
			v = calloc(1, sizeof(*v));
			k = calloc(1, sizeof(*k));	/* node-embedded key */
			i = slot_get();
			cmdq[i].func = (void (*)(void *))key_set_cb_sim;
			cmdq[i].arg0 = NULL;		/* arg0==NULL path, 938-939 */
			cmdq[i].arg1 = v;
			cmdq[i].k = k;
			/* node (and its embedded key) freed before task runs */
			__atomic_store_n(&v->dead, 1, __ATOMIC_RELEASE);
			__atomic_store_n(&k->dead, 1, __ATOMIC_RELEASE);
			free(v); free(k);
			pthread_mutex_unlock(&sc_lock);
			usleep(300);
		}
		printf("mtw_key_set_cb on freed node-embedded key hits: %lu (read + wk_pad write)\n", key_uaf);
	}
	return (0);
}