From 18b5ab35390d49f467a337361525ec37406aa385 Mon Sep 17 00:00:00 2001
From: Dave Tucker <dave@dtucker.co.uk>
Date: Fri, 16 Jan 2015 12:11:34 +0000
Subject: [PATCH] Don't forget to Shutdown the DNS Server!

Fixes #609

Signed-off-by: Dave Tucker <dave@dtucker.co.uk>
---
 command/agent/command.go |  3 +++
 command/agent/dns.go     | 24 ++++++++++++++++++------
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/command/agent/command.go b/command/agent/command.go
index b9a82e19ae..e97d588c93 100644
--- a/command/agent/command.go
+++ b/command/agent/command.go
@@ -571,6 +571,9 @@ func (c *Command) Run(args []string) int {
 	if c.rpcServer != nil {
 		defer c.rpcServer.Shutdown()
 	}
+	if c.dnsServer != nil {
+		defer c.dnsServer.Shutdown()
+	}
 
 	for _, server := range c.httpServers {
 		defer server.Shutdown()
diff --git a/command/agent/dns.go b/command/agent/dns.go
index d4b2506414..d6c2517a8e 100644
--- a/command/agent/dns.go
+++ b/command/agent/dns.go
@@ -33,6 +33,16 @@ type DNSServer struct {
 	logger       *log.Logger
 }
 
+// Shutdown stops the DNS Servers
+func (d *DNSServer) Shutdown() {
+	if err := d.dnsServer.Shutdown(); err != nil {
+		d.logger.Printf("[ERR] dns: error stopping udp server: %v", err)
+	}
+	if err := d.dnsServerTCP.Shutdown(); err != nil {
+		d.logger.Printf("[ERR] dns: error stopping tcp server: %v", err)
+	}
+}
+
 // NewDNSServer starts a new DNS server to provide an agent interface
 func NewDNSServer(agent *Agent, config *DNSConfig, logOutput io.Writer, domain string, bind string, recursors []string) (*DNSServer, error) {
 	// Make sure domain is FQDN
@@ -92,16 +102,18 @@ func NewDNSServer(agent *Agent, config *DNSConfig, logOutput io.Writer, domain s
 	// Async start the DNS Servers, handle a potential error
 	errCh := make(chan error, 1)
 	go func() {
-		err := server.ListenAndServe()
-		srv.logger.Printf("[ERR] dns: error starting udp server: %v", err)
-		errCh <- fmt.Errorf("dns udp setup failed: %v", err)
+		if err := server.ListenAndServe(); err != nil {
+			srv.logger.Printf("[ERR] dns: error starting udp server: %v", err)
+			errCh <- fmt.Errorf("dns udp setup failed: %v", err)
+		}
 	}()
 
 	errChTCP := make(chan error, 1)
 	go func() {
-		err := serverTCP.ListenAndServe()
-		srv.logger.Printf("[ERR] dns: error starting tcp server: %v", err)
-		errChTCP <- fmt.Errorf("dns tcp setup failed: %v", err)
+		if err := serverTCP.ListenAndServe(); err != nil {
+			srv.logger.Printf("[ERR] dns: error starting tcp server: %v", err)
+			errChTCP <- fmt.Errorf("dns tcp setup failed: %v", err)
+		}
 	}()
 
 	// Check the server is running, do a test lookup