Scrub: use binary search (#1543)

Use binary search to find the minimal message size, that contains whole
RRs and fits the client's buffer. This is better then just setting
entire sections to `nil`. Extend the tests to test for additional and
answer section truncation. In the first case we *don't* set the TC bit.

This function now also set Compression to true.

2 files changed, 101 insertions, 27 deletions

diff --git a/request/request.go b/request/request.go
index e2ff2008c..1105f919a 100644
--- a/request/request.go
+++ b/request/request.go
@@ -179,34 +179,86 @@ type Result int
 const (
 	// ScrubIgnored is returned when Scrub did nothing to the message.
 	ScrubIgnored Result = iota
-	// ScrubDone is returned when the reply has been scrubbed.
-	ScrubDone
+	// ScrubExtra is returned when the reply has been scrubbed by removing RRs from the additional section.
+	ScrubExtra
+	// ScrubAnswer is returned when the reply has been scrubbed by removing RRs from the answer section.
+	ScrubAnswer
 )
 
-// Scrub scrubs the reply message so that it will fit the client's buffer. If
-// even after dropping the additional section it does not fit, the answer will
-// be cleared and the TC bit will be set on the message. Note, the TC bit will
-// be set regardless of protocol, even TCP message will get the bit, the client
-// should then retry with pigeons. TODO(referral).
+// Scrub scrubs the reply message so that it will fit the client's buffer. It sets
+// reply.Compress to true.
+// Scrub uses binary search to find a save cut off point in the additional section.
+// If even *without* the additional section the reply still doesn't fit we
+// repeat this process for the answer section. If we scrub the answer section
+// we set the TC bit on the reply; indicating the client should retry over TCP.
+// Note, the TC bit will be set regardless of protocol, even TCP message will
+// get the bit, the client should then retry with pigeons.
 func (r *Request) Scrub(reply *dns.Msg) (*dns.Msg, Result) {
+	reply.Compress = true
+
 	size := r.Size()
-	l := reply.Len()
-	if size >= l {
+	rl := reply.Len()
+
+	if size >= rl {
 		return reply, ScrubIgnored
 	}
-	// TODO(miek): check for delegation
 
-	// If not delegation, drop additional section.
-	reply.Extra = nil
-	r.SizeAndDo(reply)
-	l = reply.Len()
-	if size >= l {
-		return reply, ScrubDone
+	origExtra := reply.Extra
+	re := len(reply.Extra)
+	l, m := 0, 0
+	for l < re {
+		m = (l + re) / 2
+		reply.Extra = origExtra[:m]
+		rl = reply.Len()
+		if rl < size {
+			l = m + 1
+			continue
+		}
+		if rl > size {
+			re = m - 1
+			continue
+		}
+	}
+	// 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.
+	if rl > size && m > 0 {
+		reply.Extra = origExtra[:m-1]
+		rl = reply.Len()
+	}
+
+	if rl < size {
+		r.SizeAndDo(reply)
+		return reply, ScrubExtra
 	}
 
+	origAnswer := reply.Answer
+	ra := len(reply.Answer)
+	l, m = 0, 0
+	for l < ra {
+		m = (l + ra) / 2
+		reply.Answer = origAnswer[:m]
+		rl = reply.Len()
+		if rl < size {
+			l = m + 1
+			continue
+		}
+		if rl > size {
+			ra = m - 1
+			continue
+		}
+	}
+	// 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.
+	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)
 	reply.Truncated = true
-	reply.Answer = nil
-	return reply, ScrubDone
+	return reply, ScrubAnswer
 }
 
 // Type returns the type of the question as a string. If the request is malformed
diff --git a/request/request_test.go b/request/request_test.go
index 49b825627..e6fcf6935 100644
--- a/request/request_test.go
+++ b/request/request_test.go
@@ -61,7 +61,7 @@ func TestRequestMalformed(t *testing.T) {
 	}
 }
 
-func TestRequestScrub(t *testing.T) {
+func TestRequestScrubAnswer(t *testing.T) {
 	m := new(dns.Msg)
 	m.SetQuestion("large.example.com.", dns.TypeSRV)
 	req := Request{W: &test.ResponseWriter{}, Req: m}
@@ -69,24 +69,46 @@ func TestRequestScrub(t *testing.T) {
 	reply := new(dns.Msg)
 	reply.SetReply(m)
 	for i := 1; i < 200; i++ {
-		reply.Answer = append(reply.Answer, test.SRV(fmt.Sprintf(
-			"large.example.com. 10 IN SRV 0 0 80 10-0-0-%d.default.pod.k8s.example.com.",
-			i,
-		)))
+		reply.Answer = append(reply.Answer, test.SRV(
+			fmt.Sprintf("large.example.com. 10 IN SRV 0 0 80 10-0-0-%d.default.pod.k8s.example.com.", i)))
 	}
 
-	msg, got := req.Scrub(reply)
-	if want := ScrubDone; want != got {
+	_, got := req.Scrub(reply)
+	if want := ScrubAnswer; want != got {
 		t.Errorf("want scrub result %d, got %d", want, got)
 	}
-	if want, got := req.Size(), msg.Len(); 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 !msg.Truncated {
+	if !reply.Truncated {
 		t.Errorf("want scrub to set truncated bit")
 	}
 }
 
+func TestRequestScrubExtra(t *testing.T) {
+	m := new(dns.Msg)
+	m.SetQuestion("large.example.com.", dns.TypeSRV)
+	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")
+	}
+}
+
 func BenchmarkRequestDo(b *testing.B) {
 	st := testRequest()