Bugfix for miekg/dns#748 (#749)

* Bugfix for miekg/dns#748

w.msg was being prematurely cleared prior to use by TsigVerify

* Modified patch after feedback from tmthrgd

Added a disposeBuffer method to the server that's passed a response. This wipes the reference to and frees the buffer used to store the message after TSIG validation has occured, not before. Since the pool is an attribute of the server (and the logic refers to a server UDPSize attribute), it made sense to make this a function of the server rather than a function of the response.

* Added TestServerRoundtripTsig to server_test.go

This test generates a TSIG signed query, and makes sure that server TSIG validation does not produce an error.

* Fixed data races introduced by TestServerRoundtripTsig

* Simplified error signalling in TestServerRoundtripTsig

* RunLocalUDPServerWithFinChan variadic closure argument added

This (clever hack suggested by tmthrgd) allows specifying field values (like TsigSecret) on Server instances at test time without making the race detector grouchy, but is backwards compatible with existing invocations of RunLocalUDPServerWithFinChan.

server.go

@@ -735,20 +735,23 @@ func (srv *Server) serve(w *response) {
 	}
 }
 
-func (srv *Server) serveDNS(w *response) {
-	req := new(Msg)
-	err := req.Unpack(w.msg)
+func (srv *Server) disposeBuffer(w *response) {
 	if w.udp != nil && cap(w.msg) == srv.UDPSize {
 		srv.udpPool.Put(w.msg[:srv.UDPSize])
 	}
 	w.msg = nil
+}
+
+func (srv *Server) serveDNS(w *response) {
+	req := new(Msg)
+	err := req.Unpack(w.msg)
 	if err != nil { // Send a FormatError back
 		x := new(Msg)
 		x.SetRcodeFormatError(req)
 		w.WriteMsg(x)
-		return
 	}
-	if !srv.Unsafe && req.Response {
+	if err != nil || !srv.Unsafe && req.Response {
+		srv.disposeBuffer(w)
 		return
 	}
 
@@ -765,6 +768,8 @@ func (srv *Server) serveDNS(w *response) {
 		}
 	}
 
+	srv.disposeBuffer(w)
+
 	handler := srv.Handler
 	if handler == nil {
 		handler = DefaultServeMux

server_test.go

@@ -59,7 +59,7 @@ func RunLocalUDPServer(laddr string) (*Server, string, error) {
 	return server, l, err
 }
 
-func RunLocalUDPServerWithFinChan(laddr string) (*Server, string, chan error, error) {
+func RunLocalUDPServerWithFinChan(laddr string, opts ...func(*Server)) (*Server, string, chan error, error) {
 	pc, err := net.ListenPacket("udp", laddr)
 	if err != nil {
 		return nil, "", nil, err
@@ -75,6 +75,10 @@ func RunLocalUDPServerWithFinChan(laddr string) (*Server, string, chan error, er
 	// in RunLocalUDPServer and can happen in TestShutdownUDP.
 	fin := make(chan error, 1)
 
+	for _, opt := range opts {
+		opt(server)
+	}
+
 	go func() {
 		fin <- server.ActivateAndServe()
 		pc.Close()
@@ -999,6 +1003,56 @@ func TestServerReuseport(t *testing.T) {
 	}
 }
 
+func TestServerRoundtripTsig(t *testing.T) {
+	secret := map[string]string{"test.": "so6ZGir4GPAqINNh9U5c3A=="}
+
+	s, addrstr, _, err := RunLocalUDPServerWithFinChan(":0", func(srv *Server) {
+		srv.TsigSecret = secret
+	})
+	if err != nil {
+		t.Fatalf("unable to run test server: %v", err)
+	}
+	defer s.Shutdown()
+
+	HandleFunc("example.com.", func(w ResponseWriter, r *Msg) {
+		m := new(Msg)
+		m.SetReply(r)
+		if r.IsTsig() != nil {
+			status := w.TsigStatus()
+			if status == nil {
+				// *Msg r has an TSIG record and it was validated
+				m.SetTsig("test.", HmacMD5, 300, time.Now().Unix())
+			} else {
+				// *Msg r has an TSIG records and it was not valided
+				t.Errorf("invalid TSIG: %v", status)
+			}
+		} else {
+			t.Error("missing TSIG")
+		}
+		w.WriteMsg(m)
+	})
+
+	c := new(Client)
+	m := new(Msg)
+	m.Opcode = OpcodeUpdate
+	m.SetQuestion("example.com.", TypeSOA)
+	m.Ns = []RR{&CNAME{
+		Hdr: RR_Header{
+			Name: "foo.example.com.",
+			Rrtype: TypeCNAME,
+			Class: ClassINET,
+			Ttl: 300,
+		},
+		Target: "bar.example.com.",
+	}}
+	c.TsigSecret = secret
+	m.SetTsig("test.", HmacMD5, 300, time.Now().Unix())
+	_, _, err = c.Exchange(m, addrstr)
+	if err != nil {
+		t.Fatal("failed to exchange", err)
+	}
+}
+
 type ExampleFrameLengthWriter struct {
 	Writer
 }