/
dns_local_server.go
161 lines (128 loc) · 5.23 KB
/
dns_local_server.go
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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
package dns_helper
import (
"fmt"
"log"
"net"
"strings"
"testing"
"time"
"github.com/miekg/dns"
)
var testDomain = "gruntwork.io"
// dnsDatabase stores a collection of DNSQuery with their respective DNSAnswers, to be used by a local dnsTestServer
type dnsDatabase map[DNSQuery]DNSAnswers
// dnsTestServer helper for testing this package using local DNS nameservers with test records
type dnsTestServer struct {
Server *dns.Server
DNSDatabase dnsDatabase
DNSDatabaseRetry dnsDatabase
}
// newDNSTestServer returns a new instance of dnsTestServer
func newDNSTestServer(server *dns.Server) *dnsTestServer {
return &dnsTestServer{Server: server, DNSDatabase: make(dnsDatabase), DNSDatabaseRetry: make(dnsDatabase)}
}
// Address returns the host:port string of the server listener
func (s *dnsTestServer) Address() string {
return s.Server.PacketConn.LocalAddr().String()
}
// AddEntryToDNSDatabase adds DNSAnswers to the DNSQuery in the database of the server
func (s *dnsTestServer) AddEntryToDNSDatabase(q DNSQuery, a DNSAnswers) {
s.DNSDatabase[q] = append(s.DNSDatabase[q], a...)
}
// AddEntryToDNSDatabaseRetry adds DNSAnswers to the DNSQuery in the database used when retrying
func (s *dnsTestServer) AddEntryToDNSDatabaseRetry(q DNSQuery, a DNSAnswers) {
s.DNSDatabaseRetry[q] = append(s.DNSDatabaseRetry[q], a...)
}
// setupTestDNSServers runs and returns 2x local dnsTestServer, initialized with NS records for the testDomain pointing to themselves
// it uses a handler that will send replies stored in their internal DNSDatabase
func setupTestDNSServers(t *testing.T) (s1, s2 *dnsTestServer) {
s1 = runTestDNSServer(t, "0")
s2 = runTestDNSServer(t, "0")
q := DNSQuery{"NS", testDomain}
a := DNSAnswers{{"NS", s1.Address() + "."}, {"NS", s2.Address() + "."}}
s1.AddEntryToDNSDatabase(q, a)
s2.AddEntryToDNSDatabase(q, a)
s1.Server.Handler.(*dns.ServeMux).HandleFunc(testDomain+".", func(w dns.ResponseWriter, r *dns.Msg) {
stdDNSHandler(t, w, r, s1, false)
})
s2.Server.Handler.(*dns.ServeMux).HandleFunc(testDomain+".", func(w dns.ResponseWriter, r *dns.Msg) {
stdDNSHandler(t, w, r, s2, true)
})
return s1, s2
}
// setupTestDNSServersRetry runs and returns 2x local dnsTestServer, initialized with NS records for the testDomain pointing to themselves
// it uses a handler that will send replies stored in their internal DNSDatabase, and then switch to their DNSDatabaseRetry after some time
func setupTestDNSServersRetry(t *testing.T) (s1, s2 *dnsTestServer) {
s1 = runTestDNSServer(t, "0")
s2 = runTestDNSServer(t, "0")
q := DNSQuery{"NS", testDomain}
a := DNSAnswers{{"NS", s1.Address() + "."}, {"NS", s2.Address() + "."}}
s1.AddEntryToDNSDatabase(q, a)
s2.AddEntryToDNSDatabase(q, a)
s1.AddEntryToDNSDatabaseRetry(q, a)
s2.AddEntryToDNSDatabaseRetry(q, a)
s1.Server.Handler.(*dns.ServeMux).HandleFunc(testDomain+".", func(w dns.ResponseWriter, r *dns.Msg) {
retryDNSHandler(t, w, r, s1, false)
})
s2.Server.Handler.(*dns.ServeMux).HandleFunc(testDomain+".", func(w dns.ResponseWriter, r *dns.Msg) {
retryDNSHandler(t, w, r, s2, true)
})
return s1, s2
}
// runTestDNSServer starts and returns a new dnsTestServer listening in localhost and the given UDP port
func runTestDNSServer(t *testing.T, port string) *dnsTestServer {
listener, err := net.ListenPacket("udp", "127.0.0.1:"+port)
if err != nil {
t.Fatal(err)
}
mux := dns.NewServeMux()
server := &dns.Server{PacketConn: listener, Net: "udp", Handler: mux}
go func() {
if err := server.ActivateAndServe(); err != nil {
log.Printf("Error in local DNS server: %s", err)
}
}()
return newDNSTestServer(server)
}
// doDNSAnswer sends replies to the DNS question from client, using the dnsDatabase to lookup the answers to the query
// when invertAnswers is true, reverses the order of the answers from the dnsDatabase, useful to simulate realistic nameservers behaviours
func doDNSAnswer(t *testing.T, w dns.ResponseWriter, r *dns.Msg, d dnsDatabase, invertAnswers bool) {
m := new(dns.Msg)
m.SetReply(r)
m.Authoritative = true
q := m.Question[0]
qtype := dns.TypeToString[q.Qtype]
answers := d[DNSQuery{qtype, strings.TrimSuffix(q.Name, ".")}]
var seen = make(map[DNSAnswer]bool)
for _, r := range answers {
if seen[r] {
continue
}
seen[r] = true
rr, err := dns.NewRR(fmt.Sprintf("%s %s", q.Name, r.String()))
if err != nil {
t.Fatalf("err: %s", err)
}
m.Answer = append(m.Answer, rr)
}
if invertAnswers {
for i, j := 0, len(m.Answer)-1; i < j; i, j = i+1, j-1 {
m.Answer[i], m.Answer[j] = m.Answer[j], m.Answer[i]
}
}
w.WriteMsg(m)
}
// stdDNSHandler uses the internal DNSDatabase to send answers to DNS queries
func stdDNSHandler(t *testing.T, w dns.ResponseWriter, r *dns.Msg, s *dnsTestServer, invertAnswers bool) {
doDNSAnswer(t, w, r, s.DNSDatabase, invertAnswers)
}
var startTime = time.Now()
// retryDNSHandler uses the internal DNSDatabase to send answers to DNS queries, and switches
// to using the internal DNSDatabaseRetry after 3 seconds from startup
func retryDNSHandler(t *testing.T, w dns.ResponseWriter, r *dns.Msg, s *dnsTestServer, invertAnswers bool) {
if time.Now().Sub(startTime).Seconds() > 3 {
doDNSAnswer(t, w, r, s.DNSDatabaseRetry, invertAnswers)
} else {
doDNSAnswer(t, w, r, s.DNSDatabase, invertAnswers)
}
}