1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
|
package e2e
import (
"net"
"testing"
"time"
)
// MockDNSServer is a minimal UDP server that responds to any DNS query
// with a hardcoded A record for example.com.
type MockDNSServer struct {
conn *net.UDPConn
}
func StartMockDNSServer(t *testing.T) (*MockDNSServer, int) {
addr, err := net.ResolveUDPAddr("udp", "127.0.0.1:0")
if err != nil {
t.Fatalf("failed to resolve UDP address: %v", err)
}
conn, err := net.ListenUDP("udp", addr)
if err != nil {
t.Fatalf("failed to start mock DNS server: %v", err)
}
return &MockDNSServer{conn: conn}, conn.LocalAddr().(*net.UDPAddr).Port
}
func (s *MockDNSServer) Close() {
_ = s.conn.Close()
}
// ListenAndRespond runs in a loop and responds to the first valid-looking DNS query.
// It returns the response as a channel to allow the test to synchronize.
func (s *MockDNSServer) ListenAndRespond(timeout time.Duration) (chan bool, error) {
respChan := make(chan bool, 1)
go func() {
buf := make([]byte, 1024)
_ = s.conn.SetReadDeadline(time.Now().Add(timeout))
n, remoteAddr, err := s.conn.ReadFromUDP(buf)
if err != nil {
respChan <- false
return
}
if n < 4 {
respChan <- false
return
}
// Echo the packet back to the client as a basic "I heard you" confirmation.
_, _ = s.conn.WriteToUDP(buf[:n], remoteAddr)
respChan <- true
}()
return respChan, nil
}
|
