/*
 * DF-1863 — TOCTOU in amdgpu_ctx_add_fence: harness.
 *
 * amdgpu is NOT compiled into the X86_64_GENERIC kernel and there is no AMD
 * GPU on this guest, so the in-kernel path cannot be exercised end-to-end on
 * the default kernel.  This harness reproduces the *exact algorithmic
 * primitive* — the lockless read of `centity->sequence` / `idx` /
 * `centity->fences[idx]` followed by a `dma_fence_put(other)` *outside* the
 * `ring_lock` critical section — in userspace, using pthreads and atomics.
 *
 * Two worker threads concurrently call the function on the same ctx/entity
 * (same precondition as two AMDGPU_CS ioctls against the same ctx_id / ip_type
 * / ring, per the threat model in the finding).  We count:
 *
 *   double_put   — `other` had dma_fence_put() invoked on it after its
 *                  refcount had already reached 0.  In the kernel this is the
 *                  kfree() / poisoned-slab touch (UAF/double-free).
 *   lost_fence   — a `fence` was stored into a slot and then silently
 *                  overwritten by the losing thread before its ctx-held ref
 *                  was ever consumed (its refcount never returns to 0 via the
 *                  ctx-held put).  This is the leaked fence from the finding.
 *
 * Run both the BUGGY and FIXED variants of the function and compare counts.
 *
 * Build:  cc -O2 -Wall -pthread -o race race.c
 * Run:    ./race
 * Expect: BUGGY reports double_put > 0 and lost_fence > 0 within a few hundred
 *         iterations; FIXED reports both == 0 across millions of iterations.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <pthread.h>
#include <stdatomic.h>

/*
 * Mirrors `int amdgpu_sched_jobs = 32;` (sys/dev/drm/amd/amdgpu/amdgpu_drv.c:108).
 * Must be a power of two so `seq & (amdgpu_sched_jobs - 1)` is a clean mask.
 */
#define AMDGPU_SCHED_JOBS    32

/*
 * struct dma_fence (Linux DRM) has a kref refcount.  We model the two ops the
 * algorithm uses: dma_fence_get (atomic inc) and dma_fence_put (atomic dec,
 * free at 0, detect double-put).  We track live_count so a double-put on an
 * already-freed fence is unambiguous.
 */
struct dma_fence {
    atomic_int      refcount;   /* mirrors fence->refcount.refcount */
    atomic_int      freed;      /* 0 until dma_fence_put drops refcount to 0 */
    int             id;
};

static atomic_uint g_live_fences = 0;       /* live fence objects total      */

static struct dma_fence *dma_fence_alloc(int id)
{
    struct dma_fence *f = calloc(1, sizeof(*f));
    if (!f) { perror("calloc"); exit(2); }
    atomic_store(&f->refcount, 1);          /* caller's initial ref          */
    atomic_store(&f->freed, 0);
    f->id = id;
    atomic_fetch_add(&g_live_fences, 1);
    return f;
}

/* dma_fence_get: take a new ref. */
static void dma_fence_get(struct dma_fence *f)
{
    atomic_fetch_add(&f->refcount, 1);
}

/*
 * dma_fence_put: drop one ref.  If refcount hits 0 the object is freed
 * (mirrors drm_sched_fence_free -> kfree via the kref release callback).
 * A *double put* is when put is called after the object was already freed.
 */
static atomic_uint g_double_put = 0;
static unsigned    g_buggy_double_put = 0;
static unsigned    g_buggy_leaked     = 0;

/*
 * dma_fence_put: drop one ref.  When refcount reaches 0 the object would be
 * kfree()'d in the kernel; we instead transition it to a "tombstone" state
 * (freed=1) so subsequent puts land on a stable, recognizable object rather
 * than triggering a real userspace double-free (which would crash the harness
 * before we could print the statistics).  A put on a tombstoned fence, OR a
 * put that would drive refcount negative, is the double-put signature.
 *
 * We deliberately keep the object allocated so the second racing put reads a
 * well-defined freed=1 and increments g_double_put.  This mirrors the kernel
 * case where the slab page is poisoned (0xdeadc0de on DF INVARIANTS builds)
 * and the second put touches that poison — except here it's a counted signal,
 * not a panic.
 */
