From 911891f48513eef99ba0c80fae0c1e89f70649fe Mon Sep 17 00:00:00 2001
From: Chris O'Haver <cohaver@infoblox.com>
Date: Tue, 1 Feb 2022 16:54:06 -0500
Subject: [PATCH] plugin/rewrite: Write failures with ResponseReverter (#5150)

* write failures with ResponseReverter instead of letting server write them

Signed-off-by: Chris O'Haver <cohaver@infoblox.com>

* fix comment

Signed-off-by: Chris O'Haver <cohaver@infoblox.com>
---
 plugin/rewrite/rewrite.go | 12 +++++++++++-
 test/rewrite_test.go      | 28 ++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/plugin/rewrite/rewrite.go b/plugin/rewrite/rewrite.go
index a1e474c29..c315fa73a 100644
--- a/plugin/rewrite/rewrite.go
+++ b/plugin/rewrite/rewrite.go
@@ -57,7 +57,17 @@ func (rw Rewrite) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg
 				if !rw.RevertPolicy.DoRevert() {
 					return plugin.NextOrFailure(rw.Name(), rw.Next, ctx, w, r)
 				}
-				return plugin.NextOrFailure(rw.Name(), rw.Next, ctx, wr, r)
+				rcode, err := plugin.NextOrFailure(rw.Name(), rw.Next, ctx, wr, r)
+				if plugin.ClientWrite(rcode) {
+					return rcode, err
+				}
+				// The next plugins didn't write a response, so write one now with the ResponseReverter.
+				// If server.ServeDNS does this then it will create an answer mismatch.
+				res := new(dns.Msg).SetRcode(r, rcode)
+				state.SizeAndDo(res)
+				wr.WriteMsg(res)
+				// return success, so server does not write a second error response to client
+				return dns.RcodeSuccess, err
 			}
 		}
 	}
diff --git a/test/rewrite_test.go b/test/rewrite_test.go
index 310e43efa..d9ef09b15 100644
--- a/test/rewrite_test.go
+++ b/test/rewrite_test.go
@@ -7,6 +7,34 @@ import (
 	"github.com/miekg/dns"
 )
 
+func TestRewriteFailure(t *testing.T) {
+	t.Parallel()
+	i, udp, _, err := CoreDNSServerAndPorts(`.:0 {
+        rewrite name regex (.*)\.test\.$ {1}. answer auto
+        # no next plugin to induce SERVFAIL
+    }`)
+	if err != nil {
+		t.Fatalf("Could not get CoreDNS serving instance: %s", err)
+	}
+
+	defer i.Stop()
+
+	m := new(dns.Msg)
+	m.SetQuestion("example.test.", dns.TypeMX)
+
+	r, err := dns.Exchange(m, udp)
+	if err != nil {
+		t.Fatalf("Expected to receive reply, but didn't: %s", err)
+	}
+
+	if len(r.Question) == 0 {
+		t.Error("Invalid empty question section")
+	}
+	if r.Question[0].Name != "example.test." {
+		t.Errorf("Question section mismatch. expected \"example.test.\" got %q", r.Question[0].Name)
+	}
+}
+
 func TestRewrite(t *testing.T) {
 	t.Parallel()
 	corefile := `.:0 {