2 files changed, 49 insertions, 11 deletions

diff --git a/request/request.go b/request/request.go
index f560612c8..260d73f5e 100644
--- a/request/request.go
+++ b/request/request.go
@@ -250,18 +250,16 @@ func (r *Request) Scrub(reply *dns.Msg) *dns.Msg {
 	}
 
 	// Account for the OPT record that gets added in SizeAndDo(), subtract that length.
-	sub := 0
+	re := len(reply.Extra)
 	if r.Req.IsEdns0() != nil {
-		sub = optLen
+		size -= optLen
+		// re can never be 0 because we have an OPT RR.
+		re--
 	}
 
-	// subtract to make spaces for re-added EDNS0 OPT RR.
-	re := len(reply.Extra) - sub
-	size -= sub
-
 	l, m := 0, 0
 	origExtra := reply.Extra
-	for l < re {
+	for l <= re {
 		m = (l + re) / 2
 		reply.Extra = origExtra[:m]
 		rl = reply.Len()
@@ -278,13 +276,18 @@ func (r *Request) Scrub(reply *dns.Msg) *dns.Msg {
 		}
 	}
 
-	// We may come out of this loop with one rotation too many, m makes it too large, but m-1 works.
+	// The binary search only breaks on an exact match, which will be
+	// pretty rare. Normally, the loop will exit when l > re, meaning that
+	// in the previous iteration either:
+	// rl < size: no need to do anything.
+	// rl > size: the final size is too large, and if m > 0, the preceeding
+	// iteration the size was too small. Select that preceeding size.
 	if rl > size && m > 0 {
 		reply.Extra = origExtra[:m-1]
 		rl = reply.Len()
 	}
 
-	if rl < size {
+	if rl <= size {
 		r.SizeAndDo(reply)
 		return reply
 	}
@@ -292,7 +295,7 @@ func (r *Request) Scrub(reply *dns.Msg) *dns.Msg {
 	ra := len(reply.Answer)
 	l, m = 0, 0
 	origAnswer := reply.Answer
-	for l < ra {
+	for l <= ra {
 		m = (l + ra) / 2
 		reply.Answer = origAnswer[:m]
 		rl = reply.Len()
@@ -309,7 +312,12 @@ func (r *Request) Scrub(reply *dns.Msg) *dns.Msg {
 		}
 	}
 
-	// We may come out of this loop with one rotation too many, m makes it too large, but m-1 works.
+	// The binary search only breaks on an exact match, which will be
+	// pretty rare. Normally, the loop will exit when l > ra, meaning that
+	// in the previous iteration either:
+	// rl < size: no need to do anything.
+	// rl > size: the final size is too large, and if m > 0, the preceeding
+	// iteration the size was too small. Select that preceeding size.
 	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
diff --git a/request/request_test.go b/request/request_test.go
index bfc95bc5e..5e814f76e 100644
--- a/request/request_test.go
+++ b/request/request_test.go
@@ -159,6 +159,36 @@ func TestRequestScrubExtraRegression(t *testing.T) {
 	}
 }
 
+func TestTruncation(t *testing.T) {
+	for bufsize := 1024; bufsize <= 4096; bufsize += 12 {
+		m := new(dns.Msg)
+		m.SetQuestion("http.service.tcp.srv.k8s.example.org", dns.TypeSRV)
+		m.SetEdns0(uint16(bufsize), true)
+		req := Request{W: &test.ResponseWriter{}, Req: m}
+
+		reply := new(dns.Msg)
+		reply.SetReply(m)
+
+		for i := 0; i < 61; i++ {
+			reply.Answer = append(reply.Answer, test.SRV(fmt.Sprintf("http.service.tcp.srv.k8s.example.org. 5 IN SRV 0 0 80 10-144-230-%d.default.pod.k8s.example.org.", i)))
+		}
+
+		for i := 0; i < 5; i++ {
+			reply.Extra = append(reply.Extra, test.A(fmt.Sprintf("ip-10-10-52-5%d.subdomain.example.org. 5 IN A 10.10.52.5%d", i, i)))
+		}
+
+		for i := 0; i < 5; i++ {
+			reply.Ns = append(reply.Ns, test.NS(fmt.Sprintf("srv.subdomain.example.org. 5 IN NS ip-10-10-33-6%d.subdomain.example.org.", i)))
+		}
+
+		req.Scrub(reply)
+		want, got := req.Size(), reply.Len()
+		if want < got {
+			t.Fatalf("Want scrub to reduce message length below %d bytes, got %d bytes", want, got)
+		}
+	}
+}
+
 func TestRequestScrubAnswerExact(t *testing.T) {
 	m := new(dns.Msg)
 	m.SetQuestion("large.example.com.", dns.TypeSRV)