package winpwn

import "testing"

// These tests pin heap_lfh.go's decoding against field *values* captured
// from live runs of examples/heap_lfh.exe (Windows build 10.0.26100) rather
// than raw byte dumps, the same style heap_test.go's
// TestReadHeapDecodesRealCapturedFields uses -- the offsets themselves were
// confirmed via `dt ntdll!_LFH_HEAP`/`dt ntdll!_HEAP_SUBSEGMENT` and the
// values via cdb reads cross-checked against winpwn's own output, not
// assumed.

func TestReadLFHBuckets(t *testing.T) {
	buf := make([]byte, lfhOffBuckets+lfhBucketCount*lfhBucketStride)
	// Bucket 2: 48-byte blocks (3 granularity units), the bucket that
	// actually served examples/heap_lfh's 32-byte Note/buffer allocations
	// once LFH took over that size class -- see FindLFHBucket's doc comment
	// for why 32-byte requests land in the 48-byte bucket, not the 32-byte
	// one.
	off := lfhOffBuckets + 2*lfhBucketStride
	buf[off] = 3
	buf[off+1] = 0
	buf[off+2] = 2 // SizeIndex

	buckets, err := ReadLFHBuckets(newByteReaderAt(buf), 0)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if buckets[2].BlockUnits != 3 || buckets[2].BlockSize() != 48 {
		t.Errorf("bucket[2] = %+v, want BlockUnits=3 BlockSize=48", buckets[2])
	}
	if buckets[0].BlockUnits != 0 {
		t.Errorf("bucket[0] should be unused, got %+v", buckets[0])
	}
}

func TestFindLFHBucketAccountsForBlockHeader(t *testing.T) {
	var buckets [lfhBucketCount]HeapBucket
	buckets[1] = HeapBucket{Index: 1, BlockUnits: 2} // 32 bytes total, 16 usable -- too small
	buckets[2] = HeapBucket{Index: 2, BlockUnits: 3} // 48 bytes total, 32 usable -- the real fit

	got, err := FindLFHBucket(buckets, 32)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if got.Index != 2 {
		t.Errorf("FindLFHBucket(32) chose bucket %d, want 2 (the naive BlockSize()>=wantSize check would wrongly pick bucket 1)", got.Index)
	}
}

func TestFindLFHBucketNoneActive(t *testing.T) {
	var buckets [lfhBucketCount]HeapBucket
	if _, err := FindLFHBucket(buckets, 32); err == nil {
		t.Fatal("expected an error when no bucket has BlockUnits set")
	}
}

func TestActiveSubsegment(t *testing.T) {
	const lfhHeapAddr = 0x1000
	buf := make([]byte, 0x2000)
	const segInfoAddr = 0x1500
	const subsegAddr = 0x1700
	binaryLEPutUint64(buf[lfhHeapAddr+lfhOffSegmentInfoArrays+2*8:], segInfoAddr)
	binaryLEPutUint64(buf[segInfoAddr+hlsiOffActiveSubsegment:], subsegAddr)

	got, err := ActiveSubsegment(newByteReaderAt(buf), lfhHeapAddr, 2)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if got != subsegAddr {
		t.Errorf("ActiveSubsegment = 0x%x, want 0x%x", got, subsegAddr)
	}

	if _, err := ActiveSubsegment(newByteReaderAt(buf), lfhHeapAddr, 5); err == nil {
		t.Fatal("expected an error for a bucket with no segment info")
	}
}

// TestReadLFHSubsegmentBusyBitmap pins ReadLFHSubsegment against the exact
// shape captured from a live run: bucket 2 (48-byte blocks), a 19-block
// subsegment, with one bit clear (an early, pre-LFH-activation allocation
// that never got a tracked slot -- see heap_lfh.go's package comment) and
// the rest busy.
func TestReadLFHSubsegmentBusyBitmap(t *testing.T) {
	const subsegAddr = 0x713480
	const userBlocksAddr = 0x713050
	buf := make([]byte, 0x800000)

	binaryLEPutUint64(buf[subsegAddr+subsegOffUserBlocks:], userBlocksAddr)
	buf[subsegAddr+subsegOffBlockSize] = 3 // granularity units -> 48 bytes
	buf[subsegAddr+subsegOffBlockCount] = 19

	const bitmapBufferAddr = 0x713200
	binaryLEPutUint64(buf[userBlocksAddr+userDataOffBitmapSize:], 19)
	binaryLEPutUint64(buf[userBlocksAddr+userDataOffBitmapBuffer:], bitmapBufferAddr)
	// bits: every slot busy except slot 5 (0-indexed) -- matches a real
	// captured run's BusyBitmap word.
	var word uint64 = 0
	for i := 0; i < 19; i++ {
		if i != 5 {
			word |= 1 << uint(i)
		}
	}
	binaryLEPutUint64(buf[bitmapBufferAddr:], word)

	subseg, err := ReadLFHSubsegment(newByteReaderAt(buf), subsegAddr)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if subseg.BlockSize != 48 || subseg.BlockCount != 19 {
		t.Errorf("BlockSize/BlockCount = %d/%d, want 48/19", subseg.BlockSize, subseg.BlockCount)
	}
	if subseg.UserBlocksAddr != userBlocksAddr {
		t.Errorf("UserBlocksAddr = 0x%x, want 0x%x", subseg.UserBlocksAddr, userBlocksAddr)
	}
	for i, busy := range subseg.Busy {
		want := i != 5
		if busy != want {
			t.Errorf("Busy[%d] = %v, want %v", i, busy, want)
		}
	}
}

// TestCalibrateLFHFirstBlockOffsetMatchesLiveAddress pins
// CalibrateLFHFirstBlockOffset/BlockAddress/SlotOf against a real
// before/after-free pair captured live: freeing the note at 0x7132b0
// flipped exactly Busy[11] from true to false, and slot 11's address by
// BlockAddress must equal 0x7132b0 -- not 0x713280 (one block short), which
// is what the naive "(addr-UserBlocksAddr) % BlockSize" formula gives
// before the +BlockSize correction (see CalibrateLFHFirstBlockOffset's doc
// comment for why that's wrong: it lands on _HEAP_USERDATA_HEADER's own
// reserved region, one block before any real, bit-tracked slot).
func TestCalibrateLFHFirstBlockOffsetMatchesLiveAddress(t *testing.T) {
	subseg := &LFHSubsegment{
		UserBlocksAddr: 0x713050,
		BlockSize:      48,
		BlockCount:     19,
	}
	const victimAddr = 0x7132b0

	off := CalibrateLFHFirstBlockOffset(subseg, victimAddr)
	if off != 80 {
		t.Fatalf("CalibrateLFHFirstBlockOffset = %d, want 80", off)
	}
	if got := subseg.BlockAddress(off, 11); got != victimAddr {
		t.Errorf("BlockAddress(80, 11) = 0x%x, want 0x%x", got, victimAddr)
	}
	idx, ok := subseg.SlotOf(off, victimAddr)
	if !ok || idx != 11 {
		t.Errorf("SlotOf = (%d, %v), want (11, true)", idx, ok)
	}
	if _, ok := subseg.SlotOf(off, 0xdeadbeef); ok {
		t.Error("SlotOf should reject an address outside the subsegment")
	}
}