vatrog-vm-introspection-engine/src/engine/win32/pe.c

#include "pe.h"

#include <string.h>
#include "memmodel.h"   /* gva_read, VR_* */
#include "sigscan.h"    /* mem_sub (pure matcher; engine may use it) */
#include "win32.h"      /* public surface: vmie_win32, section_desc, view_base */

/* IMAGE_SECTION_HEADER: 8-byte Name, then Misc.VirtualSize(+8), VirtualAddress
 * (+12), and Characteristics(+36); the header is 40 bytes wide. */
#define SH_SIZE        40u
#define SH_VSIZE_OFF   8u
#define SH_VADDR_OFF   12u
#define SH_CHAR_OFF    36u

/* IMAGE_SCN_MEM_* protection bits in IMAGE_SECTION_HEADER.Characteristics. */
#define SCN_MEM_EXECUTE 0x20000000u
#define SCN_MEM_READ    0x40000000u
#define SCN_MEM_WRITE   0x80000000u

/* Common PE section-table walk: validate the MZ/PE headers reachable inside `v`
 * at `module_base`, then locate the section-header array. On success fills
 * *sec_off (byte offset into v.data of the first IMAGE_SECTION_HEADER) and
 * *nsec (NumberOfSections), and returns true. This is the single header parse
 * shared by pe_find_section and pe_sections - there is no second PE parser. */
static bool pe_section_table(mem_view_t v, uint64_t module_base,
                             size_t* sec_off, uint16_t* nsec) {
    if (!v.data || module_base < v.base_va) return false;
    const size_t mo = (size_t)(module_base - v.base_va);
    if (mo + 0x40 > v.size) return false;
    if (v.data[mo] != 'M' || v.data[mo + 1] != 'Z') return false;

    int32_t e_lfanew;
    memcpy(&e_lfanew, v.data + mo + 0x3C, 4);
    const size_t nt = mo + (size_t)(uint32_t)e_lfanew;
    if (nt + 0x18 > v.size) return false;
    if (memcmp(v.data + nt, "PE\0\0", 4) != 0) return false;

    uint16_t n, opt_size;
    memcpy(&n,        v.data + nt + 6,  2);        /* NumberOfSections      */
    memcpy(&opt_size, v.data + nt + 20, 2);        /* SizeOfOptionalHeader  */

    *sec_off = nt + 24 + opt_size;                 /* first section header  */
    *nsec    = n;
    return true;
}

/* Map IMAGE_SCN_MEM_READ/WRITE/EXECUTE -> VR_R/VR_W/VR_X. Image semantics, not
 * live PTEs: VR_U is never set here (see section_desc.prot in win32.h). */
static uint32_t pe_prot(uint32_t characteristics) {
    uint32_t prot = 0;
    if (characteristics & SCN_MEM_READ)    { prot |= VR_R; }
    if (characteristics & SCN_MEM_WRITE)   { prot |= VR_W; }
    if (characteristics & SCN_MEM_EXECUTE) { prot |= VR_X; }
    return prot;
}

bool pe_find_section(mem_view_t v, uint64_t module_base, const char* name,
                     uint64_t* rva_out, uint32_t* vsize_out) {
    size_t sec; uint16_t nsec;
    if (!name || !pe_section_table(v, module_base, &sec, &nsec)) return false;

    size_t want = strlen(name);
    if (want > 8) want = 8;

    for (uint16_t i = 0; i < nsec; i++) {
        const size_t sh = sec + (size_t)i * SH_SIZE;
        if (sh + SH_SIZE > v.size) break;
        char nm[9] = {0};
        memcpy(nm, v.data + sh, 8);
        if (strncmp(nm, name, want) == 0 && (want == 8 || nm[want] == '\0')) {
            uint32_t vsize, vaddr;
            memcpy(&vsize, v.data + sh + SH_VSIZE_OFF, 4);  /* Misc.VirtualSize */
            memcpy(&vaddr, v.data + sh + SH_VADDR_OFF, 4);  /* VirtualAddress   */
            if (rva_out)   *rva_out   = vaddr;
            if (vsize_out) *vsize_out = vsize;
            return true;
        }
    }
    return false;
}

int pe_sections(mem_view_t v, uint64_t module_base, pe_secrec* out, int max) {
    size_t sec; uint16_t nsec;
    if (!pe_section_table(v, module_base, &sec, &nsec)) return -1;

    int total = 0;
    for (uint16_t i = 0; i < nsec; i++) {
        const size_t sh = sec + (size_t)i * SH_SIZE;
        if (sh + SH_SIZE > v.size) break;           /* headers truncated in view */
        if (out && total < max) {
            pe_secrec* r = &out[total];
            memset(r->name, 0, sizeof r->name);
            memcpy(r->name, v.data + sh, 8);        /* name[8] stays the NUL slot */
            uint32_t ch;
            memcpy(&r->vsize, v.data + sh + SH_VSIZE_OFF, 4);
            memcpy(&r->rva,   v.data + sh + SH_VADDR_OFF, 4);
            memcpy(&ch,       v.data + sh + SH_CHAR_OFF,  4);
            r->prot = pe_prot(ch);
        }
        total++;
    }
    return total;
}

bool pe_section(mem_view_t v, uint64_t module_base, const char* name, mem_view_t* out) {
    uint64_t rva; uint32_t vsize;
    if (!out || !pe_find_section(v, module_base, name, &rva, &vsize)) return false;
    *out = mem_sub(v, module_base + rva, vsize);
    return out->data != NULL;
}

int vmie_pe_section(vmie_mem* m, uintptr_t cr3, uint64_t module_base,
                    const char* name, uint8_t* buf, size_t bufcap, mem_view_t* out) {
    uint8_t hdr[0x1000];
    if (!out || !buf || gva_read(m, cr3, module_base, hdr, sizeof hdr)) return -1;
    const mem_view_t hv = { hdr, sizeof hdr, module_base };
    uint64_t rva; uint32_t vsize;
    if (!pe_find_section(hv, module_base, name, &rva, &vsize)) return -1;
    const size_t n = vsize < bufcap ? vsize : bufcap;
    if (gva_read(m, cr3, module_base + rva, buf, n)) return -1;
    out->data = buf; out->size = n; out->base_va = module_base + rva;
    return 0;
}

/* ---- public win32 surface: section enumeration + section views ----------- *
 * Cold paths (one-shot header parse / section gather, not a hot loop). They
 * reuse pe_sections / the shared section-table walk above - no second parser -
 * and never allocate: the section_view buffer is caller-owned. */

int vmie_win32_sections(vmie_win32* v, uint64_t cr3, uint64_t module_base,
                        section_desc* out, int max) __attribute__((cold));
int vmie_win32_sections(vmie_win32* v, uint64_t cr3, uint64_t module_base,
                        section_desc* out, int max) {
    vmie_mem* m = vmie_win32_mem(v);
    if (!m) { return -1; }

    uint8_t hdr[0x1000];
    if (gva_read(m, cr3, module_base, hdr, sizeof hdr)) { return -1; }
    const mem_view_t hv = { hdr, sizeof hdr, module_base };

    /* count first via pe_sections(out=NULL); if the caller only wants the count
     * (out==NULL) we are done. */
    const int total = pe_sections(hv, module_base, NULL, 0);
    if (total < 0 || !out) { return total; }

    pe_secrec recs[96];
    int cap = max;
    if (cap < 0) { cap = 0; }
    if (cap > (int)(sizeof recs / sizeof recs[0])) {
        cap = (int)(sizeof recs / sizeof recs[0]);
    }
    const int got = pe_sections(hv, module_base, recs, cap);
    if (got < 0) { return got; }

    int n = got < cap ? got : cap;
    for (int i = 0; i < n; i++) {
        memcpy(out[i].name, recs[i].name, sizeof out[i].name);
        out[i].rva   = recs[i].rva;
        out[i].vsize = recs[i].vsize;
        out[i].prot  = recs[i].prot;
    }
    return total;        /* total count, even if it exceeded `max` */
}

int vmie_win32_section_view(vmie_win32* v, uint64_t cr3, uint64_t module_base,
                            const section_desc* sec, view_base mode,
                            uint8_t* buf, size_t bufcap, mem_view_t* out)
    __attribute__((cold));
int vmie_win32_section_view(vmie_win32* v, uint64_t cr3, uint64_t module_base,
                            const section_desc* sec, view_base mode,
                            uint8_t* buf, size_t bufcap, mem_view_t* out) {
    vmie_mem* m = vmie_win32_mem(v);
    if (!m || !sec || !buf || !out) { return -1; }

    const size_t n = sec->vsize < bufcap ? sec->vsize : bufcap;
    if (gva_read(m, cr3, module_base + sec->rva, buf, n)) { return -1; }

    uint64_t base_va = 0;
    switch (mode) {
        case MODULE_RVA:   base_va = sec->rva;               break;
        case ABSOLUTE_VA:  base_va = module_base + sec->rva; break;
        case SECTION_LOCAL:
        default:           base_va = 0;                      break;
    }
    out->data = buf; out->size = n; out->base_va = base_va;
    return 0;
}