/*
 * pdb_scan.c - read-only reverse scan of a QEMU pc.ram backing file for the
 * 4-level page-table chain that maps a given VA to a given PA.
 *
 * Build:
 *   cc -O2 -Wall -Wextra -o pdb_scan pdb_scan.c
 *
 * Usage:
 *   sudo ./pdb_scan /dev/shm/vgpu-guest-ram 0xafef8000 0x4100000
 *
 * Arguments: <ramfile> <target_pa> [va]
 *   - target_pa : physical address the leaf PT entry must point at (page aligned)
 *   - va        : optional VA used to derive the four table indices
 *                 (default 0x4100000: PML4[0], PDPT[0], PD[8], PT[0x100])
 *
 * The scan is strictly read-only: the file is mmap()ed PROT_READ/MAP_PRIVATE
 * and no writes are issued.  It builds two reverse indexes:
 *   entry[0x000] -> pages (used for PML4[0] and PDPT[0])
 *   entry[0x040] -> pages (used for PD[8])
 * then joins them against pages whose entry[0x800] points at target_pa.
 */
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/mman.h>
#include <sys/stat.h>

#define PAGE_SHIFT 12ULL
#define PAGE_SIZE  (1ULL << PAGE_SHIFT)
#define PA_MASK    0x000FFFFFFFFFF000ULL

static uint64_t target_pa;
static uint64_t va;
static uint64_t n_pages;
static const uint8_t *ram;

static unsigned idx_pml4, idx_pdpt, idx_pd, idx_pt;
static const char *idx_names[4] = {"PML4", "PDPT", "PD", "PT"};

static inline uint64_t rd64(uint64_t pa)
{
    uint64_t v;
    memcpy(&v, ram + pa, sizeof(v));
    return v;
}

static inline uint64_t entry_mask(uint64_t e) { return e & PA_MASK; }

static const char *bits(uint64_t e)
{
    static char b[128];
    snprintf(b, sizeof(b), "P=%d RW=%d US=%d PWT=%d PCD=%d A=%d D=%d PS=%d",
             (int)(e & 1), (int)((e >> 1) & 1), (int)((e >> 2) & 1),
             (int)((e >> 3) & 1), (int)((e >> 4) & 1), (int)((e >> 5) & 1),
             (int)((e >> 6) & 1), (int)((e >> 7) & 1));
    return b;
}

/* Generic 4-level walk; returns final PA (masked) or UINT64_MAX on break. */
static uint64_t walk4(uint64_t pml4_pa, uint64_t *raw_out)
{
    uint64_t e[4], pa = pml4_pa;
    unsigned idx[4] = {idx_pml4, idx_pdpt, idx_pd, idx_pt};
    int l;
    for (l = 0; l < 4; l++) {
        if (pa + PAGE_SIZE > (n_pages << PAGE_SHIFT))
            return UINT64_MAX;
        e[l] = rd64(pa + (uint64_t)idx[l] * 8);
        if (raw_out) raw_out[l] = e[l];
        if (!(e[l] & 1))
            return UINT64_MAX;
        if (l == 2 && (e[l] & 0x80)) { /* PS=1: 2 MiB page at PD level */
            uint64_t base = entry_mask(e[l]) & ~((1ULL << 21) - 1);
            return base | (va & ((1ULL << 21) - 1));
        }
        pa = entry_mask(e[l]);
    }
    return pa;
}

