From 6cac440a8dffe6f42c2e32a919611bc1442cc611 Mon Sep 17 00:00:00 2001
From: Alex Wheeler <wheelear@gmail.com>
Date: Tue, 30 Sep 2014 15:15:36 -0400
Subject: [PATCH] Adding support for enabling the DNS truncate flag for UDP
 queries.

---
 command/agent/config.go | 10 ++++++++++
 command/agent/dns.go    |  4 ++++
 2 files changed, 14 insertions(+)

diff --git a/command/agent/config.go b/command/agent/config.go
index c9a47fa2d6..929ed012bb 100644
--- a/command/agent/config.go
+++ b/command/agent/config.go
@@ -58,6 +58,13 @@ type DNSConfig struct {
 	// only the leader.
 	AllowStale bool `mapstructure:"allow_stale"`
 
+	// EnableTruncate is used to enable setting the truncate
+	// flag for UDP DNS queries.  This allows unmodified
+	// clients to re-query the consul server using TCP
+	// when the total number of records exceeds the number
+	// returned by default for UDP.
+	EnableTruncate bool `mapstructure:"enable_truncate"`
+
 	// MaxStale is used to bound how stale of a result is
 	// accepted for a DNS lookup. This can be used with
 	// AllowStale to limit how old of a value is served up.
@@ -688,6 +695,9 @@ func MergeConfig(a, b *Config) *Config {
 	if b.DNSConfig.AllowStale {
 		result.DNSConfig.AllowStale = true
 	}
+	if b.DNSConfig.EnableTruncate {
+		result.DNSConfig.EnableTruncate = true
+	}
 	if b.DNSConfig.MaxStale != 0 {
 		result.DNSConfig.MaxStale = b.DNSConfig.MaxStale
 	}
diff --git a/command/agent/dns.go b/command/agent/dns.go
index 18e2b928b4..8a1522f75c 100644
--- a/command/agent/dns.go
+++ b/command/agent/dns.go
@@ -471,6 +471,10 @@ RPC:
 	// If the network is not TCP, restrict the number of responses
 	if network != "tcp" && len(out.Nodes) > maxServiceResponses {
 		out.Nodes = out.Nodes[:maxServiceResponses]
+		// Flag that there are more records to return in the UDP response
+		if d.config.EnableTruncate == true {
+			resp.Truncated = true
+		}
 	}
 
 	// Add various responses depending on the request