summaryrefslogtreecommitdiff
path: root/tests/e2e/dns_helpers.go
blob: 38c42a1e9f16426f5f584fc286f01b566d10a4b5 (plain)
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
}