From 7f2bf8764aadba1e6d801bac04646ccfc52c0f84 Mon Sep 17 00:00:00 2001
From: Tom Thorogood <me+github@tomthorogood.co.uk>
Date: Tue, 25 Jun 2019 01:29:43 +0930
Subject: [PATCH] Set the TC bit more aggressively in Truncate (#989)

* Set the TC bit more aggressively in Truncate

* Update Truncate documentation for TC bit changes
---
 msg_truncate.go      | 15 ++++++++++-----
 msg_truncate_test.go | 12 ++++++------
 2 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/msg_truncate.go b/msg_truncate.go
index 83780af6..89d40757 100644
--- a/msg_truncate.go
+++ b/msg_truncate.go
@@ -8,8 +8,13 @@ package dns
 // record adding as many records as possible without exceeding the
 // requested buffer size.
 //
-// The TC bit will be set if any answer records were excluded from the
-// message. This indicates to that the client should retry over TCP.
+// The TC bit will be set if any records were excluded from the message.
+// This indicates to that the client should retry over TCP.
+//
+// According to RFC 2181, the TC bit should only be set if not all of the
+// "required" RRs can be included in the response. Unfortunately, we have
+// no way of knowing which RRs are required so we set the TC bit if any RR
+// had to be omitted from the response.
 //
 // The appropriate buffer size can be retrieved from the requests OPT
 // record, if present, and is transport specific otherwise. dns.MinMsgSize
@@ -71,9 +76,9 @@ func (dns *Msg) Truncate(size int) {
 		l, numExtra = truncateLoop(dns.Extra, size, l, compression)
 	}
 
-	// According to RFC 2181, the TC bit should only be set if not all
-	// of the answer RRs can be included in the response.
-	dns.Truncated = len(dns.Answer) > numAnswer
+	// See the function documentation for when we set this.
+	dns.Truncated = len(dns.Answer) > numAnswer ||
+		len(dns.Ns) > numNS || len(dns.Extra) > numExtra
 
 	dns.Answer = dns.Answer[:numAnswer]
 	dns.Ns = dns.Ns[:numNS]
diff --git a/msg_truncate_test.go b/msg_truncate_test.go
index 075d74d3..8523a72d 100644
--- a/msg_truncate_test.go
+++ b/msg_truncate_test.go
@@ -40,8 +40,8 @@ func TestRequestTruncateExtra(t *testing.T) {
 	if want, got := MinMsgSize, reply.Len(); want < got {
 		t.Errorf("message length should be bellow %d bytes, got %d bytes", want, got)
 	}
-	if reply.Truncated {
-		t.Errorf("truncated bit should not be set")
+	if !reply.Truncated {
+		t.Errorf("truncated bit should be set")
 	}
 }
 
@@ -64,8 +64,8 @@ func TestRequestTruncateExtraEdns0(t *testing.T) {
 	if want, got := size, reply.Len(); want < got {
 		t.Errorf("message length should be bellow %d bytes, got %d bytes", want, got)
 	}
-	if reply.Truncated {
-		t.Errorf("truncated bit should not be set")
+	if !reply.Truncated {
+		t.Errorf("truncated bit should be set")
 	}
 	opt := reply.Extra[len(reply.Extra)-1]
 	if opt.Header().Rrtype != TypeOPT {
@@ -96,8 +96,8 @@ func TestRequestTruncateExtraRegression(t *testing.T) {
 	if want, got := size, reply.Len(); want < got {
 		t.Errorf("message length should be bellow %d bytes, got %d bytes", want, got)
 	}
-	if reply.Truncated {
-		t.Errorf("truncated bit should not be set")
+	if !reply.Truncated {
+		t.Errorf("truncated bit should be set")
 	}
 	opt := reply.Extra[len(reply.Extra)-1]
 	if opt.Header().Rrtype != TypeOPT {