From 8516ab7c698165d76f19654a4fe2c19875baf94b Mon Sep 17 00:00:00 2001
From: Tom Wilkie <tom@weave.works>
Date: Fri, 21 Aug 2015 11:07:11 +0000
Subject: [PATCH] Unify logic to tell if DNS response is too big.

---
 nameserver/dns.go                   | 11 +++++++----
 test/295_dns_large_response_test.sh | 11 +++++++++--
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/nameserver/dns.go b/nameserver/dns.go
index 438210e3a9..ccd2a9f368 100644
--- a/nameserver/dns.go
+++ b/nameserver/dns.go
@@ -208,7 +208,7 @@ func (h *handler) handleRecursive(w dns.ResponseWriter, req *dns.Msg) {
 			continue
 		}
 		response.Id = req.Id
-		if response.Len() > h.getMaxResponseSize(req) {
+		if h.responseTooBig(req, response) {
 			response.Compress = true
 		}
 		h.respond(w, response)
@@ -224,13 +224,12 @@ func (h *handler) makeResponse(req *dns.Msg, answers []dns.RR) *dns.Msg {
 	response.RecursionAvailable = true
 	response.Authoritative = true
 	response.Answer = answers
-
-	maxSize := h.getMaxResponseSize(req)
-	if len(answers) <= 1 || maxSize <= 0 || response.Len() <= maxSize {
+	if !h.responseTooBig(req, response) {
 		return response
 	}
 
 	// search for smallest i that is too big
+	maxSize := h.getMaxResponseSize(req)
 	i := sort.Search(len(answers), func(i int) bool {
 		// return true if too big
 		response.Answer = answers[:i+1]
@@ -252,6 +251,10 @@ func (h *handler) makeErrorResponse(req *dns.Msg, code int) *dns.Msg {
 	return response
 }
 
+func (h *handler) responseTooBig(req, response *dns.Msg) bool {
+	return len(response.Answer) > 1 && h.maxResponseSize > 0 && response.Len() > h.getMaxResponseSize(req)
+}
+
 func (h *handler) respond(w dns.ResponseWriter, response *dns.Msg) {
 	h.ns.debugf("response: %+v", response)
 	if err := w.WriteMsg(response); err != nil {
diff --git a/test/295_dns_large_response_test.sh b/test/295_dns_large_response_test.sh
index daf7e82ead..f72dcefbc1 100755
--- a/test/295_dns_large_response_test.sh
+++ b/test/295_dns_large_response_test.sh
@@ -16,8 +16,15 @@ for i in $(seq $N); do
 done
 weave_on $HOST1 dns-add $IPS $CID -h $NAME
 
-assert_dns_record $HOST1 c0 $NAME $IPS
-
 assert_raises "exec_on $HOST1 c0 dig MX $NAME | grep -q 'status: NXDOMAIN'"
 
+check() {
+    assert "exec_on $HOST1 c0 dig +short $@ $NAME A | grep -v ';;' | wc -l" $N
+}
+
+check
+check +tcp
+check      +bufsize=700
+check +tcp +bufsize=700
+
 end_suite