/*
heap_typemix -- UAF type-confusion on Windows NT Heap (no LFH).

Two struct types happen to be exactly the same size (32 bytes):

  Note:  { char title[24]; void(*onPrint)(const char*); }
  Token: { char data[24]; void(*validate)(const char*); }

The bug: Delete (D command) frees Note.title memory but leaves the note
pointer in the table. Token (T command) allocates raw 32-byte objects from
the SAME heap. Because the freed Note slot goes straight back to the NT Heap
backend freelist (no LFH: HeapSetInformation is NOT called here), the next
HeapAlloc(32) returns the same address. The caller then "prints" the stale
Note, which actually dispatches through Token.validate -- attacker controlled.

This is a pure type-confusion UAF: no grooming needed, freed chunk comes back
on the very next 32-byte allocation, deterministically (pre-LFH, backend only,
1-2 allocations total per run).

Protocol (one command per line, stdout unbuffered):
  N <text>          allocate Note{title[24], onPrint=real_print}.
                    -> "OK id=<n> addr=0x<hex>"
  D <id>            HeapFree the Note (pointer stays in table -- the bug).
                    -> "OK"
  T <hex32>         HeapAlloc 32 bytes, write attacker bytes verbatim.
                    -> "OK addr=0x<hex>"
  P <id>            call notes[id]->onPrint(notes[id]->title).
                    -> output from whatever onPrint actually points at now.
  Q                 quit.

NOTE FOR TASK AUTHORS: no LFH activation (≥18 same-size allocations) is
needed here because the freed chunk is on the NT Heap backend freelist. One
allocation reliably reclaims it on this build (10.0.26100). If you need to
adapt this for a different build, verify the reuse behavior with PoC 02
(heap/pocs/02_lfh_probe.c) before trusting this claim.
*/
#define _CRT_SECURE_NO_WARNINGS
#include <windows.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

typedef struct {
    char title[24];
    void (*onPrint)(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_print(const char *title) {
    printf("note: %s\n", title);
    fflush(stdout);
}

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

    printf("type confusion worked -- code execution via freed onPrint pointer\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 forcing -- freed chunks return to backend freelist */
    g_heap = HeapCreate(0, 0, 0);
    if (!g_heap) {
        printf("HeapCreate failed\n");
        return 1;
    }

    printf("heap_typemix ready\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 too many notes\n"); continue; }
            Note *n = (Note *)HeapAlloc(g_heap, 0, sizeof(Note));
            if (!n) { printf("ERR alloc\n"); continue; }
            strncpy(n->title, line + 2, sizeof(n->title) - 1);
            n->title[sizeof(n->title) - 1] = 0;
            n->onPrint = real_print;
            int id = g_note_count++;
            g_notes[id] = n;
            printf("OK id=%d addr=0x%p\n", id, (void *)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 in table -- classic dangling UAF */
            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 exactly 64 hex chars (32 bytes)\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;
            }
            /* note[id] may be freed -- the UAF lives here */
            g_notes[id]->onPrint(g_notes[id]->title);

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