package winpwn

import (
	"encoding/binary"
	"fmt"
	"io"
)

// This file is the foundation layer of winpwn's heap-structure parsing
// (ROADMAP.md's Phase 9): decoding the NT Heap's on-disk^Won-memory layout
// directly off a ReaderAt, the same "works identically against a file or a
// live process" design PEFile already uses for PE images. Everything here
// works against any io.ReaderAt -- a *ProcessMemory, a *Debugger (via a
// thin adapter), or a synthetic in-memory buffer in a test -- so it has no
// build tag despite being Windows-structure-specific; only *finding* a
// heap address in a live process (heap_windows.go's ListProcessHeaps) needs
// actual Windows syscalls.
//
// Every offset/field below was cross-checked against `dt ntdll!_HEAP` and
// friends via cdb (public ntdll symbols carry full type info even without
// source) on this machine's build (10.0.26100, Windows 11 24H2) -- not
// copied from a blog post and not guessed from memory. The decode logic
// specifically (DecodeHeapEntry's XOR-encoding handling and the
// Size/UnusedBytes arithmetic) was verified against `!heap -a`'s own
// ground-truth entry listing for three real allocations of different
// sizes before being trusted -- see heap_test.go's
// TestDecodeHeapEntryMatchesLiveGroundTruth for the exact captured bytes
// and the discovery that cost the most back-and-forth: UserSize is
// `Size*HeapEntrySize - UnusedBytes` with NO separate subtraction for the
// header, not the more "obvious" `Size*HeapEntrySize - HeapEntrySize -
// UnusedBytes` -- UnusedBytes already accounts for the header itself, a
// detail no amount of reading the struct definition alone would have
// caught without comparing against real numbers.
//
// Struct layouts are a moving target across Windows builds -- if you're
// reading this on a different build and something doesn't line up, redo
// the `dt ntdll!_HEAP` capture in USAGE.md's heap walkthrough rather than
// assuming these offsets still hold.

// HeapKind identifies which allocator backend a heap handle is using,
// determined the same way HeapAlloc itself effectively does: by the magic
// signature 0x10 bytes into the handle (verified earlier in this project
// against examples/heap_lfh and examples/heap_segment: 0xffeeffee for the
// classic NT Heap, 0xddeeddee for Segment Heap).
type HeapKind int

const (
	HeapKindUnknown HeapKind = iota
	HeapKindNT
	HeapKindSegment
)

func (k HeapKind) String() string {
	switch k {
	case HeapKindNT:
		return "NT Heap"
	case HeapKindSegment:
		return "Segment Heap"
	default:
		return "unknown"
	}
}

const (
	heapSignatureNT      = 0xffeeffee
	heapSignatureSegment = 0xddeeddee
)

// DetectHeapKind reads the 4-byte signature at heapAddr+0x10 -- the very
// first thing to do with any heap handle/address before parsing it any
// further, since _HEAP and _SEGMENT_HEAP are structurally unrelated past
// this point.
func DetectHeapKind(r io.ReaderAt, heapAddr uint64) (HeapKind, error) {
	sig, err := readUint32AtValue(r, int64(heapAddr)+0x10)
	if err != nil {
		return HeapKindUnknown, fmt.Errorf("reading signature at 0x%x+0x10: %w", heapAddr, err)
	}
	switch sig {
	case heapSignatureNT:
		return HeapKindNT, nil
	case heapSignatureSegment:
		return HeapKindSegment, nil
	default:
		return HeapKindUnknown, fmt.Errorf("unrecognized heap signature 0x%08x at 0x%x+0x10 (expected 0x%08x NT Heap or 0x%08x Segment Heap)",
			sig, heapAddr, heapSignatureNT, heapSignatureSegment)
	}
}

// HeapEntrySize is HEAP_GRANULARITY on x64: every _HEAP_ENTRY header is
// exactly this many bytes, and Size/PreviousSize are both counted in units
// of it, not in plain bytes. Confirmed empirically (heap_test.go), not
// just asserted from the struct definition -- see this file's top comment.
const HeapEntrySize = 0x10