int main(int argc, char **argv)
{
    const char *path;
    int fd;
    struct stat st;
    uint64_t i;
    int32_t *head0, *next0, *head40, *next40, *pt_list;
    uint64_t n_pt = 0;
    int have_va = 0;

    if (argc < 3) {
        fprintf(stderr, "usage: %s <ramfile> <target_pa> [va]\n", argv[0]);
        return 2;
    }
    path = argv[1];
    target_pa = strtoull(argv[2], NULL, 0) & PA_MASK;
    if (argc >= 4) {
        va = strtoull(argv[3], NULL, 0);
        have_va = 1;
    } else {
        va = 0x4100000ULL;
    }
    /* indices: canonical 4-level split */
    idx_pml4 = (unsigned)((va >> 39) & 0x1ff);
    idx_pdpt = (unsigned)((va >> 30) & 0x1ff);
    idx_pd   = (unsigned)((va >> 21) & 0x1ff);
    idx_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; }
    if (st.st_size == 0) { fprintf(stderr, "empty file\n"); return 1; }
    n_pages = (uint64_t)st.st_size >> PAGE_SHIFT;

    ram = mmap(NULL, (size_t)st.st_size, PROT_READ, MAP_PRIVATE, fd, 0);
    if (ram == MAP_FAILED) { perror("mmap"); return 1; }
    madvise((void *)ram, (size_t)st.st_size, MADV_SEQUENTIAL);

    fprintf(stderr, "ramfile=%s size=%llu bytes pages=%llu\n",
            path, (unsigned long long)st.st_size, (unsigned long long)n_pages);
    fprintf(stderr, "target_pa=0x%llx va=0x%llx idx: PML4[%u] PDPT[%u] PD[%u] PT[%u]\n",
            (unsigned long long)target_pa, (unsigned long long)va,
            idx_pml4, idx_pdpt, idx_pd, idx_pt);
    fprintf(stderr, "PT entry offset=0x%x PD entry offset=0x%x\n",
            idx_pt * 8, idx_pd * 8);

    head0  = malloc(n_pages * sizeof(int32_t));
    next0  = malloc(n_pages * sizeof(int32_t));
    head40 = malloc(n_pages * sizeof(int32_t));
    next40 = malloc(n_pages * sizeof(int32_t));
    pt_list = malloc(n_pages * sizeof(int32_t));
    if (!head0 || !next0 || !head40 || !next40 || !pt_list) {
        fprintf(stderr, "out of memory\n"); return 1;
    }
    for (i = 0; i < n_pages; i++) { head0[i] = -1; head40[i] = -1; }

    for (i = 0; i < n_pages; i++) {
        const uint8_t *p = ram + (i << PAGE_SHIFT);
        uint64_t e0, e40, ept, a;
        memcpy(&e0,  p + 0x000, 8);
        memcpy(&e40, p + 0x040, 8);
        memcpy(&ept, p + (uint64_t)idx_pt * 8, 8);

        if (entry_mask(ept) == target_pa)
            pt_list[n_pt++] = (int32_t)i;

        a = entry_mask(e0);
        if (a < (n_pages << PAGE_SHIFT) && (a & (PAGE_SIZE - 1)) == 0) {
            uint64_t j = a >> PAGE_SHIFT;
            next0[i] = head0[j];
            head0[j] = (int32_t)i;
        }
        a = entry_mask(e40);
        if (a < (n_pages << PAGE_SHIFT) && (a & (PAGE_SIZE - 1)) == 0) {
            uint64_t j = a >> PAGE_SHIFT;
            next40[i] = head40[j];
            head40[j] = (int32_t)i;
        }
    }
    madvise((void *)ram, (size_t)st.st_size, MADV_RANDOM);

    fprintf(stderr, "pages whose entry[0x%x] == target: %llu\n",
            idx_pt * 8, (unsigned long long)n_pt);

    uint64_t n_chain = 0;
    for (uint64_t t = 0; t < n_pt; t++) {
        int32_t pt = pt_list[t];
        int32_t pd;
        for (pd = head40[pt]; pd >= 0; pd = next40[pd]) {
            int32_t pdpt;
            for (pdpt = head0[pd]; pdpt >= 0; pdpt = next0[pdpt]) {
                int32_t pml4;
                for (pml4 = head0[pdpt]; pml4 >= 0; pml4 = next0[pml4]) {
                    uint64_t pa_pml4 = (uint64_t)pml4 << PAGE_SHIFT;
                    uint64_t pa_pdpt = (uint64_t)pdpt << PAGE_SHIFT;
                    uint64_t pa_pd   = (uint64_t)pd   << PAGE_SHIFT;
                    uint64_t pa_pt   = (uint64_t)pt   << PAGE_SHIFT;
                    uint64_t raw[4] = {0}, fin;
                    fin = walk4(pa_pml4, raw);
                    n_chain++;
                    printf("\n=== candidate #%llu ===\n", (unsigned long long)n_chain);
                    printf("pdb (PML4 PA) = 0x%llx\n", (unsigned long long)pa_pml4);
                    printf("PML4[%u]  @0x%llx+0x%03x = 0x%016llx -> PDPT 0x%llx  [%s]\n",
                           idx_pml4, (unsigned long long)pa_pml4, idx_pml4 * 8,
                           (unsigned long long)raw[0], (unsigned long long)pa_pdpt,
                           bits(raw[0]));
                    printf("PDPT[%u] @0x%llx+0x%03x = 0x%016llx -> PD   0x%llx  [%s]\n",
                           idx_pdpt, (unsigned long long)pa_pdpt, idx_pdpt * 8,
                           (unsigned long long)raw[1], (unsigned long long)pa_pd,
                           bits(raw[1]));
                    printf("PD[%u]   @0x%llx+0x%03x = 0x%016llx -> PT   0x%llx  [%s]\n",
                           idx_pd, (unsigned long long)pa_pd, idx_pd * 8,
                           (unsigned long long)raw[2], (unsigned long long)pa_pt,
                           bits(raw[2]));
                    printf("PT[%u]   @0x%llx+0x%03x = 0x%016llx -> PA   0x%llx  [%s]\n",
                           idx_pt, (unsigned long long)pa_pt, idx_pt * 8,
                           (unsigned long long)raw[3],
                           (unsigned long long)entry_mask(raw[3]), bits(raw[3]));
                    printf("walk_result = 0x%llx  %s\n",
                           (unsigned long long)fin,
                           fin == target_pa ? "MATCH" : "MISMATCH");
                }
            }
        }
    }

    /* Reference: the caller-supplied known pdb, if present in the file. */
    if (have_va) {
        uint64_t known = 0x11092a000ULL;
        if (known < (n_pages << PAGE_SHIFT)) {
            uint64_t raw[4] = {0}, fin = walk4(known, raw);
            printf("\n=== reference walk from known pdb 0x%llx ===\n",
                   (unsigned long long)known);
            printf("PML4[%u]=0x%016llx PDPT[%u]=0x%016llx PD[%u]=0x%016llx PT[%u]=0x%016llx\n",
                   idx_pml4, (unsigned long long)raw[0],
                   idx_pdpt, (unsigned long long)raw[1],
                   idx_pd, (unsigned long long)raw[2],
                   idx_pt, (unsigned long long)raw[3]);
            printf("walk_result = 0x%llx\n", (unsigned long long)fin);
        }
    }

    printf("\ntotal candidate chains: %llu\n", (unsigned long long)n_chain);

    munmap((void *)ram, (size_t)st.st_size);
    close(fd);
    free(head0); free(next0); free(head40); free(next40); free(pt_list);
    return 0;
}