From 540899743c40d65a81d2384d938734e0622e29b7 Mon Sep 17 00:00:00 2001
From: Miek Gieben <miek@miek.nl>
Date: Mon, 31 Aug 2015 17:40:56 +0100
Subject: [PATCH] Handle the last TCP connection

We currently close the connection after 128 TCP queries. But the
when the last query comes in, we close the connection immediately.
Fix this by moving the check to before we read data from the TCP
socket.

Fixes: #218.
---
 server.go | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/server.go b/server.go
index 362c6a72..40f525e3 100644
--- a/server.go
+++ b/server.go
@@ -10,6 +10,9 @@ import (
 	"time"
 )
 
+// Maximum number of TCP queries before we close the socket.
+const maxTCPQueries = 128
+
 // Handler is implemented by any value that implements ServeDNS.
 type Handler interface {
 	ServeDNS(w ResponseWriter, r *Msg)
@@ -502,7 +505,8 @@ func (srv *Server) serve(a net.Addr, h Handler, m []byte, u *net.UDPConn, s *Ses
 		w.writer = w
 	}
 
-	q := 0
+	q := 0 // counter for the amount of TCP queries we get
+
 	defer func() {
 		if u != nil {
 			srv.wgUDP.Done()
@@ -544,6 +548,12 @@ Redo:
 	h.ServeDNS(w, req) // Writes back to the client
 
 Exit:
+	// TODO(miek): make this number configurable?
+	if q > maxTCPQueries { // close socket after this many queries
+		w.Close()
+		return
+	}
+
 	if w.hijacked {
 		return // client calls Close()
 	}
@@ -558,11 +568,6 @@ Exit:
 	m, e := reader.ReadTCP(w.tcp, idleTimeout)
 	if e == nil {
 		q++
-		// TODO(miek): make this number configurable?
-		if q > 128 { // close socket after this many queries
-			w.Close()
-			return
-		}
 		goto Redo
 	}
 	w.Close()