/*
 * DF-1927 — UAF race in drm_sched_entity_fini via broken kthread_park
 *
 * Faithful pthread model of the concurrent access to a drm_sched_entity's
 * spsc_queue job_queue and last_scheduled fence between:
 *
 *   (a) the scheduler kthread running drm_sched_entity_pop_job()
 *       (sys/dev/drm/scheduler/sched_entity.c:430-456)
 *   (b) the closer thread running drm_sched_entity_fini()
 *       (sys/dev/drm/scheduler/sched_entity.c:263-292)
 *
 * The barrier separating them is the kthread_park/kthread_unpark pair at
 * sched_entity.c:277-278.  In DragonFly's LinuxKPI shim (linux_kthread.c)
 * kthread_park() at line 104-110 does set_bit + wake_up_process + RETURN --
 * it does NOT wait for the target thread to actually park, unlike upstream
 * Linux which blocks on wait_for_completion(&k->parked).
 *
 * MODEL:
 *   - Scheduler thread runs drm_sched_main's loop: each iteration, optionally
 *     parkme, then pop_job (peek + put last_scheduled + assign + pop).
 *   - Closer thread runs entity_fini: peek + park + unpark + kill_jobs + put.
 *   - A pthread barrier synchronizes the START of both threads so they begin
 *     the race simultaneously.  After that, they run freely.
 *
 * In the BUGGY build (default), kthread_park returns immediately (with a
 * sched_yield to model wake_up_process), so fini's kill_jobs + dma_fence_put
 * run concurrently with the scheduler's pop_job -> Race A (double-pop /
 * job UAF) and Race B (double dma_fence_put / fence UAF) fire.
 *
 * In the FIXED build (-DFIXED_KTHREAD_PARK), kthread_park blocks on a
 * condition variable until the scheduler has reached parkme -- the scheduler
 * therefore completes its current pop_job and parks at drm_sched_blocked
 * (sched_main.c:514) before fini can proceed.  No concurrent access.
 *
 * BUILD:
 *   cc -O2 -g -Wall -pthread -o df1927_race df1927_race.c
 *   cc -O2 -g -Wall -pthread -DFIXED_KTHREAD_PARK -o df1927_race_fixed \
 *         df1927_race.c
 *
 * RUN:
 *   ./df1927_race         # expect: race fires in most iterations, exit 0
 *   ./df1927_race_fixed   # expect: race NEVER fires,        exit 0
 */

#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdatomic.h>
#include <stdint.h>
#include <stdbool.h>
#include <unistd.h>
#include <pthread.h>
#include <sched.h>
#include <time.h>

/* ------------------------------------------------------------------ */
/* spsc_queue (faithful copy of sys/dev/drm/include/drm/spsc_queue.h)   */
/* ------------------------------------------------------------------ */

struct spsc_node { struct spsc_node *next; };
struct spsc_queue {
    struct spsc_node *head;
    struct spsc_node *tail;
};
static inline void spsc_queue_init(struct spsc_queue *q) {
    q->head = NULL;
    q->tail = (struct spsc_node *)&q->head;
}
static inline struct spsc_node *spsc_queue_peek(struct spsc_queue *q) {
    return q->head;
}
static inline bool spsc_queue_push(struct spsc_queue *q, struct spsc_node *n) {
    struct spsc_node **slot = (struct spsc_node **)q->tail;
    n->next = NULL;
    *slot = n;
    q->tail = (struct spsc_node *)&n->next;
    return slot == &q->head;
}
static inline struct spsc_node *spsc_queue_pop(struct spsc_queue *q) {
    struct spsc_node *n = q->head;
    if (!n) return NULL;
    q->head = n->next;
    if (!q->head) q->tail = (struct spsc_node *)&q->head;
    return n;
}

/* ------------------------------------------------------------------ */
/* dma_fence refcount model                                             */
/* ------------------------------------------------------------------ */

