From 727b6444ad535e85e60e43cd709f8a87861c9b59 Mon Sep 17 00:00:00 2001
From: Frank Schroeder <frank.schroeder@gmail.com>
Date: Thu, 29 Jun 2017 16:42:17 +0200
Subject: [PATCH] dns: fix data races in DNS compression tests

Make the DisableCompression value configurable at runtime
to allow tests to change it without restarting/recreating
the server.
---
 agent/dns.go       | 15 +++++++++++----
 agent/dns_test.go  |  8 ++++----
 agent/testagent.go |  7 +++++++
 3 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/agent/dns.go b/agent/dns.go
index e8bf650a4a..8f95ea137c 100644
--- a/agent/dns.go
+++ b/agent/dns.go
@@ -6,6 +6,7 @@ import (
 	"log"
 	"net"
 	"strings"
+	"sync/atomic"
 	"time"
 
 	"github.com/armon/go-metrics"
@@ -38,6 +39,11 @@ type DNSServer struct {
 	domain    string
 	recursors []string
 	logger    *log.Logger
+
+	// disableCompression is the config.DisableCompression flag that can
+	// be safely changed at runtime. It always contains a bool and is
+	// initialized with the value from config.DisableCompression.
+	disableCompression atomic.Value
 }
 
 func NewDNSServer(a *Agent) (*DNSServer, error) {
@@ -60,6 +66,7 @@ func NewDNSServer(a *Agent) (*DNSServer, error) {
 		logger:    a.logger,
 		recursors: recursors,
 	}
+	srv.disableCompression.Store(a.config.DNSConfig.DisableCompression)
 
 	return srv, nil
 }
@@ -120,7 +127,7 @@ func (d *DNSServer) handlePtr(resp dns.ResponseWriter, req *dns.Msg) {
 	// Setup the message response
 	m := new(dns.Msg)
 	m.SetReply(req)
-	m.Compress = !d.config.DisableCompression
+	m.Compress = !d.disableCompression.Load().(bool)
 	m.Authoritative = true
 	m.RecursionAvailable = (len(d.recursors) > 0)
 
@@ -195,7 +202,7 @@ func (d *DNSServer) handleQuery(resp dns.ResponseWriter, req *dns.Msg) {
 	// Setup the message response
 	m := new(dns.Msg)
 	m.SetReply(req)
-	m.Compress = !d.config.DisableCompression
+	m.Compress = !d.disableCompression.Load().(bool)
 	m.Authoritative = true
 	m.RecursionAvailable = (len(d.recursors) > 0)
 
@@ -907,7 +914,7 @@ func (d *DNSServer) handleRecurse(resp dns.ResponseWriter, req *dns.Msg) {
 			// Compress the response; we don't know if the incoming
 			// response was compressed or not, so by not compressing
 			// we might generate an invalid packet on the way out.
-			r.Compress = !d.config.DisableCompression
+			r.Compress = !d.disableCompression.Load().(bool)
 
 			// Forward the response
 			d.logger.Printf("[DEBUG] dns: recurse RTT for %v (%v)", q, rtt)
@@ -924,7 +931,7 @@ func (d *DNSServer) handleRecurse(resp dns.ResponseWriter, req *dns.Msg) {
 		q, resp.RemoteAddr().String(), resp.RemoteAddr().Network())
 	m := &dns.Msg{}
 	m.SetReply(req)
-	m.Compress = !d.config.DisableCompression
+	m.Compress = !d.disableCompression.Load().(bool)
 	m.RecursionAvailable = true
 	m.SetRcode(req, dns.RcodeServerFailure)
 	if edns := req.IsEdns0(); edns != nil {
diff --git a/agent/dns_test.go b/agent/dns_test.go
index aa5e45a391..38cf54c9da 100644
--- a/agent/dns_test.go
+++ b/agent/dns_test.go
@@ -4558,7 +4558,7 @@ func TestDNS_Compression_Query(t *testing.T) {
 		}
 
 		// Do a manual exchange with compression on (the default).
-		a.dns.config.DisableCompression = false
+		a.DNSDisableCompression(false)
 		if err := conn.WriteMsg(m); err != nil {
 			t.Fatalf("err: %v", err)
 		}
@@ -4569,7 +4569,7 @@ func TestDNS_Compression_Query(t *testing.T) {
 		}
 
 		// Disable compression and try again.
-		a.dns.config.DisableCompression = true
+		a.DNSDisableCompression(true)
 		if err := conn.WriteMsg(m); err != nil {
 			t.Fatalf("err: %v", err)
 		}
@@ -4623,7 +4623,7 @@ func TestDNS_Compression_ReverseLookup(t *testing.T) {
 	}
 
 	// Disable compression and try again.
-	a.dns.config.DisableCompression = true
+	a.DNSDisableCompression(true)
 	if err := conn.WriteMsg(m); err != nil {
 		t.Fatalf("err: %v", err)
 	}
@@ -4671,7 +4671,7 @@ func TestDNS_Compression_Recurse(t *testing.T) {
 	}
 
 	// Disable compression and try again.
-	a.dns.config.DisableCompression = true
+	a.DNSDisableCompression(true)
 	if err := conn.WriteMsg(m); err != nil {
 		t.Fatalf("err: %v", err)
 	}
diff --git a/agent/testagent.go b/agent/testagent.go
index e6a7272320..cf5d98c4db 100644
--- a/agent/testagent.go
+++ b/agent/testagent.go
@@ -245,6 +245,13 @@ func (a *TestAgent) Client() *api.Client {
 	return c
 }
 
+// DNSDisableCompression disables compression for all started DNS servers.
+func (a *TestAgent) DNSDisableCompression(b bool) {
+	for _, srv := range a.dnsServers {
+		srv.disableCompression.Store(b)
+	}
+}
+
 func (a *TestAgent) consulConfig() *consul.Config {
 	c, err := a.Agent.consulConfig()
 	if err != nil {