pkg: add dnstest (#1098)

Add a full test server impl in this new package + tests. Move
dnsrecorder into this package as well and finish up the commented out
tests that were left in the old dnsrecorder package.

Update all callers and tests.

plugin/pkg/dnstest/recorder.go

@@ -0,0 +1,54 @@
+// Package dnstest allows for easy testing of DNS client against a test server.
+package dnstest
+
+import (
+	"time"
+
+	"github.com/miekg/dns"
+)
+
+// Recorder is a type of ResponseWriter that captures
+// the rcode code written to it and also the size of the message
+// written in the response. A rcode code does not have
+// to be written, however, in which case 0 must be assumed.
+// It is best to have the constructor initialize this type
+// with that default status code.
+type Recorder struct {
+	dns.ResponseWriter
+	Rcode int
+	Len   int
+	Msg   *dns.Msg
+	Start time.Time
+}
+
+// New makes and returns a new Recorder,
+// which captures the DNS rcode from the ResponseWriter
+// and also the length of the response message written through it.
+func NewRecorder(w dns.ResponseWriter) *Recorder {
+	return &Recorder{
+		ResponseWriter: w,
+		Rcode:          0,
+		Msg:            nil,
+		Start:          time.Now(),
+	}
+}
+
+// WriteMsg records the status code and calls the
+// underlying ResponseWriter's WriteMsg method.
+func (r *Recorder) WriteMsg(res *dns.Msg) error {
+	r.Rcode = res.Rcode
+	// We may get called multiple times (axfr for instance).
+	// Save the last message, but add the sizes.
+	r.Len += res.Len()
+	r.Msg = res
+	return r.ResponseWriter.WriteMsg(res)
+}
+
+// Write is a wrapper that records the length of the message that gets written.
+func (r *Recorder) Write(buf []byte) (int, error) {
+	n, err := r.ResponseWriter.Write(buf)
+	if err == nil {
+		r.Len += n
+	}
+	return n, err
+}

plugin/pkg/dnstest/recorder_test.go

@@ -0,0 +1,50 @@
+package dnstest
+
+import (
+	"testing"
+
+	"github.com/miekg/dns"
+)
+
+type responseWriter struct{ dns.ResponseWriter }
+
+func (r *responseWriter) WriteMsg(m *dns.Msg) error     { return nil }
+func (r *responseWriter) Write(buf []byte) (int, error) { return len(buf), nil }
+
+func TestNewRecorder(t *testing.T) {
+	w := &responseWriter{}
+	record := NewRecorder(w)
+	if record.ResponseWriter != w {
+		t.Fatalf("Expected Response writer in the Recording to be same as the one sent\n")
+	}
+	if record.Rcode != dns.RcodeSuccess {
+		t.Fatalf("Expected recorded status to be dns.RcodeSuccess (%d) , but found %d\n ", dns.RcodeSuccess, record.Rcode)
+	}
+}
+
+func TestWriteMsg(t *testing.T) {
+	w := &responseWriter{}
+	record := NewRecorder(w)
+	responseTestName := "testmsg.example.org."
+	responseTestMsg := new(dns.Msg)
+	responseTestMsg.SetQuestion(responseTestName, dns.TypeA)
+
+	record.WriteMsg(responseTestMsg)
+	if record.Len != responseTestMsg.Len() {
+		t.Fatalf("Expected the bytes written counter to be %d, but instead found %d\n", responseTestMsg.Len(), record.Len)
+	}
+	if x := record.Msg.Question[0].Name; x != responseTestName {
+		t.Fatalf("Expected Msg Qname to be %s , but found %s\n", responseTestName, x)
+	}
+}
+
+func TestWrite(t *testing.T) {
+	w := &responseWriter{}
+	record := NewRecorder(w)
+	responseTest := []byte("testmsg.example.org.")
+
+	record.Write(responseTest)
+	if record.Len != len(responseTest) {
+		t.Fatalf("Expected the bytes written counter to be %d, but instead found %d\n", len(responseTest), record.Len)
+	}
+}

plugin/pkg/dnstest/server.go

@@ -0,0 +1,46 @@
+package dnstest
+
+import (
+	"net"
+
+	"github.com/miekg/dns"
+)
+
+// A Server is an DNS server listening on a system-chosen port on the local
+// loopback interface, for use in end-to-end DNS tests.
+type Server struct {
+	Addr string // Address where the server listening.
+
+	s1 *dns.Server // udp
+	s2 *dns.Server // tcp
+}
+
+// NewServer starts and returns a new Server. The caller should call Close when
+// finished, to shut it down.
+func NewServer(f dns.HandlerFunc) *Server {
+	dns.HandleFunc(".", f)
+
+	ch1 := make(chan bool)
+	ch2 := make(chan bool)
+
+	p, _ := net.ListenPacket("udp", ":0")
+	l, _ := net.Listen("tcp", p.LocalAddr().String())
+
+	s1 := &dns.Server{PacketConn: p}
+	s2 := &dns.Server{Listener: l}
+	s1.NotifyStartedFunc = func() { close(ch1) }
+	s2.NotifyStartedFunc = func() { close(ch2) }
+	go s1.ActivateAndServe()
+	go s2.ActivateAndServe()
+
+	<-ch1
+	<-ch2
+
+	return &Server{s1: s1, s2: s2, Addr: p.LocalAddr().String()}
+}
+
+// Close shuts down the server.
+func (s *Server) Close() {
+	s.s1.Shutdown()
+	s.s2.Shutdown()
+}

plugin/pkg/dnstest/server_test.go

@@ -0,0 +1,28 @@
+package dnstest
+
+import (
+	"testing"
+
+	"github.com/miekg/dns"
+)
+
+func TestNewServer(t *testing.T) {
+	s := NewServer(func(w dns.ResponseWriter, r *dns.Msg) {
+		ret := new(dns.Msg)
+		ret.SetReply(r)
+		w.WriteMsg(ret)
+	})
+	defer s.Close()
+
+	c := new(dns.Client)
+	c.Net = "tcp"
+	m := new(dns.Msg)
+	m.SetQuestion("example.org.", dns.TypeSOA)
+	ret, _, err := c.Exchange(m, s.Addr)
+	if err != nil {
+		t.Fatalf("Could not send message to dnstest.Server: %s", err)
+	}
+	if ret.Id != m.Id {
+		t.Fatalf("Msg ID's should match, expected %d, got %d", m.Id, ret.Id)
+	}
+}