package winpwn

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

// This file is the LFH (Low Fragmentation Heap) layer on top of heap.go's
// plain NT Heap decoder -- the layer ROADMAP.md's Phase 9 plan called out
// as the real acceptance test: replace examples/heap_lfh's empirically-
// discovered "free the most recently allocated object" heuristic with a
// structural answer to "which block in the active subsegment is actually
// free right now", read directly off the live process.
//
// Offsets confirmed via `dt ntdll!_LFH_HEAP` and friends on this machine's
// build (10.0.26100) -- same methodology as heap.go, see its top comment.
// Live-validated end to end against examples/heap_lfh.exe itself: spray a
// batch of Notes, free one, calibrate against a known address, and confirm
// ReadLFHSubsegment's BusyBitmap reports exactly that block's slot as free
// and every other allocated slot as busy -- see USAGE.md's heap walkthrough
// for the exact run.

const (
	lfhOffBuckets           = 0x2a4 // [129]_HEAP_BUCKET, 4 bytes each, inside _LFH_HEAP
	lfhOffSegmentInfoArrays = 0x4a8 // [129]Ptr64 _HEAP_LOCAL_SEGMENT_INFO
	lfhBucketCount          = 129
	lfhBucketStride         = 4

	hlsiOffActiveSubsegment = 0x008 // _HEAP_LOCAL_SEGMENT_INFO.ActiveSubsegment

	subsegOffUserBlocks = 0x008 // _HEAP_SUBSEGMENT.UserBlocks
	subsegOffBlockSize  = 0x024 // uint16, granularity units -- same scheme as HeapBucket.BlockUnits, NOT raw bytes (see below)
	subsegOffBlockCount = 0x028 // uint16, blocks in this subsegment

	userDataOffBitmapSize   = 0x020 // uint64, _RTL_BITMAP_EX.SizeOfBitMap (bits)
	userDataOffBitmapBuffer = 0x028 // ptr64, _RTL_BITMAP_EX.Buffer
)

// HeapBucket is one entry of _LFH_HEAP.Buckets -- which fixed block size
// this size class hands out.
type HeapBucket struct {
	Index      int
	BlockUnits uint16 // granularity units; BlockSize() = BlockUnits*HeapEntrySize
	SizeIndex  uint8
	RawFlags   uint8
}

func (b HeapBucket) BlockSize() uint64 { return uint64(b.BlockUnits) * HeapEntrySize }

// ReadLFHBuckets reads every entry of lfhHeapAddr's (a Heap.FrontEndHeap
// pointer) Buckets array.
func ReadLFHBuckets(r io.ReaderAt, lfhHeapAddr uint64) ([lfhBucketCount]HeapBucket, error) {
	var buckets [lfhBucketCount]HeapBucket
	var buf [lfhBucketCount * lfhBucketStride]byte
	if _, err := r.ReadAt(buf[:], int64(lfhHeapAddr)+lfhOffBuckets); err != nil {
		return buckets, fmt.Errorf("reading Buckets array: %w", err)
	}
	for i := 0; i < lfhBucketCount; i++ {
		off := i * lfhBucketStride
		buckets[i] = HeapBucket{
			Index:      i,
			BlockUnits: binary.LittleEndian.Uint16(buf[off : off+2]),
			SizeIndex:  buf[off+2],
			RawFlags:   buf[off+3],
		}
	}
	return buckets, nil
}