// _HEAP_ENTRY.Flags bits. Long-standing, widely published constants (every
// heap-exploitation writeup and WinDbg's own !heap extension use these same
// values), unlike the struct offsets above which are this-build-specific --
// these have been stable since before Windows 8's header encoding existed.
const (
	HeapEntryBusy         = 0x01
	HeapEntryExtraPresent = 0x02
	HeapEntryFillPattern  = 0x04
	HeapEntryVirtualAlloc = 0x08
	HeapEntryLastEntry    = 0x10
)

// HeapEntry is the decoded form of a 16-byte _HEAP_ENTRY header -- the
// thing immediately preceding every NT Heap allocation's user data,
// equally present whether or not LFH/the front-end allocator owns the
// block (see this file's WalkSegment vs heap_lfh.go's subsegment-aware
// walk for why both views matter).
type HeapEntry struct {
	Addr          uint64 // address of the header itself, i.e. UserData()-HeapEntrySize
	Size          uint16 // total block size (header+user data+padding), in HeapEntrySize units
	Flags         uint8
	SmallTagIndex uint8
	PreviousSize  uint16 // previous block's total size, same units -- lets you walk backward
	SegmentOffset uint8  // doubles as LFHFlags when this entry belongs to an LFH subsegment
	UnusedBytes   uint8
}

func (e HeapEntry) Busy() bool         { return e.Flags&HeapEntryBusy != 0 }
func (e HeapEntry) LastEntry() bool    { return e.Flags&HeapEntryLastEntry != 0 }
func (e HeapEntry) VirtualAlloc() bool { return e.Flags&HeapEntryVirtualAlloc != 0 }

// BlockSize is the entry's total physical footprint (header + user data +
// any padding), in bytes.
func (e HeapEntry) BlockSize() uint64 { return uint64(e.Size) * HeapEntrySize }

// PreviousBlockSize is the immediately preceding entry's BlockSize, in
// bytes -- lets you find where the previous entry starts without having
// walked there directly (e.Addr - e.PreviousBlockSize()).
func (e HeapEntry) PreviousBlockSize() uint64 { return uint64(e.PreviousSize) * HeapEntrySize }

// UserData is the address HeapAlloc actually returned to the caller.
func (e HeapEntry) UserData() uint64 { return e.Addr + HeapEntrySize }

// UserSize is what HeapSize() would report for this block -- the original
// requested size (rounded up to grain, header already accounted for).
// Empirically, this is BlockSize()-UnusedBytes with no separate header
// subtraction; see this file's top comment for how that was confirmed.
func (e HeapEntry) UserSize() uint64 {
	bs := e.BlockSize()
	if uint64(e.UnusedBytes) > bs {
		return 0
	}
	return bs - uint64(e.UnusedBytes)
}

// NextEntry is the address of the entry immediately following this one.
func (e HeapEntry) NextEntry() uint64 { return e.Addr + e.BlockSize() }

// DecodeHeapEntry decodes a raw 16-byte _HEAP_ENTRY read from addr,
// reversing the Windows 8+ header-encoding mitigation if encoding is
// non-nil (pass Heap.Encoding; nil only if you've separately confirmed
// EncodeFlagMask is 0 for this heap, which is rare in practice). Only
// bytes 8-15 of the entry are ever encoded -- bytes 0-7
// (PreviousBlockPrivateData) are plain, usually-stale data, not part of
// the XOR scheme at all.
func DecodeHeapEntry(addr uint64, raw [16]byte, encoding *[16]byte) HeapEntry {
	var compact [8]byte
	copy(compact[:], raw[8:16])
	if encoding != nil {
		for i := range compact {
			compact[i] ^= encoding[8+i]
		}
	}
	return HeapEntry{
		Addr:          addr,
		Size:          binary.LittleEndian.Uint16(compact[0:2]),
		Flags:         compact[2],
		SmallTagIndex: compact[3],
		PreviousSize:  binary.LittleEndian.Uint16(compact[4:6]),
		SegmentOffset: compact[6],
		UnusedBytes:   compact[7],
	}
}

