/*
heap_segment -- adjacent-chunk heap overflow on Windows' real Segment
Heap (not NT Heap/LFH). The process opts into Segment Heap via an
embedded application manifest (<heapType>SegmentHeap> -- see
heap_segment.manifest/heap_segment.rc, the only Microsoft-documented way
to force it for a specific image without touching machine-wide settings),
so GetProcessHeap() itself is Segment-Heap-backed: confirmed empirically
while building this example by reading the heap handle's own Signature
field (*(DWORD*)(GetProcessHeap()+0x10) == 0xddeeddee for Segment Heap,
0xffeeffee for classic NT Heap).

The bug: O (overflow) writes attacker-controlled, attacker-LENGTH bytes
starting at a Profile's address with no check that the length fits the
32-byte allocation -- a plain unchecked memcpy. Segment Heap's famous
mitigation (full physical isolation of heap *metadata* from user *data*,
see USAGE.md's walkthrough) means this overflow can never reach allocator
control structures, but it can still walk straight into whatever user data
happens to be allocated right after it in the same page -- and Segment
Heap's "Small" allocator packs same-size allocations densely into 4KB
pages, just at a randomized offset within the page rather than in
allocation order. Leak enough addresses (the A command leaks each one) and
some pair will be exactly 32 bytes apart (empirically: spray>=10 finds
one in every trial run while building this).

Protocol (one command per line, stdout unbuffered):
  A <text>        allocate a Profile{char name[24]; void(*describe)(const
                  char*);}, fills name (truncated to 23 chars + NUL), sets
                  describe to the real print function.
                  -> "OK id=<n> addr=0x<hex>"
  O <id> <hex>    write decode(hex) raw bytes starting at profiles[id]
                  (i.e. at name[0]) -- NOT bounds-checked against the
                  32-byte allocation.
                  -> "OK"
  D <id>          call profiles[id]->describe(profiles[id]->name).
  Q               quit.
*/
#define _CRT_SECURE_NO_WARNINGS
#include <windows.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

typedef struct {
    char name[24];
    void (*describe)(const char *);
} Profile;

#define MAX_PROFILES 4096
static Profile *g_profiles[MAX_PROFILES];
static int g_profile_count = 0;

static void real_describe(const char *name) {
    printf("profile: %s\n", name);
}

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

    printf("you just got code execution via an adjacent-chunk overflow\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;
}

/* Decodes hex into out, returns the number of bytes decoded (0 on bad input).
   No length cap here -- the caller (the O command) is the vulnerable site. */
static int unhex(const char *hex, unsigned char *out, int max_out) {
    int n = (int)strlen(hex);
    if (n % 2 != 0) return 0;
    int len = n / 2;
    if (len > max_out) return 0; /* still capped by our own receive buffer, not by the target's allocation */
    for (int i = 0; i < len; 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 len;
}

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

    unsigned int sig = *(unsigned int *)((char *)GetProcessHeap() + 0x10);
    printf("heap_segment ready (heap signature 0x%08x)\n", sig);

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

        if (line[0] == 'A' && line[1] == ' ') {
            if (g_profile_count >= MAX_PROFILES) {
                printf("ERR too many profiles\n");
                continue;
            }
            Profile *p = (Profile *)HeapAlloc(GetProcessHeap(), 0, sizeof(Profile));
            if (!p) { printf("ERR alloc failed\n"); continue; }
            strncpy(p->name, line + 2, sizeof(p->name) - 1);
            p->name[sizeof(p->name) - 1] = 0;
            p->describe = real_describe;
            int id = g_profile_count++;
            g_profiles[id] = p;
            printf("OK id=%d addr=0x%p\n", id, (void *)p);
        } else if (line[0] == 'O' && line[1] == ' ') {
            int id;
            char hexbuf[513];
            if (sscanf(line + 2, "%d %512s", &id, hexbuf) != 2) {
                printf("ERR usage: O <id> <hex>\n");
                continue;
            }
            if (id < 0 || id >= g_profile_count || !g_profiles[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(Profile). */
            memcpy(g_profiles[id], raw, n);
            printf("OK\n");
        } else if (line[0] == 'D' && line[1] == ' ') {
            int id = atoi(line + 2);
            if (id < 0 || id >= g_profile_count || !g_profiles[id]) {
                printf("ERR bad id\n");
                continue;
            }
            g_profiles[id]->describe(g_profiles[id]->name);
        } else if (line[0] == 'Q') {
            break;
        } else {
            printf("ERR unknown command\n");
        }
    }
    return 0;
}
