From cfa72d08621fdbf5fe8543955d824a4a1f6ed9ee Mon Sep 17 00:00:00 2001
From: Andrew Tunnell-Jones <andrew@tj.id.au>
Date: Thu, 26 Nov 2015 08:04:38 +0000
Subject: [PATCH 1/2] Test closing TCP ResponseWriter in a handler

---
 server_test.go | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/server_test.go b/server_test.go
index b3e0c72b..e19671bd 100644
--- a/server_test.go
+++ b/server_test.go
@@ -2,6 +2,7 @@ package dns
 
 import (
 	"fmt"
+	"io"
 	"net"
 	"runtime"
 	"sync"
@@ -389,6 +390,46 @@ func TestShutdownTCP(t *testing.T) {
 	}
 }
 
+func TestHandlerCloseTCP(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		panic(err)
+	}
+	addr := ln.Addr().String()
+
+	server := &Server{Addr: addr, Net: "tcp", Listener: ln}
+
+	hname := "testhandlerclosetcp."
+	h_triggered := false
+	HandleFunc(hname, func(w ResponseWriter, r *Msg) {
+		h_triggered = true
+		w.Close()
+	})
+	defer HandleRemove(hname)
+
+	go func() {
+		defer server.Shutdown()
+		c := &Client{Net: "tcp"}
+		m := new(Msg).SetQuestion(hname, 1)
+		tries := 0
+	exchange:
+		_, _, err := c.Exchange(m, addr)
+		if err != nil && err != io.EOF {
+			t.Logf("Exchange failed: %s\n", err)
+			if tries == 3 {
+				return
+			}
+			time.Sleep(time.Second / 10)
+			tries += 1
+			goto exchange
+		}
+	}()
+	server.ActivateAndServe()
+	if !h_triggered {
+		t.Fatalf("Handler never called")
+	}
+}
+
 func TestShutdownUDP(t *testing.T) {
 	s, _, err := RunLocalUDPServer("127.0.0.1:0")
 	if err != nil {

From 3062dcb7517f99e3a2d73d8d7379d184269dc717 Mon Sep 17 00:00:00 2001
From: Andrew Tunnell-Jones <andrew@tj.id.au>
Date: Thu, 26 Nov 2015 08:10:55 +0000
Subject: [PATCH 2/2] Check server TCP conn exists before reading

w.tcp could be nil if the conn has been closed by a handler.
---
 server.go | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/server.go b/server.go
index b2888f38..bb0d074a 100644
--- a/server.go
+++ b/server.go
@@ -535,6 +535,9 @@ Redo:
 	h.ServeDNS(w, req) // Writes back to the client
 
 Exit:
+	if w.tcp == nil {
+		return
+	}
 	// TODO(miek): make this number configurable?
 	if q > maxTCPQueries { // close socket after this many queries
 		w.Close()