package winpwn

import "testing"

// TestSprayAndFindSeededTarget exercises the examples/heap_lfh shape: a
// single pre-seeded target (the freed victim's leaked address), spraying
// replacements until one happens to equal it.
func TestSprayAndFindSeededTarget(t *testing.T) {
	const victimAddr = uint64(0xdead0000)
	replacements := []uint64{0x1111, 0x2222, victimAddr, 0x3333}

	older, newer, attempts, ok, err := SprayAndFind(
		[]SprayResult[uint64]{{ID: -1, Key: victimAddr}},
		len(replacements),
		func(attempt int) (SprayResult[uint64], error) {
			return SprayResult[uint64]{ID: attempt, Key: replacements[attempt-1]}, nil
		},
		func(a, b uint64) bool { return a == b },
	)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if !ok {
		t.Fatal("expected a match")
	}
	if attempts != 3 {
		t.Errorf("attempts = %d, want 3 (the index where victimAddr reappears)", attempts)
	}
	if older.ID != -1 || older.Key != victimAddr {
		t.Errorf("older = %+v, want the seeded victim", older)
	}
	if newer.Key != victimAddr {
		t.Errorf("newer.Key = 0x%x, want 0x%x", newer.Key, victimAddr)
	}
}

// TestSprayAndFindAdjacentPair exercises the examples/heap_segment shape: no
// seed, searching every sprayed sample against every other for a relation
// (here, "exactly 32 apart").
func TestSprayAndFindAdjacentPair(t *testing.T) {
	const profileSize = 32
	addrs := []uint64{0x1000, 0x1080, 0x1300, 0x1300 + profileSize} // last two are 32 apart

	older, newer, _, ok, err := SprayAndFind(
		nil,
		len(addrs),
		func(attempt int) (SprayResult[uint64], error) {
			return SprayResult[uint64]{ID: attempt, Key: addrs[attempt-1]}, nil
		},
		func(a, b uint64) bool {
			d := int64(b) - int64(a)
			return d == profileSize || d == -profileSize
		},
	)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if !ok {
		t.Fatal("expected an adjacent pair to be found")
	}
	if older.ID != 3 || newer.ID != 4 {
		t.Errorf("got pair ids (%d, %d), want (3, 4)", older.ID, newer.ID)
	}
}

func TestSprayAndFindExhaustsAttempts(t *testing.T) {
	_, _, attempts, ok, err := SprayAndFind(
		nil,
		5,
		func(attempt int) (SprayResult[uint64], error) {
			return SprayResult[uint64]{ID: attempt, Key: uint64(attempt)}, nil
		},
		func(a, b uint64) bool { return false }, // never matches
	)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if ok {
		t.Fatal("expected no match")
	}
	if attempts != 5 {
		t.Errorf("attempts = %d, want 5 (maxAttempts exhausted)", attempts)
	}
}

func TestSprayAndFindPropagatesSprayError(t *testing.T) {
	wantErr := errSprayTest
	_, _, attempts, ok, err := SprayAndFind(
		nil,
		5,
		func(attempt int) (SprayResult[uint64], error) {
			if attempt == 2 {
				return SprayResult[uint64]{}, wantErr
			}
			return SprayResult[uint64]{ID: attempt, Key: uint64(attempt)}, nil
		},
		func(a, b uint64) bool { return false },
	)
	if err != wantErr {
		t.Fatalf("err = %v, want %v", err, wantErr)
	}
	if ok {
		t.Fatal("ok should be false on a spray error")
	}
	if attempts != 2 {
		t.Errorf("attempts = %d, want 2 (the attempt that errored)", attempts)
	}
}

var errSprayTest = errSprayTestSentinel{}

type errSprayTestSentinel struct{}

func (errSprayTestSentinel) Error() string { return "spray error" }