package winpwn

import (
	"errors"
	"testing"
)

var errOutOfRange = errors.New("byteReaderAt: out of range")

// TestDecodeHeapEntryMatchesLiveGroundTruth pins DecodeHeapEntry against
// three real _HEAP_ENTRY headers + their heap's real Encoding key,
// captured byte-for-byte from a live process (a tiny HeapAlloc(256)/
// HeapAlloc(4)/HeapAlloc(16) probe on this machine, Windows build 10.0.26100)
// and cross-checked against WinDbg's own `!heap -a` ground-truth listing
// before being trusted -- see heap.go's top comment for the full story.
//
// This is the test that caught the real bug worth remembering: the first
// version of UserSize subtracted HeapEntrySize *again* on top of
// UnusedBytes (i.e. assumed UnusedBytes was padding *after* a separately-
// accounted-for header), which silently produced a UserSize 16 bytes
// smaller than reality for every entry. !heap -a's own "(requested size)"
// column is what caught it -- UnusedBytes already bakes the header in.
func TestDecodeHeapEntryMatchesLiveGroundTruth(t *testing.T) {
	// Heap.Encoding, captured at heap+0x80 on the live probe.
	encoding := [16]byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x1e, 0xb6, 0x56, 0xf8, 0xe5, 0xd3, 0x00, 0x00}

	cases := []struct {
		name         string
		addr         uint64
		raw          [16]byte
		wantSize     uint16 // granularity units, cross-checked against `!heap -a`'s byte-size column / HeapEntrySize
		wantFlags    uint8
		wantPrevSize uint16
		wantUnused   uint8
		wantUserSize uint64 // the requested size `!heap -a` printed in parens
	}{
		{
			// !heap -a: 00000000000d0850: 00110 . 00110 [101] - busy (100)
			name:         "256-byte allocation, no slack",
			addr:         0xd0850,
			raw:          [16]byte{0, 0, 0, 0, 0, 0, 0, 0, 0x0f, 0xb6, 0x57, 0xe8, 0xf4, 0xd3, 0x00, 0x10},
			wantSize:     17, // 17*0x10 = 0x110 = 272
			wantFlags:    HeapEntryBusy,
			wantPrevSize: 17, // previous entry (0xd0740) was also 0x110 bytes per !heap -a
			wantUnused:   16,
			wantUserSize: 0x100,
		},
		{
			// !heap -a: 00000000000d0d70: 00050 . 00020 [101] - busy (4)
			name:         "4-byte allocation, lots of slack",
			addr:         0xd0d70,
			raw:          [16]byte{0, 0, 0, 0, 0, 0, 0, 0, 0x1c, 0xb6, 0x57, 0xfb, 0xe0, 0xd3, 0x00, 0x1c},
			wantSize:     2, // 2*0x10 = 0x20
			wantFlags:    HeapEntryBusy,
			wantPrevSize: 5, // previous entry (0xd0d20) was 0x50 bytes per !heap -a
			wantUnused:   28,
			wantUserSize: 4,
		},
		{
			// !heap -a: 00000000000d0d90: 00020 . 00020 [101] - busy (10)
			name:         "16-byte allocation, header-only slack",
			addr:         0xd0d90,
			raw:          [16]byte{0, 0, 0, 0, 0, 0, 0, 0, 0x1c, 0xb6, 0x57, 0xfb, 0xe7, 0xd3, 0x00, 0x10},
			wantSize:     2,
			wantFlags:    HeapEntryBusy,
			wantPrevSize: 2, // previous entry (0xd0d70) was also 0x20 bytes
			wantUnused:   16,
			wantUserSize: 16,
		},
	}

	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			e := DecodeHeapEntry(c.addr, c.raw, &encoding)
			if e.Size != c.wantSize {
				t.Errorf("Size = %d, want %d", e.Size, c.wantSize)
			}
			if e.Flags != c.wantFlags {
				t.Errorf("Flags = 0x%x, want 0x%x", e.Flags, c.wantFlags)
			}
			if e.PreviousSize != c.wantPrevSize {
				t.Errorf("PreviousSize = %d, want %d", e.PreviousSize, c.wantPrevSize)
			}
			if e.UnusedBytes != c.wantUnused {
				t.Errorf("UnusedBytes = %d, want %d", e.UnusedBytes, c.wantUnused)
			}
			if got := e.UserSize(); got != c.wantUserSize {
				t.Errorf("UserSize() = 0x%x, want 0x%x", got, c.wantUserSize)
			}
			if !e.Busy() {
				t.Error("Busy() = false, want true (all three live samples were busy)")
			}
			if e.UserData() != c.addr+HeapEntrySize {
				t.Errorf("UserData() = 0x%x, want 0x%x", e.UserData(), c.addr+HeapEntrySize)
			}
		})
	}
}

