/*
heap_overflow -- adjacent-chunk heap overflow on NT Heap backend.

Two Note objects allocated consecutively on a clean NT Heap segment (no LFH
activation for this allocation count -- only 2-3 objects total per exploit
run). Each Note is:

  Note { char buf[24]; void(*action)(const char*); }  = 32 bytes

The overflow site: command W <id> <hex> writes hex-decoded bytes starting at
note[id]->buf with NO bounds check on the hex length. A write longer than 24
bytes spills into note[id]->action (the function pointer), and if longer still,
spills across the 16-byte NT _HEAP_ENTRY header into the next note's buf and
eventually its action pointer too.

Layout on the heap (addresses relative to note[0]'s HEADER start):
  +0x00   _HEAP_ENTRY header (16 bytes, XOR-encoded -- we don't need to care,
           we're not freeing note[1] after the overflow)
  +0x10   note[0]->buf    [0..23] (24 bytes)
  +0x28   note[0]->action [24..31] (8 bytes, function pointer)
  +0x30   _HEAP_ENTRY header for note[1] (16 bytes, gets corrupted -- OK)
  +0x40   note[1]->buf    [0..23]
  +0x58   note[1]->action [24..31]  <-- target: 0x48 bytes from note[0]->buf

So writing 0x48 = 72 bytes from note[0]->buf, with win()'s address in bytes
64-71 (0-indexed from note[0]->buf), overwrites note[1]->action.

Then command C 1 calls note[1]->action -- attacker lands at win().

Protocol (one command per line, stdout unbuffered):
  A <text>          allocate Note{buf=text (truncated to 23 chars), action=real_action}.
                    -> "OK id=<n> addr=0x<hex>"
  W <id> <hex>      write hex bytes starting at note[id]->buf -- NO BOUNDS CHECK.
                    -> "OK"
  C <id>            call note[id]->action(note[id]->buf).
  Q                 quit.

NOTE FOR TASK AUTHORS: the critical constraint is that notes[0] and notes[1]
must end up adjacent in the heap (no gap). With HeapCreate(0,0,0) and exactly
two 32-byte allocations with no interleaving frees, they are reliably adjacent
on the NT Heap backend on this build. If you increase the note struct size or
add other allocations between them, verify adjacency with the winpwn heap CLI:
`winpwn heap <pid> -walk` and look for "adjacent busy pairs". Also: the
_HEAP_ENTRY header between them uses XOR encoding and gets corrupted by the
overflow, but since we call action directly (no subsequent HeapFree on note[1]),
the corrupted header is never read back by the allocator -- so the exploit
works even though the header is garbage after the write.
*/
#define _CRT_SECURE_NO_WARNINGS
#include <windows.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

typedef struct {
    char buf[24];
    void (*action)(const char *);
} Note;

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

static void real_action(const char *buf) {
    printf("note: %s\n", buf);
    fflush(stdout);
}

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

    printf("heap overflow worked -- adjacent chunk action pointer corrupted\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]);
        int 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);

    /* Private heap, no LFH -- guaranteed backend allocation ordering */
    g_heap = HeapCreate(0, 0, 0);
    if (!g_heap) {
        printf("HeapCreate failed\n");
        return 1;
    }

    printf("heap_overflow ready\n");

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

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

        } else if (line[0] == 'W' && line[1] == ' ') {
            int id;
            char hexbuf[513];
            if (sscanf(line + 2, "%d %512s", &id, hexbuf) != 2) {
                printf("ERR usage: W <id> <hex>\n"); continue;
            }
            if (id < 0 || id >= g_note_count || !g_notes[id]) {
                printf("ERR bad id\n"); continue;
            }
            unsigned char raw[256];
            int n = unhex(hexbuf, raw, sizeof(raw));
            if (n == 0) { printf("ERR bad hex\n"); continue; }
            /* THE BUG: no check that n <= sizeof(Note) */
            memcpy(g_notes[id]->buf, raw, n);
            printf("OK\n");

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

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