// ReadHeapEntry reads and decodes the entry header at addr.
func ReadHeapEntry(r io.ReaderAt, addr uint64, encoding *[16]byte) (HeapEntry, error) {
	var raw [16]byte
	if _, err := r.ReadAt(raw[:], int64(addr)); err != nil {
		return HeapEntry{}, fmt.Errorf("reading entry at 0x%x: %w", addr, err)
	}
	return DecodeHeapEntry(addr, raw, encoding), nil
}

// Offsets within ntdll's x64 _HEAP, confirmed via `dt ntdll!_HEAP` against
// this machine's build (10.0.26100) -- see this file's top comment.
const (
	heapOffBaseAddress      = 0x030 // _HEAP_SEGMENT.BaseAddress (the embedded Segment0)
	heapOffFirstEntry       = 0x040 // _HEAP_SEGMENT.FirstEntry
	heapOffLastValidEntry   = 0x048 // _HEAP_SEGMENT.LastValidEntry
	heapOffSegmentListEntry = 0x018 // _HEAP_SEGMENT.SegmentListEntry (LIST_ENTRY)
	heapOffFlags            = 0x070
	heapOffForceFlags       = 0x074
	heapOffEncodeFlagMask   = 0x07c
	heapOffEncoding         = 0x080 // _HEAP_ENTRY-shaped, 16 bytes
	heapOffSignature        = 0x098
	heapOffSegmentList      = 0x120 // LIST_ENTRY, head of all _HEAP_SEGMENTs (including Segment0 itself)
	heapOffFrontEndHeap     = 0x198
	heapOffFrontEndHeapType = 0x1a2
)

// FrontEndHeapType values (_HEAP.FrontEndHeapType).
const (
	FrontEndHeapNone      = 0
	FrontEndHeapLookaside = 1 // legacy, rarely seen on modern Windows
	FrontEndHeapLFH       = 2
)

// Heap is the decoded subset of ntdll's _HEAP that matters for
// exploitation -- not a byte-for-byte mirror of the ~700-byte real struct
// (most of it is debug/tuning bookkeeping no script needs), the same
// "decode what's useful, expose RawStream for the rest" philosophy
// minidump.go uses.
type Heap struct {
	Addr             uint64
	Signature        uint32 // 0xeeffeeff on a real heap -- confirmed live; deliberately byte-rotated from _HEAP_SEGMENT's 0xffeeffee, a different field, not a typo if you see both
	Flags            uint32
	ForceFlags       uint32
	EncodeFlagMask   uint32
	Encoding         [16]byte // pass to DecodeHeapEntry/ReadHeapEntry
	FrontEndHeapType uint8
	FrontEndHeap     uint64 // *_LFH_HEAP when FrontEndHeapType == FrontEndHeapLFH
	BaseAddress      uint64 // Segment0's base address
	FirstEntry       uint64 // Segment0's first entry
	LastValidEntry   uint64 // Segment0's end-of-committed-range marker
}

// readUint32AtValue mirrors minidump.go's readUint32At but returns the
// value directly instead of writing through an out-pointer, matching this
// file's other readXAt helpers below.
func readUint32AtValue(r io.ReaderAt, offset int64) (uint32, error) {
	var buf [4]byte
	if _, err := r.ReadAt(buf[:], offset); err != nil {
		return 0, err
	}
	return binary.LittleEndian.Uint32(buf[:]), nil
}

func readUint64At(r io.ReaderAt, offset int64) (uint64, error) {
	var buf [8]byte
	if _, err := r.ReadAt(buf[:], offset); err != nil {
		return 0, err
	}
	return binary.LittleEndian.Uint64(buf[:]), nil
}

func readUint8At(r io.ReaderAt, offset int64) (uint8, error) {
	var buf [1]byte
	if _, err := r.ReadAt(buf[:], offset); err != nil {
		return 0, err
	}
	return buf[0], nil
}

