/*
 * DF-0843 - Deterministic race harness (PRIMARY PROOF).
 *
 * Faithful userspace TRANSCRIPTION of the unlocked race in
 * sys/vfs/ufs/ufs_dirhash.c. The DragonFly dirhash code has NO global
 * lock (no mutex / lockmgr / spinlock anywhere in the file). FreeBSD's
 * counterpart added `ufsdirhash_lock` (sx); DragonFly never did.
 *
 * The race we transcribe:
 *
 *   ufsdirhash_recycle() [sys/vfs/ufs/ufs_dirhash.c:927-970], run for a
 *   NEW inode whose build pushed memory over budget, picks a VICTIM
 *   dirhash off the GLOBAL `ufsdirhash_list` (a DIFFERENT inode) and,
 *   WITHOUT holding the victim's vnode lock or any global lock:
 *       :937  dh = TAILQ_FIRST(&ufsdirhash_list)
 *       :948  TAILQ_REMOVE(&ufsdirhash_list, dh, dh_list)
 *       :950  hash = dh->dh_hash
 *       :951  dh->dh_hash = NULL
 *       :962  kfree(hash, M_DIRHASH)        <-- hash memory freed
 *
 *   CONCURRENTLY, ufsdirhash_lookup() on the victim inode:
 *       :294  dh = ip->i_dirhash
 *       :318  if (dh->dh_hash == NULL) goto fix;   <-- unlocked NULL check
 *       :356  offset = DH_ENTRY(dh, slot)
 *             where DH_ENTRY(dh,slot) = dh->dh_hash[slot>>8][slot&255]
 *             = DOUBLE-POINTER DEREF through memory freed by recycle.
 *
 * If recycle frees `hash` between lookup's :318 check and :356 deref,
 * lookup reads freed memory through a dangling pointer = USE-AFTER-FREE.
 *
 * MODELING NOTE (important): we do NOT call real free()/munmap on the
 * hash arrays. The DragonFly slab allocator (kern_slaballoc.c, INVARIANTS)
 * does NOT unmap a page when a chunk is freed -- it keeps the page mapped
 * and writes the WEIRD_ADDR poison (0xdeadc0de) over the freed chunk so a
 * later read returns the poison pattern. We replicate that exactly: a
 * "recycle free" poisons the memory with 0xdededede and marks it logically
 * freed, but the page stays mapped. Thus the lookup thread can safely read
 * the poisoned bytes and DETECT the UAF without faulting -- exactly as the
 * kernel would observe on a no-INVARIANTS kernel, and as INVARIANTS slab
 * checks would catch on reuse. Reading the poison through the dangling
 * dh_hash pointer IS the use-after-free.
 *
 * Build:  cc -O2 -pthread -o harness harness.c
 * Run:    ./harness [iterations] [widen_ns]
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdatomic.h>
#include <pthread.h>
#include <time.h>
#include <stdint.h>
#include <sys/mman.h>

#define DH_BLKOFFSHIFT  8
#define DH_NBLKOFF      (1 << DH_BLKOFFSHIFT)
#define DH_BLKOFFMASK   (DH_NBLKOFF - 1)
#define DH_SCOREINIT    8
#define DIRHASH_EMPTY   ((int32_t)-1)

typedef int32_t doff_t;

struct dirhash {
    doff_t **dh_hash;
    int      dh_narrays;
    int      dh_hlen;
    int      dh_score;
    int      dh_onlist;
    struct { struct dirhash *tqe_next, **tqe_prev; } dh_list;
};

#define DH_ENTRY(dhp, slot) \
    ((dhp)->dh_hash[(slot) >> DH_BLKOFFSHIFT][(slot) & DH_BLKOFFMASK])

static struct dirhash *g_list_first;
#define LIST_INSERT_TAIL(dh) do {                              \
    (dh)->dh_list.tqe_next = NULL;                             \
    if (g_list_first == NULL) {                                \
        g_list_first = (dh);                                   \
        (dh)->dh_list.tqe_prev = &g_list_first;                \
    } else {                                                   \
        struct dirhash *_t = g_list_first;                     \
        while (_t->dh_list.tqe_next) _t = _t->dh_list.tqe_next;\
        _t->dh_list.tqe_next = (dh);                           \
        (dh)->dh_list.tqe_prev = &_t->dh_list.tqe_next;        \
    }                                                          \
    (dh)->dh_onlist = 1;                                       \
} while (0)
#define LIST_REMOVE(dh) do {                                   \
    if ((dh)->dh_list.tqe_next)                                \
        (dh)->dh_list.tqe_next->dh_list.tqe_prev =             \
            (dh)->dh_list.tqe_prev;                            \
    *((dh)->dh_list.tqe_prev) = (dh)->dh_list.tqe_next;        \
    (dh)->dh_onlist = 0;                                       \
} while (0)

/* Persistent pool for hash arrays. Model the slab: page stays mapped,
 * freed chunks get poisoned (0xdededede = WEIRD_ADDR analogue). */
