1 files changed, 54 insertions, 1 deletions
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}
 }