// ReadHeap decodes addr's _HEAP header. Returns an error if addr's
// signature isn't the NT Heap one (use DetectHeapKind first if you don't
// already know, or ReadSegmentHeap for a 0xddeeddee handle).
func ReadHeap(r io.ReaderAt, addr uint64) (*Heap, error) {
	kind, err := DetectHeapKind(r, addr)
	if err != nil {
		return nil, err
	}
	if kind != HeapKindNT {
		return nil, fmt.Errorf("0x%x is a %s, not an NT Heap", addr, kind)
	}

	h := &Heap{Addr: addr}
	if h.Signature, err = readUint32AtValue(r, int64(addr)+heapOffSignature); err != nil {
		return nil, fmt.Errorf("Signature: %w", err)
	}
	if h.Flags, err = readUint32AtValue(r, int64(addr)+heapOffFlags); err != nil {
		return nil, fmt.Errorf("Flags: %w", err)
	}
	if h.ForceFlags, err = readUint32AtValue(r, int64(addr)+heapOffForceFlags); err != nil {
		return nil, fmt.Errorf("ForceFlags: %w", err)
	}
	if h.EncodeFlagMask, err = readUint32AtValue(r, int64(addr)+heapOffEncodeFlagMask); err != nil {
		return nil, fmt.Errorf("EncodeFlagMask: %w", err)
	}
	if _, err = r.ReadAt(h.Encoding[:], int64(addr)+heapOffEncoding); err != nil {
		return nil, fmt.Errorf("Encoding: %w", err)
	}
	if h.FrontEndHeapType, err = readUint8At(r, int64(addr)+heapOffFrontEndHeapType); err != nil {
		return nil, fmt.Errorf("FrontEndHeapType: %w", err)
	}
	if h.FrontEndHeap, err = readUint64At(r, int64(addr)+heapOffFrontEndHeap); err != nil {
		return nil, fmt.Errorf("FrontEndHeap: %w", err)
	}
	if h.BaseAddress, err = readUint64At(r, int64(addr)+heapOffBaseAddress); err != nil {
		return nil, fmt.Errorf("BaseAddress: %w", err)
	}
	if h.FirstEntry, err = readUint64At(r, int64(addr)+heapOffFirstEntry); err != nil {
		return nil, fmt.Errorf("FirstEntry: %w", err)
	}
	if h.LastValidEntry, err = readUint64At(r, int64(addr)+heapOffLastValidEntry); err != nil {
		return nil, fmt.Errorf("LastValidEntry: %w", err)
	}
	return h, nil
}

// EncodingActive reports whether this heap actually applies the
// header-encoding mitigation -- EncodeFlagMask is occasionally zero (e.g.
// explicitly disabled via HeapSetInformation), in which case
// DecodeHeapEntry should be called with encoding=nil instead of
// h.Encoding (an all-zero or stale key would silently corrupt every
// decode otherwise).
func (h *Heap) EncodingActive() bool { return h.EncodeFlagMask != 0 }

// encodingOrNil returns &h.Encoding if encoding is actually active, else
// nil -- the one-line helper every entry-decoding call in this package
// uses instead of repeating the EncodingActive() check.
func (h *Heap) encodingOrNil() *[16]byte {
	if h.EncodingActive() {
		return &h.Encoding
	}
	return nil
}

