From 36ffedf7d0ba5efb156b6d252a9cbfac63216fbc Mon Sep 17 00:00:00 2001
From: Tom Thorogood <me+github@tomthorogood.co.uk>
Date: Sat, 6 Oct 2018 01:46:28 +0930
Subject: [PATCH] Avoid overriding aLongTimeAgo read deadline (#774)

* Avoid overriding aLongTimeAgo read deadline

The (*Server).readTCP and (*Server).readUDP methods both call
SetReadDeadline. If they race with ShutdownContext, they can override
the aLongTimeAgo read deadline that is set in ShutdownContext. That
will cause ShutdownContext to hang or timeout.

* Reword readTCP comment
---
 server.go | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/server.go b/server.go
index 6bd202bd..4b19b437 100644
--- a/server.go
+++ b/server.go
@@ -666,7 +666,14 @@ func (srv *Server) serveDNS(w *response) {
 }
 
 func (srv *Server) readTCP(conn net.Conn, timeout time.Duration) ([]byte, error) {
-	conn.SetReadDeadline(time.Now().Add(timeout))
+	if srv.isStarted() {
+		// If we race with ShutdownContext, the read deadline may
+		// have been set in the distant past to unblock the read
+		// below. We must not override it, otherwise we may block
+		// ShutdownContext.
+		conn.SetReadDeadline(time.Now().Add(timeout))
+	}
+
 	l := make([]byte, 2)
 	n, err := conn.Read(l)
 	if err != nil || n != 2 {
@@ -701,7 +708,11 @@ func (srv *Server) readTCP(conn net.Conn, timeout time.Duration) ([]byte, error)
 }
 
 func (srv *Server) readUDP(conn *net.UDPConn, timeout time.Duration) ([]byte, *SessionUDP, error) {
-	conn.SetReadDeadline(time.Now().Add(timeout))
+	if srv.isStarted() {
+		// See the comment in readTCP above.
+		conn.SetReadDeadline(time.Now().Add(timeout))
+	}
+
 	m := srv.udpPool.Get().([]byte)
 	n, s, err := ReadFromSessionUDP(conn, m)
 	if err != nil {