request.Scrub: test for rl==size case (#1631)

* request.Scrub: test for rl==size case

Make a test case for the new break statement in Scrub and also
account for the OPT record that may get re-added in SizeAndDo() -
otherwise we may break clients that expect this.

* Fix comment

2 files changed, 63 insertions, 8 deletions

diff --git a/request/request.go b/request/request.go
index 9c9512f3a..0629a2123 100644
--- a/request/request.go
+++ b/request/request.go
@@ -167,7 +167,6 @@ func (r *Request) SizeAndDo(m *dns.Msg) bool {
 	if odo {
 		o.SetDo()
 	}
-
 	m.Extra = append(m.Extra, o)
 	return true
 }
@@ -202,8 +201,13 @@ func (r *Request) Scrub(reply *dns.Msg) (*dns.Msg, Result) {
 		return reply, ScrubIgnored
 	}
 
+	// Account for the OPT record that gets added in SizeAndDo(), subtract that length.
+	sub := 0
+	if r.Do() {
+		sub = optLen
+	}
 	origExtra := reply.Extra
-	re := len(reply.Extra)
+	re := len(reply.Extra) - sub
 	l, m := 0, 0
 	for l < re {
 		m = (l + re) / 2
@@ -221,8 +225,8 @@ func (r *Request) Scrub(reply *dns.Msg) (*dns.Msg, Result) {
 			break
 		}
 	}
-	// We may come out of this loop with one rotation too many as we don't break on rl == size.
-	// I.e. m makes it too large, but m-1 works.
+
+	// We may come out of this loop with one rotation too many, m makes it too large, but m-1 works.
 	if rl > size && m > 0 {
 		reply.Extra = origExtra[:m-1]
 		rl = reply.Len()
@@ -252,16 +256,16 @@ func (r *Request) Scrub(reply *dns.Msg) (*dns.Msg, Result) {
 			break
 		}
 	}
-	// We may come out of this loop with one rotation too many as we don't break on rl == size.
-	// I.e. m makes it too large, but m-1 works.
+
+	// We may come out of this loop with one rotation too many, m makes it too large, but m-1 works.
 	if rl > size && m > 0 {
 		reply.Answer = origAnswer[:m-1]
 		// No need to recalc length, as we don't use it. We set truncated anyway. Doing
 		// this extra m-1 step does make it fit in the client's buffer however.
 	}
 
-	// It now fits, but Truncated.
-	r.SizeAndDo(reply)
+	// It now fits, but Truncated. We can't call sizeAndDo() because that adds a new record (OPT)
+	// in the additional section.
 	reply.Truncated = true
 	return reply, ScrubAnswer
 }
@@ -389,4 +393,5 @@ const (
 	// TODO(miek): make this less awkward.
 	doTrue  = 1
 	doFalse = 2
+	optLen  = 12 // OPT record length.
 )
diff --git a/request/request_test.go b/request/request_test.go
index ff9efcfd0..abb38dea3 100644
--- a/request/request_test.go
+++ b/request/request_test.go
@@ -109,6 +109,56 @@ func TestRequestScrubExtra(t *testing.T) {
 	}
 }
 
+func TestRequestScrubExtraEdns0(t *testing.T) {
+	m := new(dns.Msg)
+	m.SetQuestion("large.example.com.", dns.TypeSRV)
+	m.SetEdns0(4096, true)
+	req := Request{W: &test.ResponseWriter{}, Req: m}
+
+	reply := new(dns.Msg)
+	reply.SetReply(m)
+	for i := 1; i < 200; i++ {
+		reply.Extra = append(reply.Extra, test.SRV(
+			fmt.Sprintf("large.example.com. 10 IN SRV 0 0 80 10-0-0-%d.default.pod.k8s.example.com.", i)))
+	}
+
+	_, got := req.Scrub(reply)
+	if want := ScrubExtra; want != got {
+		t.Errorf("want scrub result %d, got %d", want, got)
+	}
+	if want, got := req.Size(), reply.Len(); want < got {
+		t.Errorf("want scrub to reduce message length below %d bytes, got %d bytes", want, got)
+	}
+	if reply.Truncated {
+		t.Errorf("want scrub to not set truncated bit")
+	}
+	opt := reply.Extra[len(reply.Extra)-1]
+	if opt.Header().Rrtype != dns.TypeOPT {
+		t.Errorf("Last RR must be OPT record")
+	}
+}
+
+func TestRequestScrubAnswerExact(t *testing.T) {
+	m := new(dns.Msg)
+	m.SetQuestion("large.example.com.", dns.TypeSRV)
+	m.SetEdns0(867, false) // Bit fiddly, but this hits the rl == size break clause in Scrub, 52 RRs should remain.
+	req := Request{W: &test.ResponseWriter{}, Req: m}
+
+	reply := new(dns.Msg)
+	reply.SetReply(m)
+	for i := 1; i < 200; i++ {
+		reply.Answer = append(reply.Answer, test.A(fmt.Sprintf("large.example.com. 10 IN A 127.0.0.%d", i)))
+	}
+
+	_, got := req.Scrub(reply)
+	if want := ScrubAnswer; want != got {
+		t.Errorf("want scrub result %d, got %d", want, got)
+	}
+	if want, got := req.Size(), reply.Len(); want < got {
+		t.Errorf("want scrub to reduce message length below %d bytes, got %d bytes", want, got)
+	}
+}
+
 func TestRequestMatch(t *testing.T) {
 	st := testRequest()
 	reply := new(dns.Msg)