struct dma_fence {
    _Atomic int refcount;
    int id;
    _Atomic int freed;            /* detects use-after-free */
};
static inline int dma_fence_put(struct dma_fence *f) {
    int prev = atomic_fetch_sub(&f->refcount, 1);
    if (prev == 1) atomic_store(&f->freed, 1);   /* last ref -> freed */
    return prev;
}

/* ------------------------------------------------------------------ */
/* drm_sched_job + entity                                               */
/* ------------------------------------------------------------------ */

struct drm_sched_job {
    struct spsc_node queue_node;
    int id;
    _Atomic int freed;
};

struct drm_sched_entity {
    struct spsc_queue job_queue;
    struct dma_fence *last_scheduled;
};

/* ------------------------------------------------------------------ */
/* task_struct + kthread_park model                                     */
/* ------------------------------------------------------------------ */

struct task_struct {
    _Atomic unsigned long kt_flags;
    pthread_mutex_t lock;
    pthread_cond_t  parked_cv;
    pthread_cond_t  unpark_cv;
    _Atomic int parked;
};
#define KTHREAD_SHOULD_PARK 2

static inline void set_bit(unsigned long nr, _Atomic unsigned long *a) {
    atomic_fetch_or(a, (1UL << nr));
}
static inline void clear_bit(unsigned long nr, _Atomic unsigned long *a) {
    atomic_fetch_and(a, ~(1UL << nr));
}
static inline bool test_bit(unsigned long nr, _Atomic unsigned long *a) {
    return (atomic_load(a) >> nr) & 1UL;
}

#ifdef FIXED_KTHREAD_PARK
/* FIXED: synchronous -- blocks until target reaches parkme (upstream Linux) */
static int kthread_park(struct task_struct *ts) {
    set_bit(KTHREAD_SHOULD_PARK, &ts->kt_flags);
    pthread_mutex_lock(&ts->lock);
    pthread_cond_broadcast(&ts->unpark_cv);
    while (!atomic_load(&ts->parked))
        pthread_cond_wait(&ts->parked_cv, &ts->lock);
    pthread_mutex_unlock(&ts->lock);
    return 0;
}
static void kthread_unpark(struct task_struct *ts) {
    clear_bit(KTHREAD_SHOULD_PARK, &ts->kt_flags);
    pthread_mutex_lock(&ts->lock);
    atomic_store(&ts->parked, 0);
    pthread_cond_broadcast(&ts->unpark_cv);
    pthread_mutex_unlock(&ts->lock);
}
static void kthread_parkme(struct task_struct *ts) {
    if (!test_bit(KTHREAD_SHOULD_PARK, &ts->kt_flags)) return;
    pthread_mutex_lock(&ts->lock);
    while (test_bit(KTHREAD_SHOULD_PARK, &ts->kt_flags)) {
        atomic_store(&ts->parked, 1);
        pthread_cond_signal(&ts->parked_cv);
        pthread_cond_wait(&ts->unpark_cv, &ts->lock);
    }
    atomic_store(&ts->parked, 0);
    pthread_mutex_unlock(&ts->lock);
}
#else
/* BUGGY: DFly master DEV linux_kthread.c:104-110 -- returns immediately */
static int kthread_park(struct task_struct *ts) {
    set_bit(KTHREAD_SHOULD_PARK, &ts->kt_flags);
    /* wake_up_process analogue: yield + brief sleep to give the target a
     * chance to run (modeling the wakeup IPI / lwkt_schedule).  Critically,
     * we do NOT wait for the target to actually park. */
    sched_yield();
    struct timespec ts25 = { .tv_sec = 0, .tv_nsec = 25 * 1000 };
    nanosleep(&ts25, NULL);
    return 0;
}
static void kthread_unpark(struct task_struct *ts) {
    clear_bit(KTHREAD_SHOULD_PARK, &ts->kt_flags);
    sched_yield();
}
static void kthread_parkme(struct task_struct *ts) {
    if (!test_bit(KTHREAD_SHOULD_PARK, &ts->kt_flags)) return;
    /* DFly: lwkt_deschedule_self once.  Modelled as a yield -- the bit
     * remains set but we don't block (matching the broken shim).  We will
     * loop back to drm_sched_blocked on the next iteration and re-check. */
    sched_yield();
}
#endif

