Default to scrubbing replies in the server (#2012)

Every plugin needs to deal with EDNS0 and should call Scrub to make a
message fit the client's buffer. Move this functionality into the server
and wrapping the ResponseWriter into a ScrubWriter that handles these
bits for us. Result:

Less code and faster, because multiple chained plugins could all be
calling scrub and SizeAndDo - now there is just one place.

Most tests in file/* and dnssec/* needed adjusting because in those unit
tests you don't see OPT RRs anymore. The DNSSEC signer was also looking
at the returned OPT RR to see if it needed to sign - as those are now
added by the server (and thus later), this needed to change slightly.

Scrub itself still exist (for backward compat reasons), but has been
made a noop. Scrub has been renamed to scrub as it should not be used by
external plugins.

Fixes: #2010

Signed-off-by: Miek Gieben <miek@miek.nl>

request/request.go

@@ -226,19 +226,11 @@ func (r *Request) SizeAndDo(m *dns.Msg) bool {
 	return true
 }
 
-// Result is the result of Scrub.
-type Result int
+// Scrub is a noop function, added for backwards compatibility reasons. The original Scrub is now called
+// automatically by the server on writing the reply. See ScrubWriter.
+func (r *Request) Scrub(reply *dns.Msg) (*dns.Msg, int) { return reply, 0 }
 
-const (
-	// ScrubIgnored is returned when Scrub did nothing to the message.
-	ScrubIgnored Result = iota
-	// 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. It will first
+// scrub scrubs the reply message so that it will fit the client's buffer. It will first
 // check if the reply fits without compression and then *with* compression.
 // Scrub will then use 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
@@ -246,19 +238,19 @@ const (
 // 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) {
+func (r *Request) scrub(reply *dns.Msg) *dns.Msg {
 	size := r.Size()
 
 	reply.Compress = false
 	rl := reply.Len()
 	if size >= rl {
-		return reply, ScrubIgnored
+		return reply
 	}
 
 	reply.Compress = true
 	rl = reply.Len()
 	if size >= rl {
-		return reply, ScrubIgnored
+		return reply
 	}
 
 	// Account for the OPT record that gets added in SizeAndDo(), subtract that length.
@@ -298,7 +290,7 @@ func (r *Request) Scrub(reply *dns.Msg) (*dns.Msg, Result) {
 
 	if rl < size {
 		r.SizeAndDo(reply)
-		return reply, ScrubExtra
+		return reply
 	}
 
 	ra := len(reply.Answer)
@@ -330,7 +322,7 @@ func (r *Request) Scrub(reply *dns.Msg) (*dns.Msg, Result) {
 
 	r.SizeAndDo(reply)
 	reply.Truncated = true
-	return reply, ScrubAnswer
+	return reply
 }
 
 // Type returns the type of the question as a string. If the request is malformed the empty string is returned.

request/request_test.go

@@ -73,10 +73,7 @@ func TestRequestScrubAnswer(t *testing.T) {
 			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 := ScrubAnswer; want != got {
-		t.Errorf("Want scrub result %d, got %d", want, got)
-	}
+	req.scrub(reply)
 	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)
 	}
@@ -97,10 +94,7 @@ func TestRequestScrubExtra(t *testing.T) {
 			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)
-	}
+	req.scrub(reply)
 	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)
 	}
@@ -122,10 +116,7 @@ func TestRequestScrubExtraEdns0(t *testing.T) {
 			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)
-	}
+	req.scrub(reply)
 	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)
 	}
@@ -155,10 +146,7 @@ func TestRequestScrubExtraRegression(t *testing.T) {
 			fmt.Sprintf("10-0-0-%d.default.pod.k8s.example.com. 10 IN A 10.0.0.%d", i, i)))
 	}
 
-	_, got := req.Scrub(reply)
-	if want := ScrubExtra; want != got {
-		t.Errorf("Want scrub result %d, got %d", want, got)
-	}
+	reply = req.scrub(reply)
 	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)
 	}
@@ -183,10 +171,7 @@ func TestRequestScrubAnswerExact(t *testing.T) {
 		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)
-	}
+	req.scrub(reply)
 	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)
 	}

request/writer.go

@@ -0,0 +1,20 @@
+package request
+
+import "github.com/miekg/dns"
+
+// ScrubWriter will, when writing the message, call scrub to make it fit the client's buffer.
+type ScrubWriter struct {
+	dns.ResponseWriter
+	req *dns.Msg // original request
+}
+
+// NewScrubWriter returns a new and initialized ScrubWriter.
+func NewScrubWriter(req *dns.Msg, w dns.ResponseWriter) *ScrubWriter { return &ScrubWriter{w, req} }
+
+// WriteMsg overrides the default implementation of the underlaying dns.ResponseWriter and calls
+// scrub on the message m and will then write it to the client.
+func (s *ScrubWriter) WriteMsg(m *dns.Msg) error {
+	state := Request{Req: s.req, W: s.ResponseWriter}
+	new, _ := state.Scrub(m)
+	return s.ResponseWriter.WriteMsg(new)
+}