func TestDetectHeapKind(t *testing.T) {
	ntHeap := make([]byte, 0x20)
	binaryLEPutUint32(ntHeap[0x10:], heapSignatureNT)

	segHeap := make([]byte, 0x20)
	binaryLEPutUint32(segHeap[0x10:], heapSignatureSegment)

	garbage := make([]byte, 0x20)

	cases := []struct {
		name string
		buf  []byte
		want HeapKind
		ok   bool
	}{
		{"nt heap", ntHeap, HeapKindNT, true},
		{"segment heap", segHeap, HeapKindSegment, true},
		{"garbage", garbage, HeapKindUnknown, false},
	}

	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			kind, err := DetectHeapKind(newByteReaderAt(c.buf), 0)
			if c.ok && err != nil {
				t.Fatalf("unexpected error: %v", err)
			}
			if !c.ok && err == nil {
				t.Fatal("expected an error for an unrecognized signature")
			}
			if kind != c.want {
				t.Errorf("kind = %v, want %v", kind, c.want)
			}
		})
	}
}

// TestReadHeapDecodesRealCapturedFields pins ReadHeap against the actual
// field bytes captured from the same live probe heap.go's top comment and
// heap_test.go's ground-truth test describe (Flags=0x2/HEAP_GROWABLE,
// EncodeFlagMask=0x100000, Signature=0xeeffeeff -- deliberately distinct
// from _HEAP_SEGMENT's 0xffeeffee, confirmed with `db` against the live
// process rather than assumed equal).
func TestReadHeapDecodesRealCapturedFields(t *testing.T) {
	buf := make([]byte, 0x300)
	binaryLEPutUint32(buf[0x10:], heapSignatureNT) // _HEAP_SEGMENT.SegmentSignature
	binaryLEPutUint32(buf[0x70:], 0x00000002)      // Flags = HEAP_GROWABLE
	binaryLEPutUint32(buf[0x7c:], 0x00100000)      // EncodeFlagMask
	copy(buf[0x80:0x90], []byte{0, 0, 0, 0, 0, 0, 0, 0, 0x1e, 0xb6, 0x56, 0xf8, 0xe5, 0xd3, 0x00, 0x00})
	binaryLEPutUint32(buf[0x98:], 0xeeffeeff) // _HEAP.Signature
	buf[0x1a2] = FrontEndHeapNone
	binaryLEPutUint64(buf[0x30:], 0x6f0000) // BaseAddress
	binaryLEPutUint64(buf[0x40:], 0x6f0740) // FirstEntry
	binaryLEPutUint64(buf[0x48:], 0x7ef000) // LastValidEntry

	h, err := ReadHeap(newByteReaderAt(buf), 0)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if h.Signature != 0xeeffeeff {
		t.Errorf("Signature = 0x%x, want 0xeeffeeff", h.Signature)
	}
	if h.Flags != 0x2 {
		t.Errorf("Flags = 0x%x, want 0x2", h.Flags)
	}
	if !h.EncodingActive() {
		t.Error("EncodingActive() = false, want true (EncodeFlagMask is nonzero)")
	}
	if h.FrontEndHeapType != FrontEndHeapNone {
		t.Errorf("FrontEndHeapType = %d, want %d", h.FrontEndHeapType, FrontEndHeapNone)
	}
	if h.BaseAddress != 0x6f0000 || h.FirstEntry != 0x6f0740 || h.LastValidEntry != 0x7ef000 {
		t.Errorf("BaseAddress/FirstEntry/LastValidEntry = 0x%x/0x%x/0x%x", h.BaseAddress, h.FirstEntry, h.LastValidEntry)
	}
}

