/*
 * pdb_scan_gmmu.c - read-only scanner for the GP10X/VER2 GMMU page-table
 * chain that maps a GPU VA to a target PA inside a live QEMU pc.ram file.
 *
 * IMPORTANT: the GPU page tables are NOT x86-64.  The format (gp10x family,
 * Pascal..Ampere) is:
 *
 *   PD3[48:47] -> PD2[46:38] -> PD1[37:29] -> PD0[28:21]
 *                  PD0 is a 16-byte dual entry: big half at +0, small at +8
 *   PT[20:12]   (4 KiB leaf)
 *
 * Entry address field: (entry >> 8) << 12 (aperture-dependent mask).
 * Directory entries: ap bits 2:1 = 0 invalid, 1 vidmem, 2 sys-coherent,
 * 3 sys-noncoherent.  Leaf PTEs: 0 vidmem, 1 peer (invalid here), 2/3 sys.
 * A directory entry has bit0 = IS_PTE = 0; a leaf has bit0 = VALID = 1.
 *
 * Also important: for this QEMU q35 guest the RAMBlock file offset is NOT
 * equal to GPA above 4 GiB.  Low RAM occupies GPA [0, low_ram) at file
 * [0, low_ram); high RAM starts at GPA 4 GiB and maps to file offset
 * low_ram (e820: low usable ends 0x7ffdf000, so low_ram = 0x80000000).
 * Hence: file_off = GPA if GPA < low_ram, else GPA - 4G + low_ram.
 *
 * Build:
 *   cc -O2 -Wall -Wextra -o pdb_scan_gmmu pdb_scan_gmmu.c
 * Run:
 *   sudo ./pdb_scan_gmmu /dev/shm/vgpu-guest-ram 0xafef8000 0x4100000
 *   sudo ./pdb_scan_gmmu /dev/shm/vgpu-guest-ram 0xa6c000 0x12006c000
 *
 * The tool is strictly read-only (PROT_READ/MAP_PRIVATE, no writes).
 */
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/mman.h>
#include <sys/stat.h>

#define SYS_MASK  0x00003FFFFFFFFFFFULL
#define VID_MASK  0x0000000001FFFFFFULL
#define LOW_RAM   0x80000000ULL      /* below-4G RAM block (override arg 4) */
#define FOUR_G    0x100000000ULL

static uint64_t va, target, low_ram = LOW_RAM;
static uint64_t fs, np;
static const uint8_t *ram;
static unsigned i_pd3, i_pd2, i_pd1, i_pd0, i_pt;

static inline uint64_t gpa_to_off(uint64_t gpa)
{
    if (gpa < low_ram) return gpa;
    if (gpa >= FOUR_G && gpa - FOUR_G < fs - low_ram)
        return gpa - FOUR_G + low_ram;
    return UINT64_MAX;
}
static inline uint64_t off_to_gpa(uint64_t off)
{
    if (off < low_ram) return off;
    return off + (FOUR_G - low_ram);
}
static inline uint64_t rd64(uint64_t gpa)
{
    uint64_t o = gpa_to_off(gpa), v;
    if (o == UINT64_MAX || o + 8 > fs) return UINT64_MAX;
    memcpy(&v, ram + o, 8);
    return v;
}
static inline unsigned apof(uint64_t e) { return (unsigned)((e >> 1) & 3); }
static inline uint64_t ent_addr(uint64_t e, unsigned ap)
{
    return ((e >> 8) & (ap == 1 ? VID_MASK : SYS_MASK)) << 12;
}
/* Directory entry -> child table.  ap 0 invalid; ap 1 vid; ap 2/3 sys. */
static int dir_child(uint64_t e, uint64_t *child)
{
    unsigned ap;
    if (e == UINT64_MAX || (e & 1)) return -1;
    ap = apof(e);
    if (ap == 0) return -1;
    *child = ent_addr(e, ap);
    return 0;
}
/* Leaf -> backing PA.  bit0 valid; ap 1 (peer) is not backed here. */
static int leaf_pa(uint64_t e, uint64_t *pa)
{
    unsigned ap;
    if (e == UINT64_MAX || !(e & 1)) return -1;
    ap = apof(e);
    if (ap == 1) return -1;
    *pa = ent_addr(e, ap);
    return 0;
}
static const char *apname(unsigned a)
{
    static const char *n[] = {"invalid", "vid", "sys-coh", "sys-nc"};
    return a < 4 ? n[a] : "?";
}
static void dump_entry(const char *lvl, uint64_t table, unsigned idx, uint64_t e)
{
    printf("  %-4s table=0x%-10llx idx=0x%-3x e=0x%016llx  V=%d ap=%u(%s)\n",
           lvl, (unsigned long long)table, idx, (unsigned long long)e,
           (int)(e & 1), apof(e), apname(apof(e)));
}

