From 3a258602700c361d956f75a95461a0179600573f Mon Sep 17 00:00:00 2001
From: Igor Dubinskiy <akkifokkusu@gmail.com>
Date: Wed, 17 Feb 2016 16:54:28 -0800
Subject: [PATCH] Make sure UDP DNS responses aren't larger than allowed

---
 command/agent/dns.go                          | 44 +++++++++++----
 command/agent/dns_test.go                     | 56 +++++++++++++++++++
 .../source/docs/agent/options.html.markdown   |  5 +-
 3 files changed, 91 insertions(+), 14 deletions(-)

diff --git a/command/agent/dns.go b/command/agent/dns.go
index a50486173e..c0ee3f773d 100644
--- a/command/agent/dns.go
+++ b/command/agent/dns.go
@@ -487,6 +487,27 @@ func (d *DNSServer) formatNodeRecord(node *structs.Node, addr, qName string, qTy
 	return records
 }
 
+// trimAnswers makes sure a UDP response is not longer than allowed by RFC 1035
+func trimAnswers(resp *dns.Msg) (trimmed bool) {
+	numAnswers := len(resp.Answer)
+
+	if numAnswers > maxServiceResponses {
+		resp.Answer = resp.Answer[:maxServiceResponses]
+	}
+
+	// Check that the response isn't more than 512 bytes
+	for respBytes := resp.Len(); respBytes > 512; respBytes = resp.Len() {
+		resp.Answer = resp.Answer[:len(resp.Answer)-1]
+
+		if len(resp.Answer) == 0 {
+			// We've done all we can
+			break
+		}
+	}
+
+	return len(resp.Answer) < numAnswers
+}
+
 // serviceLookup is used to handle a service query
 func (d *DNSServer) serviceLookup(network, datacenter, service, tag string, req, resp *dns.Msg) {
 	// Make an RPC request
@@ -547,17 +568,17 @@ RPC:
 	}
 
 	// If the network is not TCP, restrict the number of responses
-	if network != "tcp" && len(resp.Answer) > maxServiceResponses {
-		resp.Answer = resp.Answer[:maxServiceResponses]
+	if network != "tcp" {
+		wasTrimmed := trimAnswers(resp)
 
 		// Flag that there are more records to return in the UDP response
-		if d.config.EnableTruncate {
+		if wasTrimmed && d.config.EnableTruncate {
 			resp.Truncated = true
 		}
 	}
 
-	// If the answer is empty, return not found
-	if len(resp.Answer) == 0 {
+	// If the answer is empty and the response isn't truncated, return not found
+	if len(resp.Answer) == 0 && !resp.Truncated {
 		d.addSOA(d.domain, resp)
 		return
 	}
@@ -642,18 +663,17 @@ RPC:
 	}
 
 	// If the network is not TCP, restrict the number of responses.
-	if network != "tcp" && len(resp.Answer) > maxServiceResponses {
-		resp.Answer = resp.Answer[:maxServiceResponses]
+	if network != "tcp" {
+		wasTrimmed := trimAnswers(resp)
 
-		// Flag that there are more records to return in the UDP
-		// response.
-		if d.config.EnableTruncate {
+		// Flag that there are more records to return in the UDP response
+		if wasTrimmed && d.config.EnableTruncate {
 			resp.Truncated = true
 		}
 	}
 
-	// If the answer is empty, return not found.
-	if len(resp.Answer) == 0 {
+	// If the answer is empty and the response isn't truncated, return not found
+	if len(resp.Answer) == 0 && !resp.Truncated {
 		d.addSOA(d.domain, resp)
 		return
 	}
diff --git a/command/agent/dns_test.go b/command/agent/dns_test.go
index 32f4c1f14d..afbd6ebbc5 100644
--- a/command/agent/dns_test.go
+++ b/command/agent/dns_test.go
@@ -1961,6 +1961,62 @@ func TestDNS_ServiceLookup_Truncate(t *testing.T) {
 	}
 }
 
+func TestDNS_ServiceLookup_LargeResponses(t *testing.T) {
+	dir, srv := makeDNSServerConfig(t, nil, func(c *DNSConfig) {
+		c.EnableTruncate = true
+	})
+	defer os.RemoveAll(dir)
+	defer srv.agent.Shutdown()
+
+	testutil.WaitForLeader(t, srv.agent.RPC, "dc1")
+
+	longServiceName := "this-is-a-very-very-very-very-very-long-name-for-a-service"
+
+	// Register 3 nodes
+	for i := 0; i < 3; i++ {
+		args := &structs.RegisterRequest{
+			Datacenter: "dc1",
+			Node:       fmt.Sprintf("foo%d", i),
+			Address:    fmt.Sprintf("127.0.0.%d", i+1),
+			Service: &structs.NodeService{
+				Service: longServiceName,
+				Tags:    []string{"master"},
+				Port:    12345,
+			},
+		}
+
+		var out struct{}
+		if err := srv.agent.RPC("Catalog.Register", args, &out); err != nil {
+			t.Fatalf("err: %v", err)
+		}
+	}
+
+	m := new(dns.Msg)
+	m.SetQuestion("_"+longServiceName+"._master.service.consul.", dns.TypeSRV)
+
+	c := new(dns.Client)
+	addr, _ := srv.agent.config.ClientListener("", srv.agent.config.Ports.DNS)
+	in, _, err := c.Exchange(m, addr.String())
+	if err != nil && err != dns.ErrTruncated {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Make sure the response size is RFC 1035-compliant for UDP messages
+	if in.Len() > 512 {
+		t.Fatalf("Bad: %#v", in.Len())
+	}
+
+	// We should only have two answers now
+	if len(in.Answer) != 2 {
+		t.Fatalf("Bad: %#v", len(in.Answer))
+	}
+
+	// Check for the truncate bit
+	if !in.Truncated {
+		t.Fatalf("should have truncate bit")
+	}
+}
+
 func TestDNS_ServiceLookup_MaxResponses(t *testing.T) {
 	dir, srv := makeDNSServer(t)
 	defer os.RemoveAll(dir)
diff --git a/website/source/docs/agent/options.html.markdown b/website/source/docs/agent/options.html.markdown
index e227028eae..23c29c6ac8 100644
--- a/website/source/docs/agent/options.html.markdown
+++ b/website/source/docs/agent/options.html.markdown
@@ -493,8 +493,9 @@ Consul will not enable TLS for the HTTP API unless the `https` port has been ass
   setting this value.
 
   * <a name="enable_truncate"></a><a href="#enable_truncate">`enable_truncate`</a> If set to
-  true, a UDP DNS query that would return more than 3 records will set the truncated flag,
-  indicating to clients that they should re-query using TCP to get the full set of records.
+  true, a UDP DNS query that would return more 3 records, or more than would fit into a valid
+  UDP response, will set the truncated flag, indicating to clients that they should re-query
+  using TCP to get the full set of records.
 
   * <a name="only_passing"></a><a href="#only_passing">`only_passing`</a> If set to true, any
   nodes whose healthchecks are not passing will be excluded from DNS results. By default (or