func TestReadHeapRejectsSegmentHeap(t *testing.T) {
	buf := make([]byte, 0x20)
	binaryLEPutUint32(buf[0x10:], heapSignatureSegment)
	if _, err := ReadHeap(newByteReaderAt(buf), 0); err == nil {
		t.Fatal("expected ReadHeap to reject a Segment Heap signature")
	}
}

func TestHeapSegmentsSingleSegmentIsEmpty(t *testing.T) {
	// A heap with only the embedded Segment0: SegmentList's Flink/Blink
	// both point back at the list head itself (heapOffSegmentList), the
	// standard "empty list" LIST_ENTRY representation.
	buf := make([]byte, 0x200)
	headAddr := uint64(heapOffSegmentList)
	binaryLEPutUint64(buf[heapOffSegmentList:], headAddr)
	binaryLEPutUint64(buf[heapOffSegmentList+8:], headAddr)

	h := &Heap{Addr: 0}
	segs, err := h.Segments(newByteReaderAt(buf))
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if len(segs) != 0 {
		t.Errorf("got %d segments, want 0 for an empty SegmentList", len(segs))
	}
}

func TestWalkSegmentEntriesStopsAtLastEntry(t *testing.T) {
	// Three unencoded entries (no Encoding key, EncodeFlagMask=0 case):
	// 32 bytes busy, 16 bytes free, 16 bytes busy+LastEntry.
	buf := make([]byte, 64)
	putEntry := func(off int, size uint16, flags uint8, prevSize uint16) {
		buf[off+8] = byte(size)
		buf[off+9] = byte(size >> 8)
		buf[off+10] = flags
		buf[off+12] = byte(prevSize)
		buf[off+13] = byte(prevSize >> 8)
	}
	putEntry(0, 2, HeapEntryBusy, 0)
	putEntry(32, 1, 0, 2)
	putEntry(48, 1, HeapEntryBusy|HeapEntryLastEntry, 1)

	entries, err := WalkSegmentEntries(newByteReaderAt(buf), 0, 64, nil)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if len(entries) != 3 {
		t.Fatalf("got %d entries, want 3", len(entries))
	}
	if entries[0].Addr != 0 || entries[0].BlockSize() != 32 || !entries[0].Busy() {
		t.Errorf("entry 0 = %+v", entries[0])
	}
	if entries[1].Addr != 32 || entries[1].BlockSize() != 16 || entries[1].Busy() {
		t.Errorf("entry 1 = %+v", entries[1])
	}
	if entries[2].Addr != 48 || !entries[2].LastEntry() {
		t.Errorf("entry 2 = %+v", entries[2])
	}
}

func TestWalkSegmentEntriesZeroSizeIsAnError(t *testing.T) {
	buf := make([]byte, 32) // entry at 0 decodes to Size=0 -- can't make progress
	_, err := WalkSegmentEntries(newByteReaderAt(buf), 0, 32, nil)
	if err == nil {
		t.Fatal("expected an error for a zero-size entry")
	}
}

// byteReaderAt adapts a plain []byte to io.ReaderAt for synthetic tests,
// the same role buildSyntheticMinidump's bytes.Reader plays in
// minidump_test.go.
type byteReaderAt struct{ buf []byte }

func newByteReaderAt(buf []byte) *byteReaderAt { return &byteReaderAt{buf: buf} }

func (b *byteReaderAt) ReadAt(p []byte, off int64) (int, error) {
	if off < 0 || int(off) > len(b.buf) {
		return 0, errOutOfRange
	}
	n := copy(p, b.buf[off:])
	if n < len(p) {
		return n, errOutOfRange
	}
	return n, nil
}

func binaryLEPutUint32(b []byte, v uint32) {
	b[0] = byte(v)
	b[1] = byte(v >> 8)
	b[2] = byte(v >> 16)
	b[3] = byte(v >> 24)
}

func binaryLEPutUint64(b []byte, v uint64) {
	for i := 0; i < 8; i++ {
		b[i] = byte(v >> (8 * i))
	}
}