/* ------------------------------------------------------------------ */
/* Race-detection counters (global)                                     */
/* ------------------------------------------------------------------ */

static _Atomic int race_a_double_free = 0;   /* spsc_queue double-pop -> job UAF */
static _Atomic int race_b_underflow   = 0;   /* dma_fence_put kref underflow -> fence UAF */
static _Atomic int race_b_use_after_free = 0;/* dma_fence_put on already-freed fence */

/* ------------------------------------------------------------------ */
/* State per race iteration                                             */
/* ------------------------------------------------------------------ */

static struct task_struct g_thread;
static struct drm_sched_entity g_entity;
static pthread_barrier_t   g_start_barrier;

/* ------------------------------------------------------------------ */
/* Scheduler thread: drm_sched_main loop + pop_job                      */
/* ------------------------------------------------------------------ */

static void *scheduler_main(void *arg) {
    int iters = (int)(intptr_t)arg;
    pthread_barrier_wait(&g_start_barrier);
    for (int i = 0; i < iters; i++) {
        /* sched_main.c:512 drm_sched_blocked */
        if (test_bit(KTHREAD_SHOULD_PARK, &g_thread.kt_flags))
            kthread_parkme(&g_thread);

        /* sched_entity.c:435 peek */
        struct drm_sched_job *sched_job =
            (struct drm_sched_job *)spsc_queue_peek(&g_entity.job_queue);
        if (!sched_job) {
            sched_yield();
            continue;
        }

        /* sched_entity.c:451-452 put + assign last_scheduled.
         * Yield between the put and the pop to widen the race window. */
        struct dma_fence *old = g_entity.last_scheduled;
        if (old) {               /* entity may have been finalized */
            int prev = dma_fence_put(old);
            if (prev <= 0) {
                atomic_fetch_add(&race_b_underflow, 1);
                if (atomic_load(&old->freed))
                    atomic_fetch_add(&race_b_use_after_free, 1);
            }
        }
        struct dma_fence *nf = malloc(sizeof(*nf));
        atomic_init(&nf->refcount, 1);
        atomic_init(&nf->freed, 0);
        nf->id = sched_job->id;
        g_entity.last_scheduled = nf;
        sched_yield();

        /* sched_entity.c:454 pop */
        struct spsc_node *popped = spsc_queue_pop(&g_entity.job_queue);
        if (popped) {
            if (atomic_exchange(&sched_job->freed, 1) != 0)
                atomic_fetch_add(&race_a_double_free, 1);
            free(sched_job);
        } else {
            /* sched_job was peeked but pop returned NULL: it was concurrently
             * popped by kill_jobs.  sched_job is stale -- UAF. */
            if (atomic_load(&sched_job->freed))
                atomic_fetch_add(&race_a_double_free, 1);
        }
    }
    return NULL;
}

/* ------------------------------------------------------------------ */
/* drm_sched_entity_kill_jobs (sched_entity.c:223)                      */
/* ------------------------------------------------------------------ */

static void kill_jobs(struct drm_sched_entity *entity) {
    struct drm_sched_job *job;
    while ((job = (struct drm_sched_job *)spsc_queue_pop(&entity->job_queue))) {
        if (atomic_exchange(&job->freed, 1) != 0)
            atomic_fetch_add(&race_a_double_free, 1);
        free(job);
    }
}

/* ------------------------------------------------------------------ */
/* drm_sched_entity_fini (sched_entity.c:263)                           */
/* ------------------------------------------------------------------ */