// FindLFHBucket returns the smallest bucket that can actually serve a
// wantSize-byte HeapAlloc request once LFH owns this size class. Buckets
// with BlockUnits==0 are unused size classes (LFH only activates a bucket
// index after enough same-size requests) and are skipped.
//
// Confirmed live against examples/heap_lfh.exe, and worth recording
// because the "obvious" version (BlockSize() >= wantSize, no header
// accounted for) is wrong: an LFH block's BlockSize already includes its
// own 16-byte _HEAP_ENTRY-shaped header the same way a plain NT Heap
// entry's does, so a 32-byte Note allocation is actually routed to the
// 48-byte bucket (BlockSize 48, 48-16=32 usable), not the 32-byte one
// (which only has 16 bytes usable after its own header) -- verified by
// reading SegmentInfoArrays directly and seeing which bucket index
// actually had a live subsegment for a heap doing nothing but 32-byte
// allocations.
func FindLFHBucket(buckets [lfhBucketCount]HeapBucket, wantSize uint64) (HeapBucket, error) {
	for _, b := range buckets {
		if b.BlockUnits == 0 {
			continue
		}
		if b.BlockSize() >= wantSize+HeapEntrySize {
			return b, nil
		}
	}
	return HeapBucket{}, fmt.Errorf("no active LFH bucket covers %d bytes (has LFH actually taken over this size class yet?)", wantSize)
}

// ActiveSubsegment returns the address of bucketIndex's currently-active
// _HEAP_SUBSEGMENT -- the subsegment LFH is issuing new blocks from right
// now for that size class. This is the address every grooming attempt
// (examples/heap_lfh's spray loop, structurally) is actually targeting.
func ActiveSubsegment(r io.ReaderAt, lfhHeapAddr uint64, bucketIndex int) (uint64, error) {
	if bucketIndex < 0 || bucketIndex >= lfhBucketCount {
		return 0, fmt.Errorf("bucket index %d out of range [0,%d)", bucketIndex, lfhBucketCount)
	}
	arrAddr := lfhHeapAddr + lfhOffSegmentInfoArrays + uint64(bucketIndex)*8
	segInfoAddr, err := readUint64At(r, int64(arrAddr))
	if err != nil {
		return 0, fmt.Errorf("reading SegmentInfoArrays[%d]: %w", bucketIndex, err)
	}
	if segInfoAddr == 0 {
		return 0, fmt.Errorf("bucket %d has no segment info yet (LFH hasn't allocated from this size class)", bucketIndex)
	}
	subsegAddr, err := readUint64At(r, int64(segInfoAddr)+hlsiOffActiveSubsegment)
	if err != nil {
		return 0, fmt.Errorf("reading ActiveSubsegment: %w", err)
	}
	if subsegAddr == 0 {
		return 0, fmt.Errorf("bucket %d's segment info has no active subsegment", bucketIndex)
	}
	return subsegAddr, nil
}

// LFHSubsegment is a decoded _HEAP_SUBSEGMENT: a fixed-size-block arena LFH
// is handing blocks out of for one bucket. BlockSize/BlockCount/Busy are
// read directly off the subsegment's own bookkeeping with no decoding
// ambiguity; turning slot index into an address needs one calibration step
// first -- see BlockAddress.
type LFHSubsegment struct {
	Addr           uint64
	UserBlocksAddr uint64
	BlockSize      uint64 // bytes; subsegOffBlockSize*HeapEntrySize, see its comment
	BlockCount     int
	Busy           []bool // Busy[i] is slot i's state, read straight from BusyBitmap
}