#define POOL_BYTES  (16*1024*1024)
static doff_t *g_pool;            /* mmap'd, never munmap'd */
static size_t  g_pool_off;
static pthread_mutex_t g_pool_mtx = PTHREAD_MUTEX_INITIALIZER;

static void *pool_alloc(size_t bytes) {
    pthread_mutex_lock(&g_pool_mtx);
    /* round to 16 */
    bytes = (bytes + 15u) & ~(size_t)15;
    if (g_pool_off + bytes > POOL_BYTES) { abort(); }
    void *p = (char *)g_pool + g_pool_off;
    g_pool_off += bytes;
    pthread_mutex_unlock(&g_pool_mtx);
    return p;
}

static const uint32_t POISON = 0xdedededeu;
#define IS_POISON(v) (((uint32_t)(v)) == POISON)

/* logically-freed flag is implicit: poisoned contents == freed */
static struct dirhash *victim;
static atomic_int       uaf_detected;   /* outcome A: stale-ptr UAF (poison read) */
static atomic_int       nullderef;      /* outcome B: NULL deref of dh_hash */
static atomic_int       started;
static long             widen_ns = 0;
static size_t           pool_reset_off; /* reset each iteration (we never really free) */
static int              locked = 0;     /* model the fix: serialize recycle vs lookup */
static pthread_rwlock_t model_lock;     /* models ufsdirhash_lock (rd=shared, wr=excl) */

static struct dirhash *mk_dirhash(void) {
    struct dirhash *dh = pool_alloc(sizeof(*dh));
    memset(dh, 0, sizeof(*dh));
    dh->dh_narrays = 2;
    dh->dh_hlen    = dh->dh_narrays * DH_NBLKOFF;
    doff_t **arr = pool_alloc(dh->dh_narrays * sizeof(doff_t *));
    for (int i = 0; i < dh->dh_narrays; i++) {
        arr[i] = pool_alloc(DH_NBLKOFF * sizeof(doff_t));
        for (int j = 0; j < DH_NBLKOFF; j++) arr[i][j] = DIRHASH_EMPTY;
    }
    dh->dh_hash = arr;
    dh->dh_score = 1;
    return dh;
}

/* recycle's free path: detach + poison (memory stays mapped) */
static void recycle_free(struct dirhash *dh) {
    doff_t **hash = dh->dh_hash;
    int na = dh->dh_narrays;
    dh->dh_hash = NULL;            /* :951 NULL the victim's pointer FIRST */
    for (int i = 0; i < na; i++)
        memset(hash[i], 0xde, DH_NBLKOFF * sizeof(doff_t)); /* poison freed chunk */
    memset(hash, 0xde, na * sizeof(doff_t *));               /* poison the index array */
}

static void *recycler(void *arg) {
    (void)arg;
    while (!atomic_load(&started)) { /* spinwait until lookuper caches dh_hash */ }
    struct dirhash *dh = victim;
    if (!dh) return NULL;
    if (locked) pthread_rwlock_wrlock(&model_lock);   /* fix: LK_EXCLUSIVE */
    /* give the lookuper a head start so it caches dh_hash at :318 first,
     * then we race to NULL+free it before its :356 re-read. */
    if (widen_ns > 0) { struct timespec ts={0, widen_ns/4}; nanosleep(&ts,NULL); }
    if (dh->dh_onlist) LIST_REMOVE(dh);     /* :948 */
    recycle_free(dh);                       /* :951 NULL + :962 kfree (poison) */
    if (locked) pthread_rwlock_unlock(&model_lock);
    return NULL;
}