/* Full walk.  raw[0..4]=PD3,PD2,PD1,PD0-small,PT; big64/64k handled. */
static uint64_t walk(uint64_t pdb, uint64_t *raw, int *kind)
{
    uint64_t table = pdb, e, child;
    unsigned shifts[3] = {47, 38, 29}, idx;
    int l;

    *kind = 0; /* 0=4K, 1=2M big, 2=64K */
    for (l = 0; l < 3; l++) {
        idx = (unsigned)((va >> shifts[l]) & 0x1ff);
        if (l == 0) idx &= 3;               /* 4 root entries */
        e = rd64(table + 8 * idx);
        raw[l] = e;
        if (dir_child(e, &child)) return UINT64_MAX;
        table = child;
        if (gpa_to_off(table) == UINT64_MAX) return UINT64_MAX;
    }
    idx = (unsigned)((va >> 21) & 0xff);    /* PD0 is 8 bits */
    {
        uint64_t big = rd64(table + 16 * idx);
        uint64_t small = rd64(table + 16 * idx + 8);
        uint64_t pa, t64;
        raw[3] = small;
        if (big == UINT64_MAX || small == UINT64_MAX) return UINT64_MAX;
        /* big half: 2 MiB leaf (bit0=1) or 64 KiB table pointer (bit0=0) */
        if ((big & 1) && !leaf_pa(big, &pa)) {
            *kind = 1;
            return pa | (va & ((1ULL << 21) - 1));
        }
        if (!dir_child(big, &t64) && gpa_to_off(t64) != UINT64_MAX) {
            uint64_t e64 = rd64(t64 + 8 * ((va >> 16) & 0x1f));
            if (!leaf_pa(e64, &pa)) { *kind = 2; return pa | (va & 0xffff); }
        }
        /* small half: 4 KiB page table */
        if (dir_child(small, &child)) return UINT64_MAX;
        table = child;
        if (gpa_to_off(table) == UINT64_MAX) return UINT64_MAX;
    }
    idx = (unsigned)((va >> 12) & 0x1ff);
    {
        uint64_t pa;
        e = rd64(table + 8 * idx);
        raw[4] = e;
        if (leaf_pa(e, &pa)) return UINT64_MAX;
        return pa | (va & 0xfff);
    }
}

