Files
2026-07-18 21:37:15 +03:00

115 lines
3.2 KiB
Go

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" }