// ReadLFHSubsegment decodes subsegAddr's _HEAP_SUBSEGMENT, with Busy/Free
// for every slot read directly from the subsegment's BusyBitmap -- the
// structural replacement for examples/heap_lfh's "spray and see which
// leaked address repeats" technique: this answers "which slot is free" by
// reading the allocator's own bookkeeping instead of inferring it from
// outside.
//
// What it deliberately does NOT do: compute slot addresses. The natural
// place for that, _HEAP_USERDATA_HEADER.EncodedOffsets, decoded as plain
// FirstAllocationOffset(u16)|BlockStride(u16), produced a 29025-byte stride
// for a heap doing nothing but 48-byte allocations -- it's genuinely
// encoded (XORed against something derived from
// RtlpLowFragHeapRandomData/RtlpInitializeLfhRandomDataArray per this
// repo's heap/Deterministic_LFH-master reference material), and the key
// wasn't recovered in this pass. Use BlockAddress + CalibrateLFHFirstBlockOffset
// instead: calibrate against any one address you already know (which is
// also just how this kind of exploitation actually works in practice --
// correlating against a leak, not deriving addresses from nothing).
func ReadLFHSubsegment(r io.ReaderAt, subsegAddr uint64) (*LFHSubsegment, error) {
	userBlocksAddr, err := readUint64At(r, int64(subsegAddr)+subsegOffUserBlocks)
	if err != nil {
		return nil, fmt.Errorf("reading UserBlocks: %w", err)
	}
	if userBlocksAddr == 0 {
		return nil, fmt.Errorf("subsegment at 0x%x has no UserBlocks (not yet committed?)", subsegAddr)
	}

	var blockSizeUnits, blockCountBuf [2]byte
	if _, err := r.ReadAt(blockSizeUnits[:], int64(subsegAddr)+subsegOffBlockSize); err != nil {
		return nil, fmt.Errorf("reading BlockSize: %w", err)
	}
	if _, err := r.ReadAt(blockCountBuf[:], int64(subsegAddr)+subsegOffBlockCount); err != nil {
		return nil, fmt.Errorf("reading BlockCount: %w", err)
	}
	blockSize := uint64(binary.LittleEndian.Uint16(blockSizeUnits[:])) * HeapEntrySize
	blockCount := binary.LittleEndian.Uint16(blockCountBuf[:])
	if blockSize == 0 {
		return nil, fmt.Errorf("decoded BlockSize is 0 at 0x%x -- wrong offset or unsupported build (see heap.go's top comment)", subsegAddr)
	}

	bitmapSizeBits, err := readUint64At(r, int64(userBlocksAddr)+userDataOffBitmapSize)
	if err != nil {
		return nil, fmt.Errorf("reading BusyBitmap.SizeOfBitMap: %w", err)
	}
	bitmapBufferAddr, err := readUint64At(r, int64(userBlocksAddr)+userDataOffBitmapBuffer)
	if err != nil {
		return nil, fmt.Errorf("reading BusyBitmap.Buffer: %w", err)
	}
	if bitmapSizeBits < uint64(blockCount) {
		return nil, fmt.Errorf("BusyBitmap covers %d bits but BlockCount is %d", bitmapSizeBits, blockCount)
	}

	busy := make([]bool, blockCount)
	for i := uint16(0); i < blockCount; i++ {
		b, err := readBitmapBit(r, bitmapBufferAddr, uint64(i))
		if err != nil {
			return nil, fmt.Errorf("reading BusyBitmap bit %d: %w", i, err)
		}
		busy[i] = b
	}

	return &LFHSubsegment{
		Addr:           subsegAddr,
		UserBlocksAddr: userBlocksAddr,
		BlockSize:      blockSize,
		BlockCount:     int(blockCount),
		Busy:           busy,
	}, nil
}

// BlockAddress returns the address of slot index i, given firstBlockOffset
// -- the byte offset of slot 0 relative to UserBlocksAddr. Get
// firstBlockOffset from CalibrateLFHFirstBlockOffset once per subsegment;
// it's constant across every slot of the same subsegment.
func (s *LFHSubsegment) BlockAddress(firstBlockOffset uint64, index int) uint64 {
	return s.UserBlocksAddr + firstBlockOffset + uint64(index)*s.BlockSize
}

// SlotOf returns the slot index of knownAddr within s, given an already-
// calibrated firstBlockOffset (see CalibrateLFHFirstBlockOffset), or false
// if knownAddr doesn't land in this subsegment's range at all.
func (s *LFHSubsegment) SlotOf(firstBlockOffset, knownAddr uint64) (index int, ok bool) {
	base := s.UserBlocksAddr + firstBlockOffset
	if knownAddr < base {
		return 0, false
	}
	off := knownAddr - base
	if off%s.BlockSize != 0 {
		return 0, false
	}
	idx := off / s.BlockSize
	if idx >= uint64(s.BlockCount) {
		return 0, false
	}
	return int(idx), true
}

