Close dependencies implicitly

4 files changed, 10 insertions, 8 deletions

diff --git a/cmd/zdns/main.go b/cmd/zdns/main.go
index 3d7253a..e5c4019 100644
--- a/cmd/zdns/main.go
+++ b/cmd/zdns/main.go
@@ -86,9 +86,11 @@ func (c *cli) run() {
 
 	// Signal handling
 	sigHandler := signal.NewHandler(c.signal, logger)
+	sigHandler.OnClose(logger)
 
 	// Cache
 	cache := cache.New(config.DNS.CacheSize)
+	sigHandler.OnClose(cache)
 
 	// DNS server
 	proxy, err := dns.NewProxy(cache, logger, dns.ProxyOptions{
@@ -98,6 +100,7 @@ func (c *cli) run() {
 		Timeout:   config.Resolver.Timeout,
 	})
 	fatal(err)
+	sigHandler.OnClose(proxy)
 
 	dnsSrv, err := zdns.NewServer(logger, proxy, config)
 	fatal(err)
diff --git a/dns/proxy.go b/dns/proxy.go
index d482e4c..88dbb22 100644
--- a/dns/proxy.go
+++ b/dns/proxy.go
@@ -145,11 +145,9 @@ func (p *Proxy) reply(r *dns.Msg) *dns.Msg {
 // Close closes the proxy.
 func (p *Proxy) Close() error {
 	if p.server != nil {
-		if err := p.server.Shutdown(); err != nil {
-			return err
-		}
+		return p.server.Shutdown()
 	}
-	return p.cache.Close()
+	return nil
 }
 
 func answers(msg *dns.Msg) []string {
diff --git a/dns/proxy_test.go b/dns/proxy_test.go
index 003430d..85d41c9 100644
--- a/dns/proxy_test.go
+++ b/dns/proxy_test.go
@@ -144,6 +144,7 @@ func TestProxy(t *testing.T) {
 	}
 	p := testProxy(t)
 	p.Handler = h
+	defer p.Close()
 
 	m := dns.Msg{}
 	m.Id = dns.Id()
@@ -161,6 +162,7 @@ func TestProxyWithResolvers(t *testing.T) {
 	p.resolvers = []string{"resolver1"}
 	client := make(testClient)
 	p.client = client
+	defer p.Close()
 
 	// First and only resolver responds succesfully
 	reply := ReplyA("host1", net.ParseIP("192.0.2.1"))
@@ -196,6 +198,7 @@ func TestProxyWithCache(t *testing.T) {
 	p.resolvers = []string{"resolver1"}
 	client := make(testClient)
 	p.client = client
+	defer p.Close()
 
 	reply := ReplyA("host1", net.ParseIP("192.0.2.1"))
 	m := dns.Msg{}
@@ -218,6 +221,7 @@ func TestProxyWithLogging(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	defer p.Close()
 	badHost := "badhost1."
 	goodHost := "goodhost1."
 	p.resolvers = []string{"resolver1"}
diff --git a/server.go b/server.go
index ff265b4..9ec4977 100644
--- a/server.go
+++ b/server.go
@@ -171,10 +171,7 @@ func (s *Server) Close() error {
 	if s.ticker != nil {
 		s.done <- true
 	}
-	if s.proxy != nil {
-		return s.proxy.Close()
-	}
-	return s.logger.Close()
+	return nil
 }
 
 func (s *Server) hijack(r *dns.Request) *dns.Reply {