func TestSimple(t *testing.T) {
conn, err := dns.DialTimeout("udp", dnsToGrpcAddr, 1*time.Second)
if err != nil {
t.Fatalf("dns.Dial error: %v", err)
}
defer conn.Close()
err = dnsQuery(conn)
if err != nil {
t.Errorf("dns query returned error: %v", err)
}
}
func manyDNSQueries(b *testing.B, addr string) {
conn, err := dns.DialTimeout("udp", addr, 1*time.Second)
if err != nil {
b.Fatalf("dns.Dial error: %v", err)
}
defer conn.Close()
for i := 0; i < b.N; i++ {
err = dnsQuery(conn)
if err != nil {
b.Errorf("dns query returned error: %v", err)
}
}
}
func pollServer() {
backoff := time.Duration(200 * time.Millisecond)
ctx, _ := context.WithTimeout(context.Background(), time.Duration(5*time.Second))
ticker := time.NewTicker(backoff)
for {
select {
case <-ctx.Done():
fmt.Fprintln(os.Stderr, "Timeout reached while testing for the dns server to come up")
os.Exit(1)
case <-ticker.C:
conn, _ := dns.DialTimeout("tcp", dnsLoopbackAddr, backoff)
if conn != nil {
_ = conn.Close()
return
}
}
}
}
// WaitForDNSServer waits 5 seconds for a DNS server to start, and returns an
// error if it fails to do so.
// It does this by repeatedly querying the DNS server until it either replies
// or times out. Note we do not do any validation of the reply.
func WaitForDNSServer(addr string) error {
conn, err := dns.DialTimeout("udp", addr, 1*time.Second)
if err != nil {
return fmt.Errorf("dns.Dial error: %v", err)
}
defer conn.Close()
m := &dns.Msg{}
m.SetQuestion("unused.", dns.TypeA)
deadline := time.Now().Add(5 * time.Second)
tick := time.Tick(100 * time.Millisecond)
for (<-tick).Before(deadline) {
conn.SetDeadline(time.Now().Add(1 * time.Second))
conn.WriteMsg(m)
_, err := conn.ReadMsg()
if err == nil {
return nil
}
}
return fmt.Errorf("timed out")
}
func (this ClientProxy) ServeDNS(w dns.ResponseWriter, request *dns.Msg) {
// if we don't have EDNS0 in the packet, add it now
// TODO: in principle we should check packet size here, since we have made it bigger,
// but for this demo code we will just rely on most queries being really small
proxy_req := *request
opt := proxy_req.IsEdns0()
var client_buf_size uint16
if opt == nil {
proxy_req.SetEdns0(512, false)
client_buf_size = 512
_D("%s QID:%d adding EDNS0 to packet", w.RemoteAddr(), request.Id)
opt = proxy_req.IsEdns0()
} else {
client_buf_size = opt.UDPSize()
}
// add our custom EDNS0 option
local_opt := new(dns.EDNS0_LOCAL)
local_opt.Code = dns.EDNS0LOCALSTART
opt.Option = append(opt.Option, local_opt)
// create a connection to the server
// XXX: for now we will only handle UDP - this will break in unpredictable ways in production!
conn, err := dns.DialTimeout("udp", this.SERVERS[rand.Intn(len(this.SERVERS))], this.timeout)
if err != nil {
_D("%s QID:%d error setting up UDP socket: %s", w.RemoteAddr(), request.Id, err)
SRVFAIL(w, request)
return
}
defer conn.Close()
// set our timeouts
// TODO: we need to insure that our timeouts work like we expect
conn.SetReadDeadline(time.Now().Add(this.timeout))
conn.SetWriteDeadline(time.Now().Add(this.timeout))
// send our query
err = conn.WriteMsg(&proxy_req)
if err != nil {
_D("%s QID:%d error writing message: %s", w.RemoteAddr(), request.Id, err)
SRVFAIL(w, request)
return
}
// wait for our reply
response := wait_for_response(w, conn, request)
if response == nil {
return
}
// get fragment information from first response (if any)
num_frags, sequence_num := get_fragment_info(response)
// if we did not have a fragmented response, send it to the client
if num_frags == -1 {
w.WriteMsg(response)
return
}
// build a map to hold the fragments that we have received
frags := map[int]dns.Msg{sequence_num: *response}
// wait for all fragments to arrive
// duplicates overwrite previous packet, missing packets eventually timeout
for len(frags) < num_frags {
response := wait_for_response(w, conn, request)
if response == nil {
return
}
_, sequence_num := get_fragment_info(response)
// TODO: remove the extra EDNS0 option
frags[sequence_num] = *response
}
// rebuild our original packet
rebuilt_reply, ok := frags[0]
if !ok {
_D("%s QID:%d missing fragment 0", w.RemoteAddr(), request.Id)
SRVFAIL(w, request)
return
}
for n := 1; n < num_frags; n++ {
frag, ok := frags[n]
if !ok {
_D("%s QID:%d missing fragment %d", w.RemoteAddr(), request.Id, n)
SRVFAIL(w, request)
return
}
rebuilt_reply.Answer = append(rebuilt_reply.Answer, frag.Answer...)
rebuilt_reply.Ns = append(rebuilt_reply.Ns, frag.Ns...)
for _, r := range frag.Extra {
// remove EDNS0 present in fragments from final answer
if r.Header().Rrtype != dns.TypeOPT {
rebuilt_reply.Extra = append(rebuilt_reply.Extra, r)
}
}
}
// verify that we don't exceed the client buffer size
if rebuilt_reply.Len() > int(client_buf_size) {
// truncate if we need to
// TODO: test this
rebuilt_reply.MsgHdr.Truncated = true
rebuilt_reply.Answer = []dns.RR{}
rebuilt_reply.Ns = []dns.RR{}
rebuilt_reply.Extra = []dns.RR{}
}
// send our rebuilt reply
w.WriteMsg(&rebuilt_reply)
}