/*
heap_lfh -- a deliberately tiny note manager, vulnerable to a classic
use-after-free on Windows' real Low Fragmentation Heap (not a simulation):
a private heap is created and explicitly switched into LFH mode via
HeapSetInformation(HeapCompatibilityInformation, 2), the same documented
mechanism winpwn's own reference notes (heap/pocs/02_lfh_probe.c) use to
make LFH active deterministically instead of waiting on the ~17-allocation
auto-activation heuristic.

The bug: Free (F) does not clear the dangling pointer in the notes table.
Print (P) calls through Note.onPrint without checking whether the note was
freed. There's also no bounds check tying the declared length of a B
(buffer) command to the fixed 32-byte allocation it writes into.

Protocol (one command per line, stdout is unbuffered):
  A <text>        allocate a Note{char title[24]; void(*onPrint)(const
                  char*);}, fills title (truncated to 23 chars + NUL),
                  sets onPrint to the real print function.
                  -> "OK id=<n> addr=0x<hex>"
  F <id>          HeapFree the note at that id (pointer stays in the table).
                  -> "OK"
  B <hex32bytes>  allocate a raw 32-byte buffer from the SAME heap and
                  write exactly 32 attacker-supplied bytes into it
                  (hex-encoded, 64 hex chars).
                  -> "OK addr=0x<hex>"
  P <id>          call notes[id]->onPrint(notes[id]->title).
                  -> whatever onPrint prints
  Q               quit.

__declspec(dllexport) on win() so it's found via the PE export table the
same way examples/task2_rop's win() is -- no symbols needed.
*/
#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 4096
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);
}

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

    printf("you just got code execution via a freed onPrint pointer\n");

    hFile = CreateFileA("flag.txt", GENERIC_READ, FILE_SHARE_READ, NULL,
                         OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (hFile == INVALID_HANDLE_VALUE) {
        printf("Cannot open file.\n");
        ExitProcess(0);
    }
    if (ReadFile(hFile, buffer, sizeof(buffer) - 1, &bytesRead, NULL) && bytesRead > 0) {
        buffer[bytesRead] = '\0';
        printf("%s", buffer);
    }
    printf("\n");
    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 unhex32(const char *hex, unsigned char *out) {
    if (strlen(hex) != 64) return 0;
    for (int i = 0; i < 32; 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 1;
}

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;
    }
    ULONG mode = 2; /* HeapCompatibilityInformation: 2 == LFH, see heap/pocs/02_lfh_probe.c */
    HeapSetInformation(g_heap, HeapCompatibilityInformation, &mode, sizeof(mode));

    printf("heap_lfh ready\n");

    char line[256];
    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 notes\n");
                continue;
            }
            Note *n = (Note *)HeapAlloc(g_heap, 0, sizeof(Note));
            if (!n) { printf("ERR alloc failed\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] == 'F' && 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]);
            printf("OK\n");
        } else if (line[0] == 'B' && line[1] == ' ') {
            unsigned char buf[32];
            if (!unhex32(line + 2, buf)) {
                printf("ERR need exactly 64 hex chars (32 bytes)\n");
                continue;
            }
            void *p = HeapAlloc(g_heap, 0, 32);
            if (!p) { printf("ERR alloc failed\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 || !g_notes[id]) {
                printf("ERR bad id\n");
                continue;
            }
            g_notes[id]->onPrint(g_notes[id]->title);
        } else if (line[0] == 'Q') {
            break;
        } else {
            printf("ERR unknown command\n");
        }
    }
    return 0;
}