static void *lookuper(void *arg) {
    (void)arg;
    struct dirhash *dh = victim;
    /* :318 unlocked NULL check -- cache the pointer value the compiler
     * would keep in a register across the DH_ENTRY expression. */
    doff_t **cached = dh->dh_hash;
    atomic_store(&started, 1);
    if (cached == NULL) return NULL;        /* recycle already won before :318 */
    if (locked) pthread_rwlock_rdlock(&model_lock);   /* fix: LK_SHARED (downgrade) */
    /* re-validate under the lock (the fix's key addition) */
    if (dh->dh_hash == NULL) {
        if (locked) pthread_rwlock_unlock(&model_lock);
        return NULL;
    }
    /* window between :318 check and :356 deref -- with the fix, recycle
     * (exclusive) is blocked here, so this window cannot be exploited. */
    if (widen_ns > 0) { struct timespec ts={0, widen_ns}; nanosleep(&ts,NULL); }
    atomic_thread_fence(memory_order_acquire);
    /* :356  DH_ENTRY(dh,slot) = dh->dh_hash[slot>>8][slot&255].
     * The macro reads dh_hash to get the index-array base. With no
     * volatile/barrier the compiler is free to reuse the value cached at
     * :318 (CSE/register-allocation) -- we model that here by dereferencing
     * the CACHED pointer. The double-deref is two memory reads:
     *   (1) cached[slot>>8]  -> reads an entry of the freed dh_hash index
     *       array; after kfree+poison it is 0xdededededededede.
     *   (2) <that>[slot&255] -> derefs the wild pointer -> kernel PANIC.
     * In the real kernel step (2) faults and panics. Here we DETECT the
     * UAF at step (1) -- observing the poison in the freed index array --
     * and STOP before the wild deref, so the harness reports cleanly.
     * (Proving the index array is freed proves the UAF: the kernel would
     * immediately deref the poisoned pointer and crash.)
     * Separately, a fresh re-read returning NULL is outcome B (kernel
     * NULL-deref panic, recycle's :951 NULL now visible to this CPU). */
    doff_t **fresh = dh->dh_hash;
    if (fresh == NULL)
        atomic_store(&nullderef, 1);   /* outcome B noted */
    /* step (1) of the double-deref: read an index-array entry via the
     * cached (stale) pointer. Poison here == freed index array == UAF. */
    doff_t *leaf = cached[0];
    if ((uintptr_t)leaf == 0xdedededededededeUL)
        atomic_store(&uaf_detected, 1);   /* UAF: deref of freed index array */
    if (locked) pthread_rwlock_unlock(&model_lock);
    return NULL;
}

int main(int argc, char **argv) {
    long iters = (argc > 1) ? atol(argv[1]) : 200;
    if (argc > 2) widen_ns = atol(argv[2]);
    if (argc > 3) locked = atoi(argv[3]);   /* 1 = model the fix (lock) */

    g_pool = mmap(NULL, POOL_BYTES, PROT_READ|PROT_WRITE,
                  MAP_ANON|MAP_PRIVATE, -1, 0);
    if (g_pool == MAP_FAILED) { perror("mmap"); return 2; }
    /* allocate one dirhash to mark the start of the per-iteration region */
    pool_reset_off = 0;
    if (locked) pthread_rwlock_init(&model_lock, NULL);

    long hits = 0, nulld = 0;
    for (long i = 0; i < iters; i++) {
        atomic_store(&uaf_detected, 0);
        atomic_store(&nullderef, 0);
        atomic_store(&started, 0);
        g_list_first = NULL;
        g_pool_off = pool_reset_off;     /* fresh region each iter (no real free) */
        victim = mk_dirhash();
        LIST_INSERT_TAIL(victim);

        pthread_t tr, tl;
        pthread_create(&tr, NULL, recycler, NULL);
        pthread_create(&tl, NULL, lookuper, NULL);
        pthread_join(tr, NULL);
        pthread_join(tl, NULL);
        if (atomic_load(&uaf_detected)) hits++;
        if (atomic_load(&nullderef)) nulld++;
        victim = NULL;
    }
    printf("DF-0843 race harness: transcribes unlocked ufsdirhash_recycle/lookup/free.\n");
    printf("  mode: %s\n", locked ? "LOCKED (models the fix: recycle exclusive, lookup shared)"
                                  : "UNLOCKED (current kernel: no global lock)");
    printf("  outcome A (UAF): lookup derefs freed dh_hash index array -> poison 0xdededede\n");
    printf("  outcome B (panic): lookup re-reads dh_hash -> NULL (recycle nulled it) -> NULL-deref\n");
    printf("Ran %ld race iterations; UAF(stale-ptr)=%ld (%.1f%%)  NULL-deref=%ld\n",
           iters, hits, iters ? 100.0 * hits / iters : 0.0, nulld);
    printf("VERDICT: %s\n", (hits+nulld) ? "RACE REPRODUCED -- unlocked ufsdirhash_recycle/lookup "
                                          "lets lookup deref freed/NULL dh_hash (UAF/panic in kernel)"
                                  : (locked ? "RACE CLOSED -- the lock serializes recycle vs lookup (fix works)"
                                            : "no race observed"));
    /* success: unlocked-mode found the race, OR locked-mode closed it */
    if (locked) return (hits+nulld) ? 1 : 0;
    return (hits+nulld) ? 0 : 1;
}