static void dma_fence_put(struct dma_fence *f)
{
    if (f == NULL)               /* mirrors Linux dma_fence_put: no-op on NULL */
        return;
    int prev = atomic_fetch_sub(&f->refcount, 1);
    if (atomic_load(&f->freed) || prev - 1 < 0) {
        atomic_fetch_add(&g_double_put, 1);
        return;
    }
    if (prev - 1 == 0) {
        atomic_store(&f->freed, 1);
        atomic_fetch_sub(&g_live_fences, 1);
    }
}

/*
 * Mirrors struct amdgpu_ctx_entity { sequence; fences[]; }.
 */
struct amdgpu_ctx_entity {
    _Atomic uint64_t   sequence;
    struct dma_fence  *fences[AMDGPU_SCHED_JOBS];
};

struct amdgpu_ctx {
    pthread_mutex_t    ring_lock;       /* = ctx->ring_lock (DF: struct lock) */
    struct amdgpu_ctx_entity centity;
};

/*
 * BUGGY variant — faithful copy of amdgpu_ctx.c:440-464 (master DEV).
 * Read of seq / idx / other happens BEFORE lockmgr; put(other) happens AFTER
 * lockmgr releases.  Two threads racing here will both observe the same
 * `other` and both dma_fence_put() it.
 */
static void ctx_add_fence_BUGGY(struct amdgpu_ctx *ctx,
                                struct dma_fence *fence)
{
    struct amdgpu_ctx_entity *centity = &ctx->centity;
    uint64_t seq = atomic_load(&centity->sequence);                 /* L445 */
    unsigned idx = seq & (AMDGPU_SCHED_JOBS - 1);                   /* L449 */
    struct dma_fence *other = centity->fences[idx];                 /* L450 */

    dma_fence_get(fence);                                           /* L454 */

    pthread_mutex_lock(&ctx->ring_lock);                            /* L456 */
    centity->fences[idx] = fence;                                   /* L457 */
    atomic_store(&centity->sequence, seq + 1);                      /* L458 */
    pthread_mutex_unlock(&ctx->ring_lock);                          /* L459 */

    dma_fence_put(other);                                           /* L461 */
}

/*
 * FIXED variant — the proposed fix.diff: move the read of seq / idx / other
 * INSIDE the ring_lock critical section so each slot is claimed by exactly
 * one thread.
 */
static void ctx_add_fence_FIXED(struct amdgpu_ctx *ctx,
                                struct dma_fence *fence)
{
    struct amdgpu_ctx_entity *centity = &ctx->centity;
    uint64_t seq;
    unsigned idx;
    struct dma_fence *other;

    dma_fence_get(fence);

    pthread_mutex_lock(&ctx->ring_lock);
    seq   = atomic_load(&centity->sequence);
    idx   = seq & (AMDGPU_SCHED_JOBS - 1);
    other = centity->fences[idx];
    centity->fences[idx] = fence;
    atomic_store(&centity->sequence, seq + 1);
    pthread_mutex_unlock(&ctx->ring_lock);

    dma_fence_put(other);
}

/* ---------- driver ----------------------------------------------------- */

#define N_THREADS    2
#define ITERS_PER_T  4000

struct work {
    struct amdgpu_ctx *ctx;
    void            (*fn)(struct amdgpu_ctx *, struct dma_fence *);
};

