plugin/rewrite: add rcode as a rewrite option (#6204)

* plugin/forward add ignore_server_failure for masking upstream server faults

Signed-off-by: schou <pschou@users.noreply.github.com>

* Switch from a ignore_server_fail to a rewrite rcode type.

Signed-off-by: schou <pschou@users.noreply.github.com>

* trim down the tests

Signed-off-by: schou <pschou@users.noreply.github.com>

* fixing readme TTL and using map for rcode

Signed-off-by: schou <pschou@users.noreply.github.com>

* add newline

Signed-off-by: schou <pschou@users.noreply.github.com>

---------

Signed-off-by: schou <pschou@users.noreply.github.com>

4 files changed, 308 insertions, 0 deletions

diff --git a/plugin/rewrite/README.md b/plugin/rewrite/README.md
index f9826cf22..895ef6325 100644
--- a/plugin/rewrite/README.md
+++ b/plugin/rewrite/README.md
@@ -26,6 +26,7 @@ e.g., to rewrite ANY queries to HINFO, use `rewrite type ANY HINFO`.
    * `edns0` - an EDNS0 option can be appended to the request as described below in the **EDNS0 Options** section.
    * `ttl` - the TTL value in the _response_ is rewritten.
    * `cname` - the CNAME target if the response has a CNAME record
+   * `rcode` - the response code (RCODE) value in the _response_ is rewritten.
 
 * **TYPE** this optional element can be specified for a `name` or `ttl` field.
   If not given type `exact` will be assumed. If options should be specified the
@@ -335,6 +336,61 @@ rewrite ttl example.com. 30-
 rewrite ttl example.com. 30 # equivalent to rewrite ttl example.com. 30-30
 ```
 
+### RCODE Field Rewrites
+
+At times, the need to rewrite a RCODE value could arise. For example, a DNS server
+may respond with a SERVFAIL instead of NOERROR records when AAAA records are requested.
+
+In the below example, the rcode value the answer for `coredns.rocks` the replies with SERVFAIL
+is being switched to NOERROR.
+
+This example rewrites all the *.coredns.rocks domain SERVFAIL errors to NOERROR
+```
+    rewrite continue {
+        rcode regex (.*)\.coredns\.rocks SERVFAIL NOERROR
+    }
+```
+
+The same result numeric values:
+```
+    rewrite continue {
+        rcode regex (.*)\.coredns\.rocks 2 0
+    }
+```
+
+The syntax for the RCODE rewrite rule is as follows. The meaning of
+`exact|prefix|suffix|substring|regex` is the same as with the name rewrite rules.
+An omitted type is defaulted to `exact`.
+
+```
+rewrite [continue|stop] rcode [exact|prefix|suffix|substring|regex] STRING FROM TO
+```
+
+The values of FROM and TO can be any of the following, text value or numeric:
+
+```
+  0 NOERROR
+  1 FORMERR
+  2 SERVFAIL
+  3 NXDOMAIN
+  4 NOTIMP
+  5 REFUSED
+  6 YXDOMAIN
+  7 YXRRSET
+  8 NXRRSET
+  9 NOTAUTH
+  10 NOTZONE
+  16 BADSIG
+  17 BADKEY
+  18 BADTIME
+  19 BADMODE
+  20 BADNAME
+  21 BADALG
+  22 BADTRUNC
+  23 BADCOOKIE
+```
+
+
 ## EDNS0 Options
 
 Using the FIELD edns0, you can set, append, or replace specific EDNS0 options in the request.
diff --git a/plugin/rewrite/rcode.go b/plugin/rewrite/rcode.go
new file mode 100644
index 000000000..814b95fd9
--- /dev/null
+++ b/plugin/rewrite/rcode.go
@@ -0,0 +1,178 @@
+package rewrite
+
+import (
+	"context"
+	"fmt"
+	"regexp"
+	"strconv"
+	"strings"
+
+	"github.com/coredns/coredns/plugin"
+	"github.com/coredns/coredns/request"
+
+	"github.com/miekg/dns"
+)
+
+type rcodeResponseRule struct {
+	old int
+	new int
+}
+
+func (r *rcodeResponseRule) RewriteResponse(res *dns.Msg, rr dns.RR) {
+	if r.old == res.MsgHdr.Rcode {
+		res.MsgHdr.Rcode = r.new
+	}
+}
+
+type rcodeRuleBase struct {
+	nextAction string
+	response   rcodeResponseRule
+}
+
+func newRCodeRuleBase(nextAction string, old, new int) rcodeRuleBase {
+	return rcodeRuleBase{
+		nextAction: nextAction,
+		response:   rcodeResponseRule{old: old, new: new},
+	}
+}
+
+func (rule *rcodeRuleBase) responseRule(match bool) (ResponseRules, Result) {
+	if match {
+		return ResponseRules{&rule.response}, RewriteDone
+	}
+	return nil, RewriteIgnored
+}
+
+// Mode returns the processing nextAction
+func (rule *rcodeRuleBase) Mode() string { return rule.nextAction }
+
+type exactRCodeRule struct {
+	rcodeRuleBase
+	From string
+}
+
+type prefixRCodeRule struct {
+	rcodeRuleBase
+	Prefix string
+}
+
+type suffixRCodeRule struct {
+	rcodeRuleBase
+	Suffix string
+}
+
+type substringRCodeRule struct {
+	rcodeRuleBase
+	Substring string
+}
+
+type regexRCodeRule struct {
+	rcodeRuleBase
+	Pattern *regexp.Regexp
+}
+
+// Rewrite rewrites the current request based upon exact match of the name
+// in the question section of the request.
+func (rule *exactRCodeRule) Rewrite(ctx context.Context, state request.Request) (ResponseRules, Result) {
+	return rule.responseRule(rule.From == state.Name())
+}
+
+// Rewrite rewrites the current request when the name begins with the matching string.
+func (rule *prefixRCodeRule) Rewrite(ctx context.Context, state request.Request) (ResponseRules, Result) {
+	return rule.responseRule(strings.HasPrefix(state.Name(), rule.Prefix))
+}
+
+// Rewrite rewrites the current request when the name ends with the matching string.
+func (rule *suffixRCodeRule) Rewrite(ctx context.Context, state request.Request) (ResponseRules, Result) {
+	return rule.responseRule(strings.HasSuffix(state.Name(), rule.Suffix))
+}
+
+// Rewrite rewrites the current request based upon partial match of the
+// name in the question section of the request.
+func (rule *substringRCodeRule) Rewrite(ctx context.Context, state request.Request) (ResponseRules, Result) {
+	return rule.responseRule(strings.Contains(state.Name(), rule.Substring))
+}
+
+// Rewrite rewrites the current request when the name in the question
+// section of the request matches a regular expression.
+func (rule *regexRCodeRule) Rewrite(ctx context.Context, state request.Request) (ResponseRules, Result) {
+	return rule.responseRule(len(rule.Pattern.FindStringSubmatch(state.Name())) != 0)
+}
+
+// newRCodeRule creates a name matching rule based on exact, partial, or regex match
+func newRCodeRule(nextAction string, args ...string) (Rule, error) {
+	if len(args) < 3 {
+		return nil, fmt.Errorf("too few (%d) arguments for a rcode rule", len(args))
+	}
+	var oldStr, newStr string
+	if len(args) == 3 {
+		oldStr, newStr = args[1], args[2]
+	}
+	if len(args) == 4 {
+		oldStr, newStr = args[2], args[3]
+	}
+	old, valid := isValidRCode(oldStr)
+	if !valid {
+		return nil, fmt.Errorf("invalid matching RCODE '%s' for a rcode rule", oldStr)
+	}
+	new, valid := isValidRCode(newStr)
+	if !valid {
+		return nil, fmt.Errorf("invalid replacement RCODE '%s' for a rcode rule", newStr)
+	}
+	if len(args) == 4 {
+		switch strings.ToLower(args[0]) {
+		case ExactMatch:
+			return &exactRCodeRule{
+				newRCodeRuleBase(nextAction, old, new),
+				plugin.Name(args[1]).Normalize(),
+			}, nil
+		case PrefixMatch:
+			return &prefixRCodeRule{
+				newRCodeRuleBase(nextAction, old, new),
+				plugin.Name(args[1]).Normalize(),
+			}, nil
+		case SuffixMatch:
+			return &suffixRCodeRule{
+				newRCodeRuleBase(nextAction, old, new),
+				plugin.Name(args[1]).Normalize(),
+			}, nil
+		case SubstringMatch:
+			return &substringRCodeRule{
+				newRCodeRuleBase(nextAction, old, new),
+				plugin.Name(args[1]).Normalize(),
+			}, nil
+		case RegexMatch:
+			regexPattern, err := regexp.Compile(args[1])
+			if err != nil {
+				return nil, fmt.Errorf("invalid regex pattern in a rcode rule: %s", args[1])
+			}
+			return &regexRCodeRule{
+				newRCodeRuleBase(nextAction, old, new),
+				regexPattern,
+			}, nil
+		default:
+			return nil, fmt.Errorf("rcode rule supports only exact, prefix, suffix, substring, and regex name matching")
+		}
+	}
+	if len(args) > 4 {
+		return nil, fmt.Errorf("many few arguments for a rcode rule")
+	}
+	return &exactRCodeRule{
+		newRCodeRuleBase(nextAction, old, new),
+		plugin.Name(args[0]).Normalize(),
+	}, nil
+}
+
+// validRCode returns true if v is valid RCode value.
+func isValidRCode(v string) (int, bool) {
+	i, err := strconv.ParseUint(v, 10, 32)
+	// try parsing integer based rcode
+	if err == nil && i <= 23 {
+		return int(i), true
+	}
+
+	if RCodeInt, ok := dns.StringToRcode[strings.ToUpper(v)]; ok {
+		return RCodeInt, true
+	}
+	return 0, false
+}
diff --git a/plugin/rewrite/rcode_test.go b/plugin/rewrite/rcode_test.go
new file mode 100644
index 000000000..e40260719
--- /dev/null
+++ b/plugin/rewrite/rcode_test.go
@@ -0,0 +1,72 @@
+package rewrite
+
+import (
+	"testing"
+
+	"github.com/coredns/coredns/plugin/test"
+	"github.com/coredns/coredns/request"
+
+	"github.com/miekg/dns"
+)
+
+func TestNewRCodeRule(t *testing.T) {
+	tests := []struct {
+		next         string
+		args         []string
+		expectedFail bool
+	}{
+		{"stop", []string{"numeric.rcode.coredns.rocks", "2", "0"}, false},
+		{"stop", []string{"too.few.rcode.coredns.rocks", "2"}, true},
+		{"stop", []string{"exact", "too.many.rcode.coredns.rocks", "2", "1", "0"}, true},
+		{"stop", []string{"exact", "match.string.rcode.coredns.rocks", "SERVFAIL", "NOERROR"}, false},
+		{"continue", []string{"regex", `(regex)\.rcode\.(coredns)\.(rocks)`, "FORMERR", "NOERROR"}, false},
+		{"stop", []string{"invalid.rcode.coredns.rocks", "random", "nothing"}, true},
+	}
+	for i, tc := range tests {
+		failed := false
+		rule, err := newRCodeRule(tc.next, tc.args...)
+		if err != nil {
+			failed = true
+		}
+		if !failed && !tc.expectedFail {
+			continue
+		}
+		if failed && tc.expectedFail {
+			continue
+		}
+		t.Fatalf("Test %d: FAIL, expected fail=%t, but received fail=%t: (%s) %s, rule=%v, err=%v", i, tc.expectedFail, failed, tc.next, tc.args, rule, err)
+	}
+	for i, tc := range tests {
+		failed := false
+		tc.args = append([]string{tc.next, "rcode"}, tc.args...)
+		rule, err := newRule(tc.args...)
+		if err != nil {
+			failed = true
+		}
+		if !failed && !tc.expectedFail {
+			continue
+		}
+		if failed && tc.expectedFail {
+			continue
+		}
+		t.Fatalf("Test %d: FAIL, expected fail=%t, but received fail=%t: (%s) %s, rule=%v, err=%v", i, tc.expectedFail, failed, tc.next, tc.args, rule, err)
+	}
+}
+
+func TestRCodeRewrite(t *testing.T) {
+	rule, err := newRCodeRule("stop", []string{"exact", "srv1.coredns.rocks", "SERVFAIL", "FORMERR"}...)
+
+	m := new(dns.Msg)
+	m.SetQuestion("srv1.coredns.rocks.", dns.TypeA)
+	m.Question[0].Qclass = dns.ClassINET
+	m.Answer = []dns.RR{test.A("srv1.coredns.rocks.  5   IN  A  10.0.0.1")}
+	m.MsgHdr.Rcode = dns.RcodeServerFailure
+	request := request.Request{Req: m}
+
+	rcRule, _ := rule.(*exactRCodeRule)
+	var rr dns.RR
+	rcRule.response.RewriteResponse(request.Req, rr)
+	if request.Req.MsgHdr.Rcode != dns.RcodeFormatError {
+		t.Fatalf("RCode rewrite did not apply changes, request=%#v, err=%v", request.Req, err)
+	}
+}
diff --git a/plugin/rewrite/rewrite.go b/plugin/rewrite/rewrite.go
index d991c7f02..edb11813e 100644
--- a/plugin/rewrite/rewrite.go
+++ b/plugin/rewrite/rewrite.go
@@ -141,6 +141,8 @@ func newRule(args ...string) (Rule, error) {
 		return newTTLRule(mode, args[startArg:]...)
 	case "cname":
 		return newCNAMERule(mode, args[startArg:]...)
+	case "rcode":
+		return newRCodeRule(mode, args[startArg:]...)
 	default:
 		return nil, fmt.Errorf("invalid rule type %q", args[0])
 	}