From dcd7d9b015fd208078f8b076e2dfac0447a7116e Mon Sep 17 00:00:00 2001
From: Siva Prasad <MagnumOpus21@users.noreply.github.com>
Date: Thu, 2 Aug 2018 10:12:52 -0400
Subject: [PATCH] DNS : Fixes recursors answering the DNS query to properly
 return the correct response. (#4461)

* Fixes the DNS recursor properly resolving the requests

* Added a test case for the recursor bug

* Refactored code && added a test case for all failing recursors

* Inner indentation moved into else if check
---
 agent/dns.go      | 10 +++++--
 agent/dns_test.go | 69 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/agent/dns.go b/agent/dns.go
index 49028020a..af3a80115 100644
--- a/agent/dns.go
+++ b/agent/dns.go
@@ -1302,14 +1302,20 @@ func (d *DNSServer) handleRecurse(resp dns.ResponseWriter, req *dns.Msg) {
 	var err error
 	for _, recursor := range d.recursors {
 		r, rtt, err = c.Exchange(req, recursor)
-		if err == nil || err == dns.ErrTruncated {
+		// Check if the response is valid and has the desired Response code
+		if r != nil && (r.Rcode != dns.RcodeSuccess && r.Rcode != dns.RcodeNameError) {
+			d.logger.Printf("[DEBUG] dns: recurse RTT for %v (%v) Recursor queried: %v Status returned: %v", q, rtt, recursor, dns.RcodeToString[r.Rcode])
+			// If we still have recursors to forward the query to,
+			// we move forward onto the next one else the loop ends
+			continue
+		} else if err == nil || err == dns.ErrTruncated {
 			// 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.disableCompression.Load().(bool)
 
 			// Forward the response
-			d.logger.Printf("[DEBUG] dns: recurse RTT for %v (%v)", q, rtt)
+			d.logger.Printf("[DEBUG] dns: recurse RTT for %v (%v) Recursor queried: %v", q, rtt, recursor)
 			if err := resp.WriteMsg(r); err != nil {
 				d.logger.Printf("[WARN] dns: failed to respond: %v", err)
 			}
diff --git a/agent/dns_test.go b/agent/dns_test.go
index f2c702739..043fa87b0 100644
--- a/agent/dns_test.go
+++ b/agent/dns_test.go
@@ -37,9 +37,15 @@ const (
 // the provided reply. This is useful for mocking a DNS recursor with
 // an expected result.
 func makeRecursor(t *testing.T, answer dns.Msg) *dns.Server {
+	a := answer
 	mux := dns.NewServeMux()
 	mux.HandleFunc(".", func(resp dns.ResponseWriter, msg *dns.Msg) {
+		// The SetReply function sets the return code of the DNS
+		// query to SUCCESS
+		// We need a way to copy the variables not addressed
+		// in SetReply
 		answer.SetReply(msg)
+		answer.Rcode = a.Rcode
 		if err := resp.WriteMsg(&answer); err != nil {
 			t.Fatalf("err: %s", err)
 		}
@@ -371,6 +377,69 @@ func TestDNS_NodeLookup_AAAA(t *testing.T) {
 	}
 }
 
+func TestDNSCycleRecursorCheck(t *testing.T) {
+	t.Parallel()
+	// Start a DNS recursor that returns a SERVFAIL
+	server1 := makeRecursor(t, dns.Msg{
+		MsgHdr: dns.MsgHdr{Rcode: dns.RcodeServerFailure},
+	})
+	// Start a DNS recursor that returns the result
+	defer server1.Shutdown()
+	server2 := makeRecursor(t, dns.Msg{
+		MsgHdr: dns.MsgHdr{Rcode: dns.RcodeSuccess},
+		Answer: []dns.RR{
+			dnsA("www.google.com", "172.21.45.67"),
+		},
+	})
+	defer server2.Shutdown()
+	//Mock the agent startup with the necessary configs
+	agent := NewTestAgent(t.Name(),
+		`recursors = ["`+server1.Addr+`", "`+server2.Addr+`"]
+		`)
+	defer agent.Shutdown()
+	// DNS Message init
+	m := new(dns.Msg)
+	m.SetQuestion("google.com.", dns.TypeA)
+	// Agent request
+	client := new(dns.Client)
+	in, _, _ := client.Exchange(m, agent.DNSAddr())
+	wantAnswer := []dns.RR{
+		&dns.A{
+			Hdr: dns.RR_Header{Name: "www.google.com.", Rrtype: dns.TypeA, Class: dns.ClassINET, Rdlength: 0x4},
+			A:   []byte{0xAC, 0x15, 0x2D, 0x43}, // 172 , 21, 45, 67
+		},
+	}
+	verify.Values(t, "Answer", in.Answer, wantAnswer)
+}
+func TestDNSCycleRecursorCheckAllFail(t *testing.T) {
+	t.Parallel()
+	// Start 3 DNS recursors that returns a REFUSED status
+	server1 := makeRecursor(t, dns.Msg{
+		MsgHdr: dns.MsgHdr{Rcode: dns.RcodeRefused},
+	})
+	defer server1.Shutdown()
+	server2 := makeRecursor(t, dns.Msg{
+		MsgHdr: dns.MsgHdr{Rcode: dns.RcodeRefused},
+	})
+	defer server2.Shutdown()
+	server3 := makeRecursor(t, dns.Msg{
+		MsgHdr: dns.MsgHdr{Rcode: dns.RcodeRefused},
+	})
+	defer server3.Shutdown()
+	//Mock the agent startup with the necessary configs
+	agent := NewTestAgent(t.Name(),
+		`recursors = ["`+server1.Addr+`", "`+server2.Addr+`","`+server3.Addr+`"]
+		`)
+	defer agent.Shutdown()
+	// DNS dummy message initialization
+	m := new(dns.Msg)
+	m.SetQuestion("google.com.", dns.TypeA)
+	// Agent request
+	client := new(dns.Client)
+	in, _, _ := client.Exchange(m, agent.DNSAddr())
+	//Verify if we hit SERVFAIL from Consul
+	verify.Values(t, "Answer", in.Rcode, dns.RcodeServerFailure)
+}
 func TestDNS_NodeLookup_CNAME(t *testing.T) {
 	t.Parallel()
 	recursor := makeRecursor(t, dns.Msg{