From f9129942587e2343988f55a4776ea913dff5d7cd Mon Sep 17 00:00:00 2001
From: Miek Gieben <miek@miek.nl>
Date: Tue, 19 Aug 2014 21:46:00 +0100
Subject: [PATCH] Don't relay on deadline to stop the server.

We just send ourselves an UDP or TCP packet so that the listener fails
through and picks up the quit bool send on the channel.
---
 server.go      | 23 +++++++----------------
 server_test.go |  7 ++++---
 udp.go         |  3 +--
 3 files changed, 12 insertions(+), 21 deletions(-)

diff --git a/server.go b/server.go
index db2e7201..3e25126f 100644
--- a/server.go
+++ b/server.go
@@ -223,8 +223,6 @@ type Server struct {
 	WriteTimeout time.Duration
 	// TCP idle timeout for multiple queries, if nil, defaults to 8 * time.Second (RFC 5966).
 	IdleTimeout func() time.Duration
-	// Listener deadline timeout, defaults to 1 * time.Second.
-	Deadline time.Duration
 	// Secret(s) for Tsig map[<zonename>]<base64 secret>.
 	TsigSecret map[string]string
 
@@ -305,12 +303,15 @@ func (srv *Server) ActivateAndServe() error {
 // Shutdown shuts down a server. After a call to Shutdown, ListenAndServe and
 // ActivateAndServe will return.
 func (srv *Server) Shutdown() {
+	c := new(Client)
 	switch srv.Net {
 	case "tcp", "tcp4", "tcp6":
-		srv.stopTCP <- true
+		c.Net = "tcp"
+		go func() { srv.stopTCP <- true }()
+		c.Exchange(new(Msg), srv.Addr)
 	case "udp", "udp4", "udp6":
-		// TODO(miek): does not work for udp
-	//		srv.stopUDP <- true
+		go func() { srv.stopUDP <- true }()
+		c.Exchange(new(Msg), srv.Addr)
 	}
 }
 
@@ -326,12 +327,7 @@ func (srv *Server) serveTCP(l *net.TCPListener) error {
 	if srv.ReadTimeout != 0 {
 		rtimeout = srv.ReadTimeout
 	}
-	deadline := 1 * time.Second
-	if srv.Deadline != 0 {
-		deadline = srv.Deadline
-	}
 	for {
-		l.SetDeadline(time.Now().Add(deadline))
 		rw, e := l.AcceptTCP()
 		select {
 		case <-srv.stopTCP:
@@ -365,12 +361,8 @@ func (srv *Server) serveUDP(l *net.UDPConn) error {
 	if srv.ReadTimeout != 0 {
 		rtimeout = srv.ReadTimeout
 	}
-	deadline := 1 * time.Second
-	if srv.Deadline != 0 {
-		deadline = srv.Deadline
-	}
+	// deadline is not used here
 	for {
-		l.SetDeadline(time.Now().Add(deadline))
 		m, s, e := srv.readUDP(l, rtimeout)
 		select {
 		case <-srv.stopUDP:
@@ -488,7 +480,6 @@ func (srv *Server) readTCP(conn *net.TCPConn, timeout time.Duration) ([]byte, er
 func (srv *Server) readUDP(conn *net.UDPConn, timeout time.Duration) ([]byte, *sessionUDP, error) {
 	conn.SetReadDeadline(time.Now().Add(timeout))
 	m := make([]byte, srv.UDPSize)
-	// TODO(miek): deadline and timeout seem not to be honered
 	n, s, e := readFromSessionUDP(conn, m)
 	if e != nil || n == 0 {
 		if e != nil {
diff --git a/server_test.go b/server_test.go
index c25d7068..5d2700bd 100644
--- a/server_test.go
+++ b/server_test.go
@@ -273,6 +273,8 @@ func TestServingLargeResponses(t *testing.T) {
 	}
 }
 
+// TODO(miek): These tests should actually fail when the server does
+// not shut down.
 func TestShutdownTCP(t *testing.T) {
 	server := Server{Addr: ":8055", Net: "tcp"}
 	go func() {
@@ -285,10 +287,9 @@ func TestShutdownTCP(t *testing.T) {
 	}()
 	time.Sleep(4e8)
 	server.Shutdown()
-	time.Sleep(1 * time.Second)
+	time.Sleep(4e8)
 }
 
-// TODO(miek): does not work for udp
 func TestShutdownUDP(t *testing.T) {
 	server := Server{Addr: ":8054", Net: "udp"}
 	go func() {
@@ -301,5 +302,5 @@ func TestShutdownUDP(t *testing.T) {
 	}()
 	time.Sleep(4e8)
 	server.Shutdown()
-	time.Sleep(1 * time.Second)
+	time.Sleep(4e8)
 }
diff --git a/udp.go b/udp.go
index dad7c684..c5bf33f6 100644
--- a/udp.go
+++ b/udp.go
@@ -49,8 +49,7 @@ func readFromSessionUDP(conn *net.UDPConn, b []byte) (int, *sessionUDP, error) {
 	if err != nil {
 		return n, nil, err
 	}
-	session := &sessionUDP{raddr, oob[:oobn]}
-	return n, session, err
+	return n, &sessionUDP{raddr, oob[:oobn]}, err
 }
 
 // writeToSessionUDP acts just like net.UDPConn.WritetTo(), but uses a *sessionUDP instead of a net.Addr.