// WalkSegmentEntries decodes every _HEAP_ENTRY from firstEntry up to (not
// including) lastValidEntry -- the same chain `!heap -a` itself walks to
// print its block-by-block summary, which is what this decoder was
// cross-checked against (see heap_test.go). Stops early if an entry
// reports LastEntry(), and returns an error rather than looping forever if
// an entry's Size decodes to zero (a corrupted heap or a wrong/missing
// encoding key can't make forward progress otherwise).
//
// lastValidEntry is misleadingly named for this purpose: it marks the end
// of the segment's *reserved* address range, not its *committed* one, and
// `!heap -a` itself shows real heaps routinely ending with an uncommitted
// tail before that address is reached -- confirmed live against a real
// process while building this (see USAGE.md's heap walkthrough). A read
// failure partway through is therefore treated as "the committed entry
// chain ended here", not a hard error -- everything decoded up to that
// point is still returned.
func WalkSegmentEntries(r io.ReaderAt, firstEntry, lastValidEntry uint64, encoding *[16]byte) ([]HeapEntry, error) {
	var entries []HeapEntry
	addr := firstEntry
	for addr < lastValidEntry {
		e, err := ReadHeapEntry(r, addr, encoding)
		if err != nil {
			return entries, nil
		}
		entries = append(entries, e)
		if e.Size == 0 {
			return entries, fmt.Errorf("zero-size entry at 0x%x -- stopping to avoid an infinite loop (corrupted heap, or wrong/missing encoding key?)", addr)
		}
		if e.LastEntry() {
			break
		}
		addr = e.NextEntry()
	}
	return entries, nil
}

// WalkSegment0 walks this heap's embedded first segment -- the common
// case for any heap that hasn't grown past one segment. Use Segments +
// ReadSegmentRange + WalkSegmentEntries directly for a heap with more
// than one.
func (h *Heap) WalkSegment0(r io.ReaderAt) ([]HeapEntry, error) {
	return WalkSegmentEntries(r, h.FirstEntry, h.LastValidEntry, h.encodingOrNil())
}

// ReadSegmentRange reads a _HEAP_SEGMENT's FirstEntry/LastValidEntry/
// BaseAddress -- segAddr is anything Segments returns (Segment0's address
// equals the owning Heap's own address, since _HEAP embeds it at offset 0).
func ReadSegmentRange(r io.ReaderAt, segAddr uint64) (baseAddress, firstEntry, lastValidEntry uint64, err error) {
	if baseAddress, err = readUint64At(r, int64(segAddr)+heapOffBaseAddress); err != nil {
		return 0, 0, 0, fmt.Errorf("BaseAddress: %w", err)
	}
	if firstEntry, err = readUint64At(r, int64(segAddr)+heapOffFirstEntry); err != nil {
		return 0, 0, 0, fmt.Errorf("FirstEntry: %w", err)
	}
	if lastValidEntry, err = readUint64At(r, int64(segAddr)+heapOffLastValidEntry); err != nil {
		return 0, 0, 0, fmt.Errorf("LastValidEntry: %w", err)
	}
	return baseAddress, firstEntry, lastValidEntry, nil
}

// Segments returns the address of every _HEAP_SEGMENT belonging to this
// heap (walking the SegmentList LIST_ENTRY), including Segment0 (whose
// address is h.Addr itself, since _HEAP embeds its first segment at
// offset 0 -- the same struct-embedding pattern _HEAP_SEGMENT.Entry/
// SegmentSignature being literally at offset 0/0x10 of _HEAP relies on).
// Most heaps never grow past one segment; HeapCreate(0,0,0)-style growable
// heaps under sustained allocation pressure can.
func (h *Heap) Segments(r io.ReaderAt) ([]uint64, error) {
	headAddr := h.Addr + heapOffSegmentList
	cur, err := readUint64At(r, int64(headAddr))
	if err != nil {
		return nil, fmt.Errorf("reading SegmentList head at 0x%x: %w", headAddr, err)
	}

	var segments []uint64
	for cur != headAddr && cur != 0 {
		segments = append(segments, cur-heapOffSegmentListEntry)
		if len(segments) > 4096 {
			return segments, fmt.Errorf("SegmentList walk exceeded 4096 entries, stopping (corrupted list?)")
		}
		next, err := readUint64At(r, int64(cur)) // Flink is LIST_ENTRY's first field
		if err != nil {
			return segments, fmt.Errorf("walking SegmentList at 0x%x: %w", cur, err)
		}
		cur = next
	}
	return segments, nil
}

