From 171b43e198e41971829e92b220002de0ec2b9995 Mon Sep 17 00:00:00 2001
From: Seth Vargo <seth@sethvargo.com>
Date: Tue, 13 Jun 2017 08:15:00 +0100
Subject: [PATCH] Add EDNS0 test for trimming

---
 agent/dns.go      | 11 ++++---
 agent/dns_test.go | 79 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 5 deletions(-)

diff --git a/agent/dns.go b/agent/dns.go
index c150b15d093b..e9dbdb878cc6 100644
--- a/agent/dns.go
+++ b/agent/dns.go
@@ -567,11 +567,12 @@ func trimUDPResponse(config *DNSConfig, req, resp *dns.Msg) (trimmed bool) {
 		}
 	}
 
-	// This enforces the hard limit of 512 bytes per the RFC. Note that we
-	// temporarily switch to uncompressed so that we limit to a response
-	// that will not exceed 512 bytes uncompressed, which is more
-	// conservative and will allow our responses to be compliant even if
-	// some downstream server uncompresses them.
+	// This enforces the given limit on the number bytes. The default is 512 as
+	// per the RFC, but EDNS0 allows for the user to specify larger sizes. Note
+	// that we temporarily switch to uncompressed so that we limit to a response
+	// that will not exceed 512 bytes uncompressed, which is more conservative and
+	// will allow our responses to be compliant even if some downstream server
+	// uncompresses them.
 	compress := resp.Compress
 	resp.Compress = false
 	for len(resp.Answer) > 0 && resp.Len() > maxSize {
diff --git a/agent/dns_test.go b/agent/dns_test.go
index f1ba37049e4e..fd69dcb5d6d2 100644
--- a/agent/dns_test.go
+++ b/agent/dns_test.go
@@ -4190,6 +4190,85 @@ func TestDNS_trimUDPResponse_TrimSize(t *testing.T) {
 	}
 }
 
+func TestDNS_trimUDPResponse_TrimSizeEDNS(t *testing.T) {
+	t.Parallel()
+	config := &DefaultConfig().DNSConfig
+
+	req, resp := &dns.Msg{}, &dns.Msg{}
+
+	for i := 0; i < 100; i++ {
+		target := fmt.Sprintf("ip-10-0-1-%d.node.dc1.consul.", 150+i)
+		srv := &dns.SRV{
+			Hdr: dns.RR_Header{
+				Name:   "redis-cache-redis.service.consul.",
+				Rrtype: dns.TypeSRV,
+				Class:  dns.ClassINET,
+			},
+			Target: target,
+		}
+		a := &dns.A{
+			Hdr: dns.RR_Header{
+				Name:   target,
+				Rrtype: dns.TypeA,
+				Class:  dns.ClassINET,
+			},
+			A: net.ParseIP(fmt.Sprintf("10.0.1.%d", 150+i)),
+		}
+
+		resp.Answer = append(resp.Answer, srv)
+		resp.Extra = append(resp.Extra, a)
+	}
+
+	// Copy over to a new slice since we are trimming both.
+	reqEDNS, respEDNS := &dns.Msg{}, &dns.Msg{}
+	reqEDNS.SetEdns0(2048, true)
+	respEDNS.Answer = append(respEDNS.Answer, resp.Answer...)
+	respEDNS.Extra = append(respEDNS.Extra, resp.Extra...)
+
+	// Trim each response
+	if trimmed := trimUDPResponse(config, req, resp); !trimmed {
+		t.Errorf("expected response to be trimmed: %#v", resp)
+	}
+	if trimmed := trimUDPResponse(config, reqEDNS, respEDNS); !trimmed {
+		t.Errorf("expected edns to be trimmed: %#v", resp)
+	}
+
+	// Check answer lengths
+	if len(resp.Answer) == 0 || len(resp.Answer) != len(resp.Extra) {
+		t.Errorf("bad response answer length: %#v", resp)
+	}
+	if len(respEDNS.Answer) == 0 || len(respEDNS.Answer) != len(respEDNS.Extra) {
+		t.Errorf("bad edns answer length: %#v", resp)
+	}
+
+	// Due to the compression, we can't check exact equality of sizes, but we can
+	// make two requests and ensure that the edns one returns a larger payload
+	// than the non-edns0 one.
+	if len(resp.Answer) >= len(respEDNS.Answer) {
+		t.Errorf("expected edns have larger answer: %#v\n%#v", resp, respEDNS)
+	}
+	if len(resp.Extra) >= len(respEDNS.Extra) {
+		t.Errorf("expected edns have larger extra: %#v\n%#v", resp, respEDNS)
+	}
+
+	// Verify that the things point where they should
+	for i := range resp.Answer {
+		srv, ok := resp.Answer[i].(*dns.SRV)
+		if !ok {
+			t.Errorf("%d should be an SRV", i)
+		}
+
+		a, ok := resp.Extra[i].(*dns.A)
+		if !ok {
+			t.Errorf("%d should be an A", i)
+		}
+
+		if srv.Target != a.Header().Name {
+			t.Errorf("%d: bad %#v vs. %#v", i, srv, a)
+		}
+	}
+}
+
 func TestDNS_syncExtra(t *testing.T) {
 	t.Parallel()
 	resp := &dns.Msg{