/* Reverse chain: PT -> PD0 -> PD1 -> PD2 -> PD3, keyed by GPA page. */
static void reverse_scan(void)
{
    uint64_t max_gpa_page = (off_to_gpa(fs) + 0xfff) >> 12;
    int32_t *head0 = malloc((max_gpa_page + 1) * 4);
    int32_t *next0 = malloc(np * 4);
    int32_t *headpd = malloc((max_gpa_page + 1) * 4);
    int32_t *nextpd = malloc(np * 4);
    int32_t *pt_list = malloc(np * 4);
    uint64_t i, n_pt = 0, n_chain = 0;

    if (!head0 || !next0 || !headpd || !nextpd || !pt_list) {
        fprintf(stderr, "reverse scan: out of memory\n");
        return;
    }
    for (i = 0; i <= max_gpa_page; i++) { head0[i] = -1; headpd[i] = -1; }

    for (i = 0; i < np; i++) {
        const uint8_t *p = ram + (i << 12);
        uint64_t e0, epd, ept, child, pa;
        memcpy(&e0,  p + 0, 8);
        memcpy(&epd, p + 16 * i_pd0 + 8, 8);
        memcpy(&ept, p + 8 * i_pt, 8);

        if (!leaf_pa(ept, &pa) && pa == target)
            pt_list[n_pt++] = (int32_t)i;

        if (!dir_child(e0, &child) && (child >> 12) <= max_gpa_page) {
            next0[i] = head0[child >> 12];
            head0[child >> 12] = (int32_t)i;
        }
        if (!dir_child(epd, &child) && (child >> 12) <= max_gpa_page) {
            nextpd[i] = headpd[child >> 12];
            headpd[child >> 12] = (int32_t)i;
        }
    }
    printf("\n[reverse scan] PT pages with leaf[0x%x] -> 0x%llx : %llu\n",
           i_pt, (unsigned long long)target, (unsigned long long)n_pt);

    for (uint64_t t = 0; t < n_pt; t++) {
        int32_t pt = pt_list[t], pd0, pd1, pd2, pd3;
        uint64_t pt_gp = off_to_gpa((uint64_t)pt << 12) >> 12;
        for (pd0 = headpd[pt_gp]; pd0 >= 0; pd0 = nextpd[pd0]) {
            uint64_t pd0_gp = off_to_gpa((uint64_t)pd0 << 12) >> 12;
            for (pd1 = head0[pd0_gp]; pd1 >= 0; pd1 = next0[pd1]) {
                uint64_t pd1_gp = off_to_gpa((uint64_t)pd1 << 12) >> 12;
                for (pd2 = head0[pd1_gp]; pd2 >= 0; pd2 = next0[pd2]) {
                    uint64_t pd2_gp = off_to_gpa((uint64_t)pd2 << 12) >> 12;
                    for (pd3 = head0[pd2_gp]; pd3 >= 0; pd3 = next0[pd3]) {
                        uint64_t raw[5] = {0}, fin;
                        int kind;
                        fin = walk(off_to_gpa((uint64_t)pd3 << 12), raw, &kind);
                        if (fin != target) continue;
                        n_chain++;
                        printf("\n=== reverse chain #%llu (pdb GPA 0x%llx) ===\n",
                               (unsigned long long)n_chain,
                               (unsigned long long)off_to_gpa((uint64_t)pd3 << 12));
                        dump_entry("PD3", off_to_gpa((uint64_t)pd3 << 12), i_pd3, raw[0]);
                        dump_entry("PD2", ent_addr(raw[0], apof(raw[0])), i_pd2, raw[1]);
                        dump_entry("PD1", ent_addr(raw[1], apof(raw[1])), i_pd1, raw[2]);
                        dump_entry("PD0", ent_addr(raw[2], apof(raw[2])), i_pd0, raw[3]);
                        dump_entry("PT",  ent_addr(raw[3], apof(raw[3])), i_pt, raw[4]);
                        printf("  walk result = 0x%llx (%s)\n",
                               (unsigned long long)fin,
                               kind == 1 ? "2M big" : kind == 2 ? "64K" : "4K");
                    }
                }
            }
        }
    }
    printf("[reverse scan] matching chains: %llu\n", (unsigned long long)n_chain);
    free(head0); free(next0); free(headpd); free(nextpd); free(pt_list);
}

typedef struct { uint64_t pa, pdb; int kind; } Hit;
static int hit_cmp(const void *a, const void *b)
{
    const Hit *x = a, *y = b;
    return x->pa < y->pa ? -1 : x->pa > y->pa ? 1 : 0;
}

/* Brute force: treat every 4K page as a candidate pdb, walk the VA, tally. */
static void brute_force(void)
{
    Hit *h = malloc(np * sizeof(Hit));
    uint64_t i, nh = 0, plausible = 0, matches = 0;
    if (!h) { fprintf(stderr, "brute force: oom\n"); return; }

    for (i = 0; i < np; i++) {
        uint64_t pdb = off_to_gpa(i << 12), e0, child, fin;
        int kind;
        uint64_t raw[5] = {0};
        e0 = rd64(pdb);
        if (dir_child(e0, &child) || gpa_to_off(child) == UINT64_MAX) continue;
        plausible++;
        fin = walk(pdb, raw, &kind);
        if (fin == UINT64_MAX) continue;
        h[nh].pa = fin; h[nh].pdb = pdb; h[nh].kind = kind; nh++;
        if (fin == target) {
            matches++;
            printf("\n=== brute-force MATCH pdb GPA 0x%llx -> 0x%llx (%s) ===\n",
                   (unsigned long long)pdb, (unsigned long long)fin,
                   kind == 1 ? "2M big" : kind == 2 ? "64K" : "4K");
            dump_entry("PD3", pdb, i_pd3, raw[0]);
            dump_entry("PD2", ent_addr(raw[0], apof(raw[0])), i_pd2, raw[1]);
            dump_entry("PD1", ent_addr(raw[1], apof(raw[1])), i_pd1, raw[2]);
            dump_entry("PD0", ent_addr(raw[2], apof(raw[2])), i_pd0, raw[3]);
            dump_entry("PT",  ent_addr(raw[3], apof(raw[3])), i_pt, raw[4]);
        }
    }
    qsort(h, nh, sizeof(Hit), hit_cmp);
    printf("\n[brute force] plausible pdbs=%llu valid walks=%llu target matches=%llu\n",
           (unsigned long long)plausible, (unsigned long long)nh,
           (unsigned long long)matches);
    i = 0;
    while (i < nh) {
        uint64_t j = i;
        while (j < nh && h[j].pa == h[i].pa) j++;
        printf("  VA 0x%llx -> PA 0x%llx : %llu pdb(s), %s\n",
               (unsigned long long)va, (unsigned long long)h[i].pa,
               (unsigned long long)(j - i),
               h[i].kind == 1 ? "2M big" : h[i].kind == 2 ? "64K" : "4K");
        for (uint64_t k = i; k < j && k < i + 8; k++)
            printf("      pdb 0x%llx\n", (unsigned long long)h[k].pdb);
        if (j - i > 8) printf("      ...\n");
        i = j;
    }
    free(h);
}

int main(int argc, char **argv)
{
    const char *path;
    int fd;
    struct stat st;

    if (argc < 3) {
        fprintf(stderr,
            "usage: %s <ramfile> <target_pa> [va] [low_ram_bytes]\n"
            "  default va=0x4100000 low_ram=0x80000000\n", argv[0]);
        return 2;
    }
    path = argv[1];
    target = strtoull(argv[2], NULL, 0);
    va = (argc >= 4) ? strtoull(argv[3], NULL, 0) : 0x4100000ULL;
    if (argc >= 5) low_ram = strtoull(argv[4], NULL, 0);

    i_pd3 = (unsigned)((va >> 47) & 3);
    i_pd2 = (unsigned)((va >> 38) & 0x1ff);
    i_pd1 = (unsigned)((va >> 29) & 0x1ff);
    i_pd0 = (unsigned)((va >> 21) & 0xff);
    i_pt  = (unsigned)((va >> 12) & 0x1ff);

    fd = open(path, O_RDONLY);
    if (fd < 0) { perror("open"); return 1; }
    if (fstat(fd, &st) < 0) { perror("fstat"); return 1; }
    fs = (uint64_t)st.st_size;
    np = fs >> 12;
    ram = mmap(NULL, (size_t)fs, PROT_READ, MAP_PRIVATE, fd, 0);
    if (ram == MAP_FAILED) { perror("mmap"); return 1; }

    printf("file=%s size=%llu pages=%llu\n", path,
           (unsigned long long)fs, (unsigned long long)np);
    printf("target_pa=0x%llx va=0x%llx low_ram=0x%llx\n",
           (unsigned long long)target, (unsigned long long)va,
           (unsigned long long)low_ram);
    printf("GMMU idx: PD3[%u] PD2[%u] PD1[%u] PD0[%u] PT[%u]\n",
           i_pd3, i_pd2, i_pd1, i_pd0, i_pt);

    reverse_scan();
    brute_force();

    munmap((void *)ram, (size_t)fs);
    close(fd);
    return 0;
}