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.

1 files changed, 70 insertions, 18 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