static void entity_fini(struct drm_sched_entity *entity) {
    if (spsc_queue_peek(&entity->job_queue)) {
        kthread_park(&g_thread);
        kthread_unpark(&g_thread);
        kill_jobs(entity);
    }
    struct dma_fence *f = entity->last_scheduled;
    int prev = dma_fence_put(f);
    if (prev <= 0) {
        atomic_fetch_add(&race_b_underflow, 1);
        if (atomic_load(&f->freed))
            atomic_fetch_add(&race_b_use_after_free, 1);
    }
    entity->last_scheduled = NULL;
}

static void *fini_thread(void *arg) {
    (void)arg;
    pthread_barrier_wait(&g_start_barrier);
    entity_fini(&g_entity);
    return NULL;
}

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

static int run_one(void) {
    atomic_store(&race_a_double_free, 0);
    atomic_store(&race_b_underflow, 0);
    atomic_store(&race_b_use_after_free, 0);

    spsc_queue_init(&g_entity.job_queue);
    g_entity.last_scheduled = malloc(sizeof(struct dma_fence));
    atomic_init(&g_entity.last_scheduled->refcount, 1);
    atomic_init(&g_entity.last_scheduled->freed, 0);
    g_entity.last_scheduled->id = -1;

    for (int i = 0; i < 4; i++) {
        struct drm_sched_job *job = malloc(sizeof(*job));
        job->id = i;
        atomic_init(&job->freed, 0);
        spsc_queue_push(&g_entity.job_queue, &job->queue_node);
    }

    atomic_init(&g_thread.kt_flags, 0);
    atomic_init(&g_thread.parked, 0);
    pthread_mutex_init(&g_thread.lock, NULL);
    pthread_cond_init(&g_thread.parked_cv, NULL);
    pthread_cond_init(&g_thread.unpark_cv, NULL);

    pthread_barrier_init(&g_start_barrier, NULL, 2);
    pthread_t t_sched, t_fini;
    pthread_create(&t_sched, NULL, scheduler_main, (void *)(intptr_t)6);
    pthread_create(&t_fini,  NULL, fini_thread,   NULL);

    pthread_join(t_sched, NULL);
    pthread_join(t_fini,  NULL);
    pthread_barrier_destroy(&g_start_barrier);

    int a = atomic_load(&race_a_double_free);
    int b = atomic_load(&race_b_underflow);
    int c = atomic_load(&race_b_use_after_free);
    if (a || b || c) return 1;
    return 0;
}

int main(void) {
    /* Pin both threads to one CPU to maximize interleaving (single CPU
     * round-robin = realistic model of an interrupt-driven preemption
     * between two threads on the same core). */
    cpu_set_t cs;
    CPU_ZERO(&cs);
    CPU_SET(0, &cs);
    sched_setaffinity(0, sizeof(cs), &cs);

    int trips = 0;
    int iters = 100;
    for (int i = 0; i < iters; i++) {
        int rc = run_one();
        if (rc) trips++;
    }
    int a_total = atomic_load(&race_a_double_free);
    int b_total = atomic_load(&race_b_underflow);
    int c_total = atomic_load(&race_b_use_after_free);
#ifdef FIXED_KTHREAD_PARK
    const char *variant = "FIXED_KTHREAD_PARK (synchronous, matches upstream Linux)";
#else
    const char *variant = "BUGGY (DFly master DEV linux_kthread.c:104-110)";
#endif
    (void)a_total; (void)b_total; (void)c_total;
    printf("==== %s ====\n", variant);
    printf("race tripped in %d/%d iterations\n", trips, iters);
    printf("  Race A (spsc_queue double-pop / job UAF):       %d events total\n", a_total);
    printf("  Race B (dma_fence_put kref underflow):          %d events total\n", b_total);
    printf("  Race B' (dma_fence_put use-after-free):         %d events total\n", c_total);
#ifdef FIXED_KTHREAD_PARK
    /* Fixed variant: race should NEVER fire */
    return (trips == 0 && a_total == 0 && b_total == 0 && c_total == 0) ? 0 : 2;
#else
    /* Buggy variant: race SHOULD fire */
    return (trips > 0) ? 0 : 2;
#endif
}