feature: plugin/rewrite: rewrite ANSWER SECTION (#1318)

Resolves: #1313

plugin/rewrite/class.go

@@ -41,3 +41,8 @@ func (rule *classRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
 func (rule *classRule) Mode() string {
 	return rule.NextAction
 }
+
+// GetResponseRule return a rule to rewrite the response with. Currently not implemented.
+func (rule *classRule) GetResponseRule() ResponseRule {
+	return ResponseRule{}
+}

plugin/rewrite/edns0.go

@@ -78,6 +78,11 @@ func (rule *edns0NsidRule) Mode() string {
 	return rule.mode
 }
 
+// GetResponseRule return a rule to rewrite the response with. Currently not implemented.
+func (rule *edns0NsidRule) GetResponseRule() ResponseRule {
+	return ResponseRule{}
+}
+
 // Rewrite will alter the request EDNS0 local options
 func (rule *edns0LocalRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
 	result := RewriteIgnored
@@ -115,6 +120,11 @@ func (rule *edns0LocalRule) Mode() string {
 	return rule.mode
 }
 
+// GetResponseRule return a rule to rewrite the response with. Currently not implemented.
+func (rule *edns0LocalRule) GetResponseRule() ResponseRule {
+	return ResponseRule{}
+}
+
 // newEdns0Rule creates an EDNS0 rule of the appropriate type based on the args
 func newEdns0Rule(mode string, args ...string) (Rule, error) {
 	if len(args) < 2 {
@@ -312,6 +322,11 @@ func (rule *edns0VariableRule) Mode() string {
 	return rule.mode
 }
 
+// GetResponseRule return a rule to rewrite the response with. Currently not implemented.
+func (rule *edns0VariableRule) GetResponseRule() ResponseRule {
+	return ResponseRule{}
+}
+
 func isValidVariable(variable string) bool {
 	switch variable {
 	case
@@ -423,6 +438,11 @@ func (rule *edns0SubnetRule) Mode() string {
 	return rule.mode
 }
 
+// GetResponseRule return a rule to rewrite the response with. Currently not implemented.
+func (rule *edns0SubnetRule) GetResponseRule() ResponseRule {
+	return ResponseRule{}
+}
+
 // These are all defined actions.
 const (
 	Replace = "replace"

plugin/rewrite/name.go

@@ -37,6 +37,7 @@ type regexNameRule struct {
 	NextAction  string
 	Pattern     *regexp.Regexp
 	Replacement string
+	ResponseRule
 }
 
 const (
@@ -113,9 +114,6 @@ func newNameRule(nextAction string, args ...string) (Rule, error) {
 	if len(args) < 2 {
 		return nil, fmt.Errorf("too few arguments for a name rule")
 	}
-	if len(args) > 3 {
-		return nil, fmt.Errorf("exceeded the number of arguments for a name rule")
-	}
 	if len(args) == 3 {
 		switch strings.ToLower(args[0]) {
 		case ExactMatch:
@@ -131,11 +129,45 @@ func newNameRule(nextAction string, args ...string) (Rule, error) {
 			if err != nil {
 				return nil, fmt.Errorf("Invalid regex pattern in a name rule: %s", args[1])
 			}
-			return &regexNameRule{nextAction, regexPattern, plugin.Name(args[2]).Normalize()}, nil
+			return &regexNameRule{nextAction, regexPattern, plugin.Name(args[2]).Normalize(), ResponseRule{}}, nil
 		default:
 			return nil, fmt.Errorf("A name rule supports only exact, prefix, suffix, substring, and regex name matching")
 		}
 	}
+	if len(args) == 7 {
+		if strings.ToLower(args[0]) == RegexMatch {
+			if args[3] != "answer" {
+				return nil, fmt.Errorf("exceeded the number of arguments for a regex name rule")
+			}
+			switch strings.ToLower(args[4]) {
+			case "name":
+			default:
+				return nil, fmt.Errorf("exceeded the number of arguments for a regex name rule")
+			}
+			regexPattern, err := regexp.Compile(args[1])
+			if err != nil {
+				return nil, fmt.Errorf("Invalid regex pattern in a name rule: %s", args)
+			}
+			responseRegexPattern, err := regexp.Compile(args[5])
+			if err != nil {
+				return nil, fmt.Errorf("Invalid regex pattern in a name rule: %s", args)
+			}
+			return &regexNameRule{
+				nextAction,
+				regexPattern,
+				plugin.Name(args[2]).Normalize(),
+				ResponseRule{
+					Active:      true,
+					Pattern:     responseRegexPattern,
+					Replacement: plugin.Name(args[6]).Normalize(),
+				},
+			}, nil
+		}
+		return nil, fmt.Errorf("the rewrite of response is supported only for name regex rule")
+	}
+	if len(args) > 3 && len(args) != 7 {
+		return nil, fmt.Errorf("exceeded the number of arguments for a name rule")
+	}
 	return &nameRule{nextAction, plugin.Name(args[0]).Normalize(), plugin.Name(args[1]).Normalize()}, nil
 }
 
@@ -159,3 +191,28 @@ func (rule *substringNameRule) Mode() string {
 func (rule *regexNameRule) Mode() string {
 	return rule.NextAction
 }
+
+// GetResponseRule return a rule to rewrite the response with. Currently not implemented.
+func (rule *nameRule) GetResponseRule() ResponseRule {
+	return ResponseRule{}
+}
+
+// GetResponseRule return a rule to rewrite the response with. Currently not implemented.
+func (rule *prefixNameRule) GetResponseRule() ResponseRule {
+	return ResponseRule{}
+}
+
+// GetResponseRule return a rule to rewrite the response with. Currently not implemented.
+func (rule *suffixNameRule) GetResponseRule() ResponseRule {
+	return ResponseRule{}
+}
+
+// GetResponseRule return a rule to rewrite the response with. Currently not implemented.
+func (rule *substringNameRule) GetResponseRule() ResponseRule {
+	return ResponseRule{}
+}
+
+// GetResponseRule return a rule to rewrite the response with.
+func (rule *regexNameRule) GetResponseRule() ResponseRule {
+	return rule.ResponseRule
+}

plugin/rewrite/reverter.go

@@ -1,27 +1,61 @@
 package rewrite
 
-import "github.com/miekg/dns"
+import (
+	"github.com/miekg/dns"
+	"regexp"
+	"strconv"
+	"strings"
+)
+
+// ResponseRule contains a rule to rewrite a response with.
+type ResponseRule struct {
+	Active      bool
+	Pattern     *regexp.Regexp
+	Replacement string
+}
 
 // ResponseReverter reverses the operations done on the question section of a packet.
 // This is need because the client will otherwise disregards the response, i.e.
 // dig will complain with ';; Question section mismatch: got miek.nl/HINFO/IN'
 type ResponseReverter struct {
 	dns.ResponseWriter
-	original dns.Question
+	originalQuestion dns.Question
+	ResponseRewrite  bool
+	ResponseRules    []ResponseRule
 }
 
 // NewResponseReverter returns a pointer to a new ResponseReverter.
 func NewResponseReverter(w dns.ResponseWriter, r *dns.Msg) *ResponseReverter {
 	return &ResponseReverter{
 		ResponseWriter:   w,
-		original:       r.Question[0],
+		originalQuestion: r.Question[0],
 	}
 }
 
 // WriteMsg records the status code and calls the
 // underlying ResponseWriter's WriteMsg method.
 func (r *ResponseReverter) WriteMsg(res *dns.Msg) error {
-	res.Question[0] = r.original
+	res.Question[0] = r.originalQuestion
+	if r.ResponseRewrite {
+		for _, rr := range res.Answer {
+			name := rr.(*dns.A).Hdr.Name
+			for _, rule := range r.ResponseRules {
+				regexGroups := rule.Pattern.FindStringSubmatch(name)
+				if len(regexGroups) == 0 {
+					continue
+				}
+				s := rule.Replacement
+				for groupIndex, groupValue := range regexGroups {
+					groupIndexStr := "{" + strconv.Itoa(groupIndex) + "}"
+					if strings.Contains(s, groupIndexStr) {
+						s = strings.Replace(s, groupIndexStr, groupValue, -1)
+					}
+				}
+				name = s
+			}
+			rr.(*dns.A).Hdr.Name = name
+		}
+	}
 	return r.ResponseWriter.WriteMsg(res)
 }
 

plugin/rewrite/rewrite.go

@@ -45,6 +45,11 @@ func (rw Rewrite) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg
 	for _, rule := range rw.Rules {
 		switch result := rule.Rewrite(w, r); result {
 		case RewriteDone:
+			respRule := rule.GetResponseRule()
+			if respRule.Active == true {
+				wr.ResponseRewrite = true
+				wr.ResponseRules = append(wr.ResponseRules, respRule)
+			}
 			if rule.Mode() == Stop {
 				if rw.noRevert {
 					return plugin.NextOrFailure(rw.Name(), rw.Next, ctx, w, r)
@@ -70,8 +75,10 @@ func (rw Rewrite) Name() string { return "rewrite" }
 type Rule interface {
 	// Rewrite rewrites the current request.
 	Rewrite(dns.ResponseWriter, *dns.Msg) Result
-	// Mode returns the processing mode stop or continue
+	// Mode returns the processing mode stop or continue.
 	Mode() string
+	// GetResponseRule returns the rule to rewrite response with, if any.
+	GetResponseRule() ResponseRule
 }
 
 func newRule(args ...string) (Rule, error) {

plugin/rewrite/rewrite_test.go

@@ -36,6 +36,13 @@ func TestNewRule(t *testing.T) {
 		{[]string{"name", "substring", "a.com", "b.com"}, false, reflect.TypeOf(&substringNameRule{})},
 		{[]string{"name", "regex", "([a])\\.com", "new-{1}.com"}, false, reflect.TypeOf(&regexNameRule{})},
 		{[]string{"name", "regex", "([a]\\.com", "new-{1}.com"}, true, nil},
+		{[]string{"name", "regex", "(dns)\\.(core)\\.(rocks)", "{2}.{1}.{3}", "answer", "name", "(core)\\.(dns)\\.(rocks)", "{2}.{1}.{3}"}, false, reflect.TypeOf(&regexNameRule{})},
+		{[]string{"name", "regex", "(adns)\\.(core)\\.(rocks)", "{2}.{1}.{3}", "answer", "name", "(core)\\.(adns)\\.(rocks)", "{2}.{1}.{3}", "too.long", "way.too.long"}, true, nil},
+		{[]string{"name", "regex", "(bdns)\\.(core)\\.(rocks)", "{2}.{1}.{3}", "NoAnswer", "name", "(core)\\.(bdns)\\.(rocks)", "{2}.{1}.{3}"}, true, nil},
+		{[]string{"name", "regex", "(cdns)\\.(core)\\.(rocks)", "{2}.{1}.{3}", "answer", "ttl", "(core)\\.(cdns)\\.(rocks)", "{2}.{1}.{3}"}, true, nil},
+		{[]string{"name", "regex", "(ddns)\\.(core)\\.(rocks)", "{2}.{1}.{3}", "answer", "name", "\xecore\\.(ddns)\\.(rocks)", "{2}.{1}.{3}"}, true, nil},
+		{[]string{"name", "regex", "\xedns\\.(core)\\.(rocks)", "{2}.{1}.{3}", "answer", "name", "(core)\\.(edns)\\.(rocks)", "{2}.{1}.{3}"}, true, nil},
+		{[]string{"name", "substring", "fcore.dns.rocks", "dns.fcore.rocks", "answer", "name", "(fcore)\\.(dns)\\.(rocks)", "{2}.{1}.{3}"}, true, nil},
 		{[]string{"name", "substring", "a.com", "b.com", "c.com"}, true, nil},
 		{[]string{"type"}, true, nil},
 		{[]string{"type", "a"}, true, nil},
@@ -152,6 +159,8 @@ func TestRewrite(t *testing.T) {
 	rules := []Rule{}
 	r, _ := newNameRule("stop", "from.nl.", "to.nl.")
 	rules = append(rules, r)
+	r, _ = newNameRule("stop", "regex", "(core)\\.(dns)\\.(rocks)\\.(nl)", "{2}.{1}.{3}.{4}", "answer", "name", "(dns)\\.(core)\\.(rocks)\\.(nl)", "{2}.{1}.{3}.{4}")
+	rules = append(rules, r)
 	r, _ = newNameRule("stop", "exact", "from.exact.nl.", "to.nl.")
 	rules = append(rules, r)
 	r, _ = newNameRule("stop", "prefix", "prefix", "to")
@@ -203,6 +212,7 @@ func TestRewrite(t *testing.T) {
 		{"a.nl.", dns.TypeANY, dns.ClassCHAOS, "a.nl.", dns.TypeANY, dns.ClassINET},
 		// class gets rewritten twice because of continue/stop logic: HS to CH, CH to IN
 		{"a.nl.", dns.TypeANY, 4, "a.nl.", dns.TypeANY, dns.ClassINET},
+		{"core.dns.rocks.nl.", dns.TypeA, dns.ClassINET, "dns.core.rocks.nl.", dns.TypeA, dns.ClassINET},
 	}
 
 	ctx := context.TODO()
@@ -224,6 +234,13 @@ func TestRewrite(t *testing.T) {
 		if resp.Question[0].Qclass != tc.toC {
 			t.Errorf("Test %d: Expected Class to be '%d' but was '%d'", i, tc.toC, resp.Question[0].Qclass)
 		}
+		if tc.fromT == dns.TypeA && tc.toT == dns.TypeA {
+			if len(resp.Answer) > 0 {
+				if resp.Answer[0].(*dns.A).Hdr.Name != tc.to {
+					t.Errorf("Test %d: Expected Answer Name to be %q but was %q", i, tc.to, resp.Answer[0].(*dns.A).Hdr.Name)
+				}
+			}
+		}
 	}
 }
 

plugin/rewrite/setup.go

@@ -3,7 +3,6 @@ package rewrite
 import (
 	"github.com/coredns/coredns/core/dnsserver"
 	"github.com/coredns/coredns/plugin"
-
 	"github.com/mholt/caddy"
 )
 
@@ -32,6 +31,12 @@ func rewriteParse(c *caddy.Controller) ([]Rule, error) {
 
 	for c.Next() {
 		args := c.RemainingArgs()
+		if len(args) < 2 {
+			// Handles rules out of nested instructions, i.e. the ones enclosed in curly brackets
+			for c.NextBlock() {
+				args = append(args, c.Val())
+			}
+		}
 		rule, err := newRule(args...)
 		if err != nil {
 			return nil, err

plugin/rewrite/type.go

@@ -42,3 +42,8 @@ func (rule *typeRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
 func (rule *typeRule) Mode() string {
 	return rule.nextAction
 }
+
+// GetResponseRule return a rule to rewrite the response with. Currently not implemented.
+func (rule *typeRule) GetResponseRule() ResponseRule {
+	return ResponseRule{}
+}