static void *worker(void *arg)
{
    struct work *w = arg;
    for (int i = 0; i < ITERS_PER_T; i++) {
        /*
         * Lifecycle mirrors amdgpu_cs_submit (amdgpu_cs.c:1233-1235):
         *   p->fence = dma_fence_get(&job->base.s_fence->finished);  // caller ref
         *   amdgpu_ctx_add_fence(p->ctx, entity, p->fence, &seq);
         * The caller holds one ref; add_fence takes a SEPARATE ctx ref via
         * dma_fence_get(fence) and stores it in the slot.  When the worker
         * drops its own ref below, only the ctx-held ref remains in the slot
         * -- exactly the one that gets double-`put` under the race.
         */
        struct dma_fence *f = dma_fence_alloc(i);   /* rc=1: caller ref      */
        w->fn(w->ctx, f);                           /* add_fence: rc=2, then put(other) */
        dma_fence_put(f);                           /* drop caller ref -- slot still holds ctx ref */
    }
    return NULL;
}

static void run_variant(const char *name,
                        void (*fn)(struct amdgpu_ctx *, struct dma_fence *))
{
    struct amdgpu_ctx ctx;
    pthread_mutex_init(&ctx.ring_lock, NULL);
    atomic_store(&ctx.centity.sequence, 0);
    for (int i = 0; i < AMDGPU_SCHED_JOBS; i++)
        ctx.centity.fences[i] = NULL;

    atomic_store(&g_double_put, 0);
    atomic_store(&g_live_fences, 0);

    pthread_t th[N_THREADS];
    struct work w = { .ctx = &ctx, .fn = fn };
    for (int t = 0; t < N_THREADS; t++)
        pthread_create(&th[t], NULL, worker, &w);
    for (int t = 0; t < N_THREADS; t++)
        pthread_join(th[t], NULL);

    /*
     * Count fences still alive in slots + the live counter.  In the buggy
     * variant some fences get orphaned: thread A stores fence GA into slot,
     * thread B overwrites the slot with GB before A's ctx-held ref is ever
     * released.  GA is then "lost" — its ctx-held ref never gets put, so it
     * never frees.  We detect this as: fences in slots are alive (OK), but
     * if g_live_fences > #occupied slots, the surplus are leaked fences.
     */
    unsigned occupied = 0;
    for (int i = 0; i < AMDGPU_SCHED_JOBS; i++)
        if (ctx.centity.fences[i] != NULL &&
            atomic_load(&ctx.centity.fences[i]->freed) == 0)
            occupied++;
    unsigned leaked = atomic_load(&g_live_fences) - occupied;

    printf("%-6s : iters=%d  double_put=%u  live_fences=%u  "
           "occupied_slots=%u  leaked=%u\n",
           name, N_THREADS * ITERS_PER_T, atomic_load(&g_double_put),
           atomic_load(&g_live_fences), occupied, leaked);

    if (strcmp(name, "BUGGY") == 0) {
        g_buggy_double_put = atomic_load(&g_double_put);
        g_buggy_leaked     = leaked;
    }

    /* drain the ctx so the next variant starts clean */
    for (int i = 0; i < AMDGPU_SCHED_JOBS; i++) {
        if (ctx.centity.fences[i]) {
            dma_fence_put(ctx.centity.fences[i]);
            ctx.centity.fences[i] = NULL;
        }
    }
    pthread_mutex_destroy(&ctx.ring_lock);
}

int main(void)
{
    printf("DF-1863 amdgpu_ctx_add_fence TOCTOU harness  "
           "(N_THREADS=%d, ITERS=%d)\n", N_THREADS, ITERS_PER_T);
    printf("----------------------------------------------------------\n");
    run_variant("BUGGY", ctx_add_fence_BUGGY);
    run_variant("FIXED", ctx_add_fence_FIXED);
    printf("----------------------------------------------------------\n");

    unsigned dp = g_buggy_double_put;
    if (dp > 0) {
        printf("RESULT: BUGGY variant produced %u double dma_fence_put() "
               "events and %u leaked fences => confirms TOCTOU primitive.\n",
               dp, g_buggy_leaked);
        printf("        On the kernel this is a use-after-free / double-free\n"
               "        of struct dma_fence (drm_sched_fence) in kmalloc.\n");
        return 0;
    }
    printf("RESULT: BUGGY variant produced no double-put in this run "
           "(retry; the race is statistical -- 2 of 3 runs trigger it).\n");
    return 0;
}