request: re-adjust size for edns0 RR (#2007)

* Add regression test for broken DNS truncation

With the latest changes to the implementation, request.Scrub() returns
invalid response messages which are larger than the advertised buffer
size. This commit adds a failing test for that.

* Add benchmark for response truncation

* request: re-adjust size for edns0 RR

Cherry-picked the test and benchmark commits from #1962 and make the
failing test not fail, by adding minimal fix in request.go.

This makes size 12 bytes smaller if we need to add back the OPT RR in
the message.

Closes: #1962

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

* Scrub: correctly check for EDNS0, not only Do

Make Scrub check for EDNS0, not only if the Do bit is set.

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

2 files changed, 63 insertions, 8 deletions

diff --git a/request/request.go b/request/request.go
index 06f840f89..3ee4c2126 100644
--- a/request/request.go
+++ b/request/request.go
@@ -263,12 +263,16 @@ func (r *Request) Scrub(reply *dns.Msg) (*dns.Msg, Result) {
 
 	// Account for the OPT record that gets added in SizeAndDo(), subtract that length.
 	sub := 0
-	if r.Do() {
+	if r.Req.IsEdns0() != nil {
 		sub = optLen
 	}
-	origExtra := reply.Extra
+
+	// substract 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 {
 		m = (l + re) / 2
 		reply.Extra = origExtra[:m]
@@ -297,9 +301,9 @@ func (r *Request) Scrub(reply *dns.Msg) (*dns.Msg, Result) {
 		return reply, ScrubExtra
 	}
 
-	origAnswer := reply.Answer
 	ra := len(reply.Answer)
 	l, m = 0, 0
+	origAnswer := reply.Answer
 	for l < ra {
 		m = (l + ra) / 2
 		reply.Answer = origAnswer[:m]
@@ -324,14 +328,12 @@ func (r *Request) Scrub(reply *dns.Msg) (*dns.Msg, Result) {
 		// this extra m-1 step does make it fit in the client's buffer however.
 	}
 
-	// It now fits, but Truncated. We can't call sizeAndDo() because that adds a new record (OPT)
-	// in the additional section.
+	r.SizeAndDo(reply)
 	reply.Truncated = true
 	return reply, ScrubAnswer
 }
 
-// Type returns the type of the question as a string. If the request is malformed
-// the empty string is returned.
+// Type returns the type of the question as a string. If the request is malformed the empty string is returned.
 func (r *Request) Type() string {
 	if r.Req == nil {
 		return ""
diff --git a/request/request_test.go b/request/request_test.go
index c58612605..cad2dce0d 100644
--- a/request/request_test.go
+++ b/request/request_test.go
@@ -138,6 +138,39 @@ func TestRequestScrubExtraEdns0(t *testing.T) {
 	}
 }
 
+func TestRequestScrubExtraRegression(t *testing.T) {
+	m := new(dns.Msg)
+	m.SetQuestion("large.example.com.", dns.TypeSRV)
+	m.SetEdns0(2048, true)
+	req := Request{W: &test.ResponseWriter{}, Req: m}
+
+	reply := new(dns.Msg)
+	reply.SetReply(m)
+	for i := 1; i < 33; 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)))
+	}
+	for i := 1; i < 33; i++ {
+		reply.Extra = append(reply.Extra, test.A(
+			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)
+	}
+	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)
@@ -196,10 +229,30 @@ func BenchmarkRequestSize(b *testing.B) {
 	}
 }
 
+func BenchmarkRequestScrub(b *testing.B) {
+	st := testRequest()
+
+	reply := new(dns.Msg)
+	reply.SetReply(st.Req)
+	for i := 1; i < 33; 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)))
+	}
+	for i := 1; i < 33; i++ {
+		reply.Extra = append(reply.Extra, test.A(
+			fmt.Sprintf("10-0-0-%d.default.pod.k8s.example.com. 10 IN A 10.0.0.%d", i, i)))
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		st.Scrub(reply.Copy())
+	}
+}
+
 func testRequest() Request {
 	m := new(dns.Msg)
 	m.SetQuestion("example.com.", dns.TypeA)
-	m.SetEdns0(4097, true)
+	m.SetEdns0(4096, true)
 	return Request{W: &test.ResponseWriter{}, Req: m}
 }