/*
 * DF-5246 - Simulated reproduction of the if_run(4) beacon_mbuf race
 * ==================================================================
 *
 * Kernel code modeled (sys/bus/u4b/wlan/if_run.c, DragonFly 6.5-DEVEL):
 *
 *  run_newstate()  [thread A: net80211 nstate task]        if_run.c:2067-2106
 *      IEEE80211_UNLOCK(ic);    (drops ic lock)
 *      RUN_LOCK(sc);            (takes sc lock only)
 *      ...
 *      case IEEE80211_S_RUN:
 *          m_freem(rvp->beacon_mbuf);       <-- FREES, holds only sc lock
 *          rvp->beacon_mbuf = NULL;
 *
 *  run_update_beacon()  [thread B: tx/rx path, e.g. ieee80211_set_tim
 *                        -> vap->iv_update_beacon, ieee80211_power.c:316,
 *                        executed with the ic lock held, sc lock NOT held]
 *                                                          if_run.c:4861-4901
 *      if (rvp->beacon_mbuf == NULL) { ... alloc ... }     <-- check 4887
 *      ieee80211_beacon_update(ni, rvp->beacon_mbuf, mcast); <-- USE 4892
 *          -> writes wh->i_seq (2 bytes @ offset 22), bo_caps, TIM bytes,
 *             and on the CSA path rewrites the whole beacon (beacon_construct)
 *             through the mbuf data pointer  (ieee80211_output.c:3339+)
 *
 *  The two paths share NO common lock (ic lock vs sc lock), so the
 *  check-then-use in run_update_beacon races with the free in
 *  run_newstate / run_vap_delete (if_run.c:1013) -> write-after-free on a
 *  mbuf cluster (kernel heap).
 *
 * This harness reproduces that interleaving in userland with two threads
 * and the same two independent mutexes. A "hit" is counted when thread B
 * passes the NULL-check on a beacon mbuf that thread A frees before B's
 * write phase. (The write is withheld in the harness so the process
 * survives; in the kernel it is a real write into freed memory.)
 *
 * build: cc -O2 -pthread -o beacon_race beacon_race.c
 * run:   ./beacon_race [seconds]
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <pthread.h>
#include <unistd.h>

/* ---- the two kernel lock domains ---- */
static pthread_mutex_t ic_lock = PTHREAD_MUTEX_INITIALIZER; /* IEEE80211_LOCK */
static pthread_mutex_t sc_lock = PTHREAD_MUTEX_INITIALIZER; /* RUN_LOCK       */

/* ---- rvp->beacon_mbuf stand-in ---- */
struct mbuf_sim {
    volatile int live;         /* 1 while not m_freem'd */
    unsigned char data[512];
};
static struct mbuf_sim *beacon_mbuf;

static long uaf_hits;          /* thread B used a freed mbuf */
static long allocs;            /* beacon_alloc calls */
static long frees;             /* m_freem calls */
static long updates;           /* beacon_update calls that saw live mbuf */
static volatile int stop;

static struct mbuf_sim *beacon_alloc(void)
{
    struct mbuf_sim *m = calloc(1, sizeof(*m));
    m->live = 1;
    __sync_fetch_and_add(&allocs, 1);
    return m;
}

/* m_freem(): the kernel frees the cluster here. We mark dead then free. */
static void m_freem(struct mbuf_sim *m)
{
    m->live = 0;
    __sync_fetch_and_add(&frees, 1);
    free(m);
}

/* thread B: run_update_beacon() - executes holding ONLY ic_lock */
static void run_update_beacon(void)
{
    struct mbuf_sim *m;

    pthread_mutex_lock(&ic_lock);            /* IEEE80211_LOCK(ic) held by caller */
    m = beacon_mbuf;                         /* if_run.c:4887 check */
    if (m == NULL) {
        m = beacon_alloc();                  /* if_run.c:4888 */
        beacon_mbuf = m;
    }
    /*
     * ieee80211_beacon_update(ni, m, mcast): writes seq/capinfo/TIM
     * through the mbuf. We mirror the write set of the kernel function
     * but first test liveness: if the mbuf was freed between the
     * NULL-check above and here, the kernel would have written into
     * freed heap. Count it and skip the actual write.
     */
    if (!m->live) {
        __sync_fetch_and_add(&uaf_hits, 1);
    } else {
        m->data[22] = 0x11; m->data[23] = 0x22;   /* wh->i_seq */
        memset(m->data + 60, 0x5a, 8);            /* TIM bytes */
        __sync_fetch_and_add(&updates, 1);
    }
    pthread_mutex_unlock(&ic_lock);
}

/* thread B body: hammer the TIM-update path (ieee80211_set_tim) */
static void *beacon_updater(void *arg)
{
    (void)arg;
    while (!stop)
        run_update_beacon();
    return NULL;
}

/* thread A: run_newstate() RUN->INIT - executes holding ONLY sc_lock */
static void run_newstate_to_init(void)
{
    pthread_mutex_lock(&sc_lock);
    /* if_run.c:2105-2106 */
    if (beacon_mbuf != NULL) {
        m_freem(beacon_mbuf);
        beacon_mbuf = NULL;
    }
    pthread_mutex_unlock(&sc_lock);
}

/* thread A body: state transitions RUN->INIT->RUN forever */
static void *state_machine(void *arg)
{
    (void)arg;
    while (!stop) {
        run_newstate_to_init();
        usleep(50);          /* re-enter RUN, beacon re-allocated by updater */
    }
    return NULL;
}

int main(int argc, char **argv)
{
    int secs = (argc > 1) ? atoi(argv[1]) : 5;
    pthread_t ta, tb, tb2;

    beacon_mbuf = beacon_alloc();

    pthread_create(&ta, NULL, state_machine, NULL);
    pthread_create(&tb, NULL, beacon_updater, NULL);
    pthread_create(&tb2, NULL, beacon_updater, NULL);

    sleep(secs);
    stop = 1;
    pthread_join(ta, NULL);
    pthread_join(tb, NULL);
    pthread_join(tb2, NULL);

    printf("DF-5246 beacon_mbuf race simulation over %ds\n", secs);
    printf("  beacon_allocs : %ld\n", allocs);
    printf("  m_freems      : %ld\n", frees);
    printf("  clean updates : %ld\n", updates);
    printf("  UAF hits      : %ld   (write-after-free in kernel equivalent)\n",
        uaf_hits);
    if (uaf_hits > 0) {
        printf("RESULT: RACE OBSERVED - run_update_beacon used a freed mbuf\n");
        return 0;
    }
    printf("RESULT: no hit this run (widen load / rerun)\n");
    return 1;
}