plugin/rewrite: add closing dot for suffix rewrite rule (#2070)

* add closing dot for suffix rewrite rule
* improve rule syntax checks

Resolves: #1881

4 files changed, 196 insertions, 34 deletions

diff --git a/plugin/rewrite/README.md b/plugin/rewrite/README.md
index d72f6f930..05b3efa6d 100644
--- a/plugin/rewrite/README.md
+++ b/plugin/rewrite/README.md
@@ -185,6 +185,16 @@ follows:
 rewrite [continue|stop] name regex STRING STRING answer name STRING STRING
 ```
 
+When using `exact` name rewrite rules, answer gets re-written automatically,
+and there is no need defining `answer name` instruction. The below rule
+rewrites the name in a request from `RED` to `BLUE`, and subsequently
+rewrites the name in a corresponding response from `BLUE` to `RED`. The
+client in the request would see only `RED` and no `BLUE`.
+
+```
+rewrite [continue|stop] name exact RED BLUE
+```
+
 ### TTL Field Rewrites
 
 At times, the need for rewriting TTL value could arise. For example, a DNS server
diff --git a/plugin/rewrite/name.go b/plugin/rewrite/name.go
index 7c3371b8f..3be1aec4d 100644
--- a/plugin/rewrite/name.go
+++ b/plugin/rewrite/name.go
@@ -13,10 +13,11 @@ import (
 	"github.com/miekg/dns"
 )
 
-type nameRule struct {
+type exactNameRule struct {
 	NextAction string
 	From       string
 	To         string
+	ResponseRule
 }
 
 type prefixNameRule struct {
@@ -59,7 +60,7 @@ const (
 
 // Rewrite rewrites the current request based upon exact match of the name
 // in the question section of the request.
-func (rule *nameRule) Rewrite(ctx context.Context, state request.Request) Result {
+func (rule *exactNameRule) Rewrite(ctx context.Context, state request.Request) Result {
 	if rule.From == state.Name() {
 		state.Req.Question[0].Name = rule.To
 		return RewriteDone
@@ -115,76 +116,141 @@ func (rule *regexNameRule) Rewrite(ctx context.Context, state request.Request) R
 
 // newNameRule creates a name matching rule based on exact, partial, or regex match
 func newNameRule(nextAction string, args ...string) (Rule, error) {
+	var matchType, rewriteQuestionFrom, rewriteQuestionTo string
+	var rewriteAnswerField, rewriteAnswerFrom, rewriteAnswerTo string
 	if len(args) < 2 {
 		return nil, fmt.Errorf("too few arguments for a name rule")
 	}
-	if len(args) == 3 {
-		switch strings.ToLower(args[0]) {
+	if len(args) == 2 {
+		matchType = "exact"
+		rewriteQuestionFrom = plugin.Name(args[0]).Normalize()
+		rewriteQuestionTo = plugin.Name(args[1]).Normalize()
+	}
+	if len(args) >= 3 {
+		matchType = strings.ToLower(args[0])
+		rewriteQuestionFrom = plugin.Name(args[1]).Normalize()
+		rewriteQuestionTo = plugin.Name(args[2]).Normalize()
+	}
+	if matchType == RegexMatch {
+		rewriteQuestionFrom = args[1]
+		rewriteQuestionTo = args[2]
+	}
+	if matchType == ExactMatch || matchType == SuffixMatch {
+		if !hasClosingDot(rewriteQuestionFrom) {
+			rewriteQuestionFrom = rewriteQuestionFrom + "."
+		}
+		if !hasClosingDot(rewriteQuestionTo) {
+			rewriteQuestionTo = rewriteQuestionTo + "."
+		}
+	}
+
+	if len(args) > 3 && len(args) != 7 {
+		return nil, fmt.Errorf("response rewrites must consist only of a name rule with 3 arguments and an answer rule with 3 arguments")
+	}
+
+	if len(args) < 7 {
+		switch matchType {
 		case ExactMatch:
-			return &nameRule{nextAction, plugin.Name(args[1]).Normalize(), plugin.Name(args[2]).Normalize()}, nil
+			rewriteAnswerFromPattern, err := isValidRegexPattern(rewriteQuestionTo, rewriteQuestionFrom)
+			if err != nil {
+				return nil, err
+			}
+			return &exactNameRule{
+				nextAction,
+				rewriteQuestionFrom,
+				rewriteQuestionTo,
+				ResponseRule{
+					Active:      true,
+					Type:        "name",
+					Pattern:     rewriteAnswerFromPattern,
+					Replacement: rewriteQuestionFrom,
+				},
+			}, nil
 		case PrefixMatch:
-			return &prefixNameRule{nextAction, plugin.Name(args[1]).Normalize(), plugin.Name(args[2]).Normalize()}, nil
+			return &prefixNameRule{
+				nextAction,
+				rewriteQuestionFrom,
+				rewriteQuestionTo,
+			}, nil
 		case SuffixMatch:
-			return &suffixNameRule{nextAction, plugin.Name(args[1]).Normalize(), plugin.Name(args[2]).Normalize()}, nil
+			return &suffixNameRule{
+				nextAction,
+				rewriteQuestionFrom,
+				rewriteQuestionTo,
+			}, nil
 		case SubstringMatch:
-			return &substringNameRule{nextAction, plugin.Name(args[1]).Normalize(), plugin.Name(args[2]).Normalize()}, nil
+			return &substringNameRule{
+				nextAction,
+				rewriteQuestionFrom,
+				rewriteQuestionTo,
+			}, nil
 		case RegexMatch:
-			regexPattern, err := regexp.Compile(args[1])
+			rewriteQuestionFromPattern, err := isValidRegexPattern(rewriteQuestionFrom, rewriteQuestionTo)
 			if err != nil {
-				return nil, fmt.Errorf("Invalid regex pattern in a name rule: %s", args[1])
+				return nil, err
 			}
-			return &regexNameRule{nextAction, regexPattern, plugin.Name(args[2]).Normalize(), ResponseRule{Type: "name"}}, nil
+			rewriteQuestionTo := plugin.Name(args[2]).Normalize()
+			return &regexNameRule{
+				nextAction,
+				rewriteQuestionFromPattern,
+				rewriteQuestionTo,
+				ResponseRule{
+					Type: "name",
+				},
+			}, nil
 		default:
-			return nil, fmt.Errorf("A name rule supports only exact, prefix, suffix, substring, and regex name matching")
+			return nil, fmt.Errorf("A name rule supports only exact, prefix, suffix, substring, and regex name matching, received: %s", matchType)
 		}
 	}
 	if len(args) == 7 {
-		if strings.ToLower(args[0]) == RegexMatch {
+		if matchType == RegexMatch {
 			if args[3] != "answer" {
 				return nil, fmt.Errorf("exceeded the number of arguments for a regex name rule")
 			}
-			switch strings.ToLower(args[4]) {
+			rewriteQuestionFromPattern, err := isValidRegexPattern(rewriteQuestionFrom, rewriteQuestionTo)
+			if err != nil {
+				return nil, err
+			}
+			rewriteAnswerField = strings.ToLower(args[4])
+			switch rewriteAnswerField {
 			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])
+			rewriteAnswerFrom = args[5]
+			rewriteAnswerTo = args[6]
+			rewriteAnswerFromPattern, err := isValidRegexPattern(rewriteAnswerFrom, rewriteAnswerTo)
 			if err != nil {
-				return nil, fmt.Errorf("Invalid regex pattern in a name rule: %s", args)
+				return nil, err
 			}
+			rewriteQuestionTo = plugin.Name(args[2]).Normalize()
+			rewriteAnswerTo = plugin.Name(args[6]).Normalize()
 			return &regexNameRule{
 				nextAction,
-				regexPattern,
-				plugin.Name(args[2]).Normalize(),
+				rewriteQuestionFromPattern,
+				rewriteQuestionTo,
 				ResponseRule{
 					Active:      true,
 					Type:        "name",
-					Pattern:     responseRegexPattern,
-					Replacement: plugin.Name(args[6]).Normalize(),
+					Pattern:     rewriteAnswerFromPattern,
+					Replacement: rewriteAnswerTo,
 				},
 			}, 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("response rewrites must consist only of a name rule with 3 arguments and an answer rule with 3 arguments")
-	}
-	return &nameRule{nextAction, plugin.Name(args[0]).Normalize(), plugin.Name(args[1]).Normalize()}, nil
+	return nil, fmt.Errorf("the rewrite rule is invalid: %s", args)
 }
 
 // Mode returns the processing nextAction
-func (rule *nameRule) Mode() string          { return rule.NextAction }
+func (rule *exactNameRule) Mode() string     { return rule.NextAction }
 func (rule *prefixNameRule) Mode() string    { return rule.NextAction }
 func (rule *suffixNameRule) Mode() string    { return rule.NextAction }
 func (rule *substringNameRule) Mode() string { return rule.NextAction }
 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{} }
+func (rule *exactNameRule) GetResponseRule() ResponseRule { return rule.ResponseRule }
 
 // GetResponseRule return a rule to rewrite the response with. Currently not implemented.
 func (rule *prefixNameRule) GetResponseRule() ResponseRule { return ResponseRule{} }
@@ -210,3 +276,34 @@ func validName(s string) bool {
 
 	return true
 }
+
+// hasClosingDot return true if s has a closing dot at the end.
+func hasClosingDot(s string) bool {
+	if strings.HasSuffix(s, ".") {
+		return true
+	}
+	return false
+}
+
+// getSubExprUsage return the number of subexpressions used in s.
+func getSubExprUsage(s string) int {
+	subExprUsage := 0
+	for i := 0; i <= 100; i++ {
+		if strings.Contains(s, "{"+strconv.Itoa(i)+"}") {
+			subExprUsage++
+		}
+	}
+	return subExprUsage
+}
+
+// isValidRegexPattern return a regular expression for pattern matching or errors, if any.
+func isValidRegexPattern(rewriteFrom, rewriteTo string) (*regexp.Regexp, error) {
+	rewriteFromPattern, err := regexp.Compile(rewriteFrom)
+	if err != nil {
+		return nil, fmt.Errorf("Invalid regex matching pattern: %s", rewriteFrom)
+	}
+	if getSubExprUsage(rewriteTo) > rewriteFromPattern.NumSubexp() {
+		return nil, fmt.Errorf("The rewrite regex pattern (%s) uses more subexpressions than its corresponding matching regex pattern (%s)", rewriteTo, rewriteFrom)
+	}
+	return rewriteFromPattern, nil
+}
diff --git a/plugin/rewrite/name_test.go b/plugin/rewrite/name_test.go
index 57fb1eaa9..fe2a2eb06 100644
--- a/plugin/rewrite/name_test.go
+++ b/plugin/rewrite/name_test.go
@@ -31,3 +31,58 @@ func TestRewriteIllegalName(t *testing.T) {
 		t.Errorf("Expected invalid name, got %s", err.Error())
 	}
 }
+
+func TestNewNameRule(t *testing.T) {
+	tests := []struct {
+		next         string
+		args         []string
+		expectedFail bool
+	}{
+		{"stop", []string{"exact", "srv3.coredns.rocks", "srv4.coredns.rocks"}, false},
+		{"stop", []string{"srv1.coredns.rocks", "srv2.coredns.rocks"}, false},
+		{"stop", []string{"suffix", "coredns.rocks", "coredns.rocks."}, false},
+		{"stop", []string{"suffix", "coredns.rocks.", "coredns.rocks"}, false},
+		{"stop", []string{"suffix", "coredns.rocks.", "coredns.rocks."}, false},
+		{"stop", []string{"regex", "srv1.coredns.rocks", "10"}, false},
+		{"stop", []string{"regex", "(.*).coredns.rocks", "10"}, false},
+		{"stop", []string{"regex", "(.*).coredns.rocks", "{1}.coredns.rocks"}, false},
+		{"stop", []string{"regex", "(.*).coredns.rocks", "{1}.{2}.coredns.rocks"}, true},
+		{"stop", []string{"regex", "staging.mydomain.com", "aws-loadbalancer-id.us-east-1.elb.amazonaws.com"}, false},
+	}
+	for i, tc := range tests {
+		failed := false
+		rule, err := newNameRule(tc.next, tc.args...)
+		if err != nil {
+			failed = true
+		}
+		if !failed && !tc.expectedFail {
+			t.Logf("Test %d: PASS, passed as expected: (%s) %s", i, tc.next, tc.args)
+			continue
+		}
+		if failed && tc.expectedFail {
+			t.Logf("Test %d: PASS, failed as expected: (%s) %s: %s", i, tc.next, tc.args, err)
+			continue
+		}
+		if failed && !tc.expectedFail {
+			t.Fatalf("Test %d: FAIL, expected fail=%t, but received fail=%t: (%s) %s, rule=%v, error=%s", i, tc.expectedFail, failed, tc.next, tc.args, rule, err)
+		}
+		t.Fatalf("Test %d: FAIL, expected fail=%t, but received fail=%t: (%s) %s, rule=%v", i, tc.expectedFail, failed, tc.next, tc.args, rule)
+	}
+	for i, tc := range tests {
+		failed := false
+		tc.args = append([]string{tc.next, "name"}, tc.args...)
+		rule, err := newRule(tc.args...)
+		if err != nil {
+			failed = true
+		}
+		if !failed && !tc.expectedFail {
+			t.Logf("Test %d: PASS, passed as expected: (%s) %s", i, tc.next, tc.args)
+			continue
+		}
+		if failed && tc.expectedFail {
+			t.Logf("Test %d: PASS, failed as expected: (%s) %s: %s", i, tc.next, tc.args, err)
+			continue
+		}
+		t.Fatalf("Test %d: FAIL, expected fail=%t, but received fail=%t: (%s) %s, rule=%v", i, tc.expectedFail, failed, tc.next, tc.args, rule)
+	}
+}
diff --git a/plugin/rewrite/rewrite_test.go b/plugin/rewrite/rewrite_test.go
index 670467664..316921954 100644
--- a/plugin/rewrite/rewrite_test.go
+++ b/plugin/rewrite/rewrite_test.go
@@ -31,8 +31,8 @@ func TestNewRule(t *testing.T) {
 		{[]string{"name"}, true, nil},
 		{[]string{"name", "a.com"}, true, nil},
 		{[]string{"name", "a.com", "b.com", "c.com"}, true, nil},
-		{[]string{"name", "a.com", "b.com"}, false, reflect.TypeOf(&nameRule{})},
-		{[]string{"name", "exact", "a.com", "b.com"}, false, reflect.TypeOf(&nameRule{})},
+		{[]string{"name", "a.com", "b.com"}, false, reflect.TypeOf(&exactNameRule{})},
+		{[]string{"name", "exact", "a.com", "b.com"}, false, reflect.TypeOf(&exactNameRule{})},
 		{[]string{"name", "prefix", "a.com", "b.com"}, false, reflect.TypeOf(&prefixNameRule{})},
 		{[]string{"name", "suffix", "a.com", "b.com"}, false, reflect.TypeOf(&suffixNameRule{})},
 		{[]string{"name", "substring", "a.com", "b.com"}, false, reflect.TypeOf(&substringNameRule{})},
@@ -105,8 +105,8 @@ func TestNewRule(t *testing.T) {
 		{[]string{"edns0", "subnet", "append", "24", "56"}, false, reflect.TypeOf(&edns0SubnetRule{})},
 		{[]string{"edns0", "subnet", "replace", "24", "56"}, false, reflect.TypeOf(&edns0SubnetRule{})},
 		{[]string{"unknown-action", "name", "a.com", "b.com"}, true, nil},
-		{[]string{"stop", "name", "a.com", "b.com"}, false, reflect.TypeOf(&nameRule{})},
-		{[]string{"continue", "name", "a.com", "b.com"}, false, reflect.TypeOf(&nameRule{})},
+		{[]string{"stop", "name", "a.com", "b.com"}, false, reflect.TypeOf(&exactNameRule{})},
+		{[]string{"continue", "name", "a.com", "b.com"}, false, reflect.TypeOf(&exactNameRule{})},
 		{[]string{"unknown-action", "type", "any", "a"}, true, nil},
 		{[]string{"stop", "type", "any", "a"}, false, reflect.TypeOf(&typeRule{})},
 		{[]string{"continue", "type", "any", "a"}, false, reflect.TypeOf(&typeRule{})},