/*
heap_info_leak -- two-stage Windows heap exploit: OOB read to defeat ASLR,
then UAF function-pointer overwrite to get code execution.

STAGE 1 (info leak): The S command prints up to `len` bytes from a Note's
data field -- but puts no upper bound on `len`. A Note is:

  struct Note { char data[24]; void(*onShow)(const char*); }  = 32 bytes

Requesting S <id> 32 prints all 32 bytes: the 24 bytes of actual data AND the
8-byte function pointer onShow (= real_show, a known-offset function inside
this binary). Since the binary has ASLR enabled (DynamicBase set -- try
`winpwn checksec`), the runtime address is random per run. But it's a fixed
offset from the binary's base. Reading those 8 bytes at runtime reveals the
actual load address, and from there the player computes win()'s address as:
    win_va = leaked_real_show_va + (win_rva - real_show_rva)
where the RVA difference is a static property of the binary, readable via
`winpwn exports heap_info_leak.exe`.

STAGE 2 (UAF): The D command frees note[id] (pointer stays in table). The T
command allocates a raw 32-byte token from the same private heap. Since LFH is
NOT active (≤1 prior allocation of this size), the freed Note slot comes back
immediately on the next 32-byte allocation. The player places win()'s computed
address at offset 24 in the token -- exactly where onShow lives. Then P <id>
dispatches through the dangling note pointer, calling win().

Protocol (one command per line, stdout unbuffered):
  N <text>          allocate Note{data=text (truncated 23 chars), onShow=real_show}.
                    -> "OK id=<n> addr=0x<hex>"
  S <id> <len>      print exactly <len> bytes of note[id]->data as hex.
                    Bug: no bounds check on len (valid: 1..255).
                    -> "HEX <hexbytes>"
  D <id>            HeapFree note (pointer stays -- the UAF bug).
                    -> "OK"
  T <hex32>         allocate raw 32-byte token with attacker content.
                    -> "OK addr=0x<hex>"
  P <id>            call note[id]->onShow(note[id]->data).
                    -> output (or crash if onShow is garbage)
  Q                 quit.

NOTE FOR TASK AUTHORS: This is the standard two-step pattern used in:
- justCTF 2024 "Baby Heap but Windows": heap UAF + ntdll leak via heap+0x2c0
- ECW CTF 2024 "Address Book": type confusion OOB read leaks binary pointer
The binary has ASLR enabled (unlike the other heap_* examples in this repo
which all have ASLR: No) because learning to handle info-leak → address-compute
→ exploit is the explicit teaching goal. The player MUST read the leaked function
pointer bytes and calculate win()' runtime address. The solve script in main.go
demonstrates the full chain in Go using winpwn's OpenPE to get the RVA offset.
*/
#define _CRT_SECURE_NO_WARNINGS
#include <windows.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

typedef struct Note {
    char data[24];
    void (*onShow)(const char *);
} Note;

#define MAX_NOTES 256
static Note *g_notes[MAX_NOTES];
static int g_note_count = 0;
static HANDLE g_heap;

__declspec(dllexport) void real_show(const char *data) {
    printf("note: %.*s\n", 23, data);
    fflush(stdout);
}

__declspec(dllexport) void win(const char *ignored) {
    HANDLE hFile;
    char buffer[256];
    DWORD bytesRead;

    printf("two-stage exploit worked: OOB read defeated ASLR, UAF gave code exec\n");
    fflush(stdout);

    hFile = CreateFileA("flag.txt", GENERIC_READ, FILE_SHARE_READ, NULL,
                        OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (hFile == INVALID_HANDLE_VALUE) {
        printf("Cannot open flag.txt\n"); fflush(stdout); ExitProcess(0);
    }
    if (ReadFile(hFile, buffer, sizeof(buffer) - 1, &bytesRead, NULL) && bytesRead > 0) {
        buffer[bytesRead] = '\0';
        printf("%s", buffer); fflush(stdout);
    }
    CloseHandle(hFile);
    ExitProcess(0);
}

static int hexval(char c) {
    if (c >= '0' && c <= '9') return c - '0';
    if (c >= 'a' && c <= 'f') return c - 'a' + 10;
    if (c >= 'A' && c <= 'F') return c - 'A' + 10;
    return -1;
}
static int unhex(const char *hex, unsigned char *out, int maxout) {
    int n = (int)strlen(hex);
    if (n % 2 != 0 || n / 2 > maxout) return 0;
    for (int i = 0; i < n / 2; i++) {
        int hi = hexval(hex[i * 2]), lo = hexval(hex[i * 2 + 1]);
        if (hi < 0 || lo < 0) return 0;
        out[i] = (unsigned char)((hi << 4) | lo);
    }
    return n / 2;
}

int main(void) {
    setvbuf(stdout, NULL, _IONBF, 0);
    setvbuf(stdin,  NULL, _IONBF, 0);

    g_heap = HeapCreate(0, 0, 0);
    if (!g_heap) { printf("HeapCreate failed\n"); return 1; }
    printf("heap_info_leak ready (ASLR: enabled, DynamicBase set)\n");

    char line[512];
    while (fgets(line, sizeof(line), stdin)) {
        line[strcspn(line, "\r\n")] = 0;

        if (line[0] == 'N' && line[1] == ' ') {
            if (g_note_count >= MAX_NOTES) { printf("ERR\n"); continue; }
            Note *n = (Note *)HeapAlloc(g_heap, 0, sizeof(Note));
            if (!n) { printf("ERR alloc\n"); continue; }
            memset(n->data, 0, sizeof(n->data));
            strncpy(n->data, line + 2, sizeof(n->data) - 1);
            n->onShow = real_show;
            int id = g_note_count++;
            g_notes[id] = n;
            printf("OK id=%d addr=0x%p\n", id, (void *)n);

        } else if (line[0] == 'S' && line[1] == ' ') {
            int id, len;
            if (sscanf(line + 2, "%d %d", &id, &len) != 2 || len < 1 || len > 255) {
                printf("ERR usage: S <id> <len 1-255>\n"); continue;
            }
            if (id < 0 || id >= g_note_count || !g_notes[id]) {
                printf("ERR bad id\n"); continue;
            }
            /* BUG: prints len bytes from data but no check that len <= 24 */
            printf("HEX ");
            const unsigned char *p = (const unsigned char *)g_notes[id]->data;
            for (int i = 0; i < len; i++) printf("%02x", p[i]);
            printf("\n");

        } else if (line[0] == 'D' && line[1] == ' ') {
            int id = atoi(line + 2);
            if (id < 0 || id >= g_note_count || !g_notes[id]) {
                printf("ERR bad id\n"); continue;
            }
            HeapFree(g_heap, 0, g_notes[id]);
            /* BUG: pointer stays */
            printf("OK\n");

        } else if (line[0] == 'T' && line[1] == ' ') {
            unsigned char buf[32];
            int n = unhex(line + 2, buf, 32);
            if (n != 32) { printf("ERR need 64 hex chars\n"); continue; }
            void *p = HeapAlloc(g_heap, 0, 32);
            if (!p) { printf("ERR alloc\n"); continue; }
            memcpy(p, buf, 32);
            printf("OK addr=0x%p\n", p);

        } else if (line[0] == 'P' && line[1] == ' ') {
            int id = atoi(line + 2);
            if (id < 0 || id >= g_note_count) { printf("ERR bad id\n"); continue; }
            g_notes[id]->onShow(g_notes[id]->data);

        } else if (line[0] == 'Q') {
            break;
        } else {
            printf("ERR unknown\n");
        }
        fflush(stdout);
    }
    return 0;
}