// WalkAllHeapEntries decodes every entry in every segment of h -- calls
// WalkSegment0 for the embedded first segment, then enumerates any
// additional registered segments from Segments() and walks each one. The
// returned slice is in address order, one segment after another. Segment
// walk errors are returned immediately (unlike the uncommitted-tail
// read-failure inside WalkSegmentEntries itself, which is graceful).
func (h *Heap) WalkAllHeapEntries(r io.ReaderAt) ([]HeapEntry, error) {
	entries, err := h.WalkSegment0(r)
	if err != nil {
		return nil, fmt.Errorf("segment0 walk: %w", err)
	}
	segs, err := h.Segments(r)
	if err != nil {
		return entries, fmt.Errorf("Segments(): %w", err)
	}
	for _, segAddr := range segs {
		if segAddr == h.Addr {
			continue // segment0 already walked above
		}
		_, first, last, err := ReadSegmentRange(r, segAddr)
		if err != nil {
			return entries, fmt.Errorf("ReadSegmentRange(0x%x): %w", segAddr, err)
		}
		more, err := WalkSegmentEntries(r, first, last, h.encodingOrNil())
		if err != nil {
			return entries, fmt.Errorf("walking segment 0x%x: %w", segAddr, err)
		}
		entries = append(entries, more...)
	}
	return entries, nil
}

// HeapStats summarises an NT Heap's entry layout across a slice of decoded
// entries (typically from WalkAllHeapEntries or a per-segment walk).
type HeapStats struct {
	TotalEntries int
	BusyEntries  int
	FreeEntries  int
	BusyBytes    uint64 // sum of UserSize() for busy entries
	FreeBytes    uint64 // sum of BlockSize() for free entries
}

// SummariseEntries computes a HeapStats over any entry slice -- useful after
// WalkAllHeapEntries, WalkSegment0, or any filtered subset.
func SummariseEntries(entries []HeapEntry) HeapStats {
	var s HeapStats
	s.TotalEntries = len(entries)
	for _, e := range entries {
		if e.Busy() {
			s.BusyEntries++
			s.BusyBytes += e.UserSize()
		} else {
			s.FreeEntries++
			s.FreeBytes += e.BlockSize()
		}
	}
	return s
}

// EntriesInRange returns the subset of entries whose header address falls
// within [lo, hi) -- useful for filtering down to a known allocation region
// (e.g. one specific segment) without re-walking.
func EntriesInRange(entries []HeapEntry, lo, hi uint64) []HeapEntry {
	var out []HeapEntry
	for _, e := range entries {
		if e.Addr >= lo && e.Addr < hi {
			out = append(out, e)
		}
	}
	return out
}

// EntriesWithUserData returns every entry whose UserData() (the address
// returned by HeapAlloc) equals any address in the addrs set -- direct
// reverse lookup from leaked heap pointer to its decoded entry.
func EntriesWithUserData(entries []HeapEntry, addrs ...uint64) []HeapEntry {
	set := make(map[uint64]struct{}, len(addrs))
	for _, a := range addrs {
		set[a] = struct{}{}
	}
	var out []HeapEntry
	for _, e := range entries {
		if _, ok := set[e.UserData()]; ok {
			out = append(out, e)
		}
	}
	return out
}

// AdjacentBusyPairs returns every pair of busy entries that are physically
// consecutive with no intervening free chunk -- i.e. entries[i+1].Addr ==
// entries[i].NextEntry(), both busy. This is the structural replacement for
// examples/heap_segment's "spray many, look for a 32-byte gap" technique:
// instead of spraying and comparing leaked addresses, read the allocator's
// own chain to find which two allocations are adjacent before overflowing.
//
// Note: the returned pairs are in chain order, not insertion order. On NT
// Heap the chain order matches allocation order within a given segment;
// whether that's also true for Segment Heap's small-block allocator is not
// yet confirmed on this build -- see heap_segment.go.
func AdjacentBusyPairs(entries []HeapEntry) [][2]HeapEntry {
	var pairs [][2]HeapEntry
	for i := 0; i+1 < len(entries); i++ {
		a, b := entries[i], entries[i+1]
		if a.Busy() && b.Busy() && b.Addr == a.NextEntry() {
			pairs = append(pairs, [2]HeapEntry{a, b})
		}
	}
	return pairs
}