// CalibrateLFHFirstBlockOffset computes BlockAddress's firstBlockOffset
// from one address you already know lies inside this subsegment -- e.g.
// one of your own freshly leaked allocations. _HEAP_USERDATA_HEADER does
// store this value (as part of the encoded EncodedOffsets field) but it's
// genuinely obfuscated and wasn't decoded in this pass; calibrating
// against a known address sidesteps that entirely, and is also simply how
// you'd correlate against a real target in practice.
//
// The +BlockSize matters and was the second real bug found empirically: the
// naive "(addr-UserBlocksAddr) % BlockSize" residue is the right modular
// class but the wrong absolute offset -- it points at a reserved region
// belonging to _HEAP_USERDATA_HEADER itself (one block-size's worth of
// space with no corresponding BusyBitmap bit at all), one full block before
// where Busy[0]'s real address actually is. Confirmed by freeing a known
// address and watching which bit actually flipped: it was the bit for
// address_index-1 under the naive offset, i.e. exactly one block short.
func CalibrateLFHFirstBlockOffset(s *LFHSubsegment, knownBlockAddr uint64) uint64 {
	return (knownBlockAddr-s.UserBlocksAddr)%s.BlockSize + s.BlockSize
}

// AllSubsegments returns the addresses of every non-null _HEAP_SUBSEGMENT
// associated with bucketIndex in lfhHeapAddr's _LFH_HEAP, not just the
// currently active one. The active subsegment is always first (index 0)
// when present; any additional cached/full ones follow in the order they
// appear in the SegmentInfoArrays chain.
//
// LFH doesn't maintain a traditional linked list of subsegments per
// bucket -- it uses _HEAP_LOCAL_SEGMENT_INFO which has a single
// ActiveSubsegment pointer and an optional cached slot. This walk follows
// ActiveSubsegment (via ActiveSubsegment()) and leaves deeper enumeration
// (e.g. walking the CachedItems or InfoArrays of retired subsegments) for
// a future pass where that extra complexity is actually needed by a CTF task.
func AllSubsegments(r io.ReaderAt, lfhHeapAddr uint64, bucketIndex int) ([]uint64, error) {
	if bucketIndex < 0 || bucketIndex >= lfhBucketCount {
		return nil, fmt.Errorf("bucket index %d out of range [0,%d)", bucketIndex, lfhBucketCount)
	}
	arrAddr := lfhHeapAddr + lfhOffSegmentInfoArrays + uint64(bucketIndex)*8
	segInfoAddr, err := readUint64At(r, int64(arrAddr))
	if err != nil {
		return nil, fmt.Errorf("reading SegmentInfoArrays[%d]: %w", bucketIndex, err)
	}
	if segInfoAddr == 0 {
		return nil, nil // bucket not yet activated, not an error
	}

	var subsegAddrs []uint64
	active, err := readUint64At(r, int64(segInfoAddr)+hlsiOffActiveSubsegment)
	if err != nil {
		return nil, fmt.Errorf("reading ActiveSubsegment: %w", err)
	}
	if active != 0 {
		subsegAddrs = append(subsegAddrs, active)
	}

	// _HEAP_LOCAL_SEGMENT_INFO.CachedItems (an array of 2 ptr slots
	// immediately after ActiveSubsegment for this build -- empirically
	// observed at hlsiOffActiveSubsegment+8 and +16; if it's wrong they'll
	// simply be 0 and get skipped).
	for i := 1; i <= 2; i++ {
		cached, err := readUint64At(r, int64(segInfoAddr)+hlsiOffActiveSubsegment+int64(i)*8)
		if err != nil || cached == 0 || cached == active {
			continue
		}
		subsegAddrs = append(subsegAddrs, cached)
	}

	return subsegAddrs, nil
}

// readBitmapBit reads bit index bitIndex of an RTL_BITMAP_EX-style bitmap
// (an array of 64-bit words starting at bufferAddr).
func readBitmapBit(r io.ReaderAt, bufferAddr uint64, bitIndex uint64) (bool, error) {
	word, err := readUint64At(r, int64(bufferAddr)+int64(bitIndex/64)*8)
	if err != nil {
		return false, err
	}
	return (word>>(bitIndex%64))&1 != 0, nil
}