plugin/rewrite: regular expression and substring match/replace (#1296) (#1297)

4 files changed, 209 insertions, 10 deletions

diff --git a/plugin/rewrite/README.md b/plugin/rewrite/README.md
index 00f263fb9..199d8a5ec 100644
--- a/plugin/rewrite/README.md
+++ b/plugin/rewrite/README.md
@@ -96,3 +96,48 @@ rewrite edns0 subnet set 24 56
 
 * If the query has source IP as IPv4, the first 24 bits in the IP will be the network subnet.
 * If the query has source IP as IPv6, the first 56 bits in the IP will be the network subnet.
+
+### Name Field Rewrites
+
+The `rewrite` plugin offers the ability to match on the name in the question section of
+a DNS request. The match could be exact, substring, or based on a prefix, suffix, or regular
+expression.
+
+The syntax for the name re-writing is as follows:
+
+```
+rewrite [continue|stop] name [exact|prefix|suffix|substring|regex] STRING STRING
+```
+
+The match type, i.e. `exact`, `substring`, etc., triggers re-write: 
+
+* **exact** (default): on exact match of the name in the question section of a request
+* **substring**: on a partial match of the name in the question section of a request
+* **prefix**: when the name begins with the matching string
+* **suffix**: when the name ends with the matching string
+* **regex**: when the name in the question section of a request matches a regular expression
+
+If the match type is omitted, the `exact` match type is being assumed.
+
+The following instruction allows re-writing the name in the query that
+contains `service.us-west-1.example.org` substring.
+
+```
+rewrite name substring service.us-west-1.example.org service.us-west-1.consul
+```
+
+Thus:
+* Incoming Request Name: `ftp.service.us-west-1.example.org`
+* Re-written Request Name: `ftp.service.us-west-1.consul`
+
+The following instruction uses regular expressions. The name in a request
+matching `(.*)-(us-west-1)\.example\.org` regular expression is being replaces with
+`{1}.service.{2}.consul`, where `{1}` and `{2}` are regular expression match groups.
+
+```
+rewrite name regex (.*)-(us-west-1)\.example\.org {1}.service.{2}.consul
+```
+
+Thus:
+* Incoming Request Name: `ftp-us-west-1.example.org`
+* Re-written Request Name: `ftp.service.us-west-1.consul`
diff --git a/plugin/rewrite/name.go b/plugin/rewrite/name.go
index 016ae4deb..102c8ad09 100644
--- a/plugin/rewrite/name.go
+++ b/plugin/rewrite/name.go
@@ -1,20 +1,59 @@
 package rewrite
 
 import (
+	"fmt"
 	"github.com/coredns/coredns/plugin"
-
 	"github.com/miekg/dns"
+	"regexp"
+	"strconv"
+	"strings"
 )
 
 type nameRule struct {
-	From, To string
+	NextAction string
+	From       string
+	To         string
+}
+
+type prefixNameRule struct {
+	NextAction  string
+	Prefix      string
+	Replacement string
 }
 
-func newNameRule(from, to string) (Rule, error) {
-	return &nameRule{plugin.Name(from).Normalize(), plugin.Name(to).Normalize()}, nil
+type suffixNameRule struct {
+	NextAction  string
+	Suffix      string
+	Replacement string
 }
 
-// Rewrite rewrites the the current request.
+type substringNameRule struct {
+	NextAction  string
+	Substring   string
+	Replacement string
+}
+
+type regexNameRule struct {
+	NextAction  string
+	Pattern     *regexp.Regexp
+	Replacement string
+}
+
+const (
+	// ExactMatch matches only on exact match of the name in the question section of a request
+	ExactMatch = "exact"
+	// PrefixMatch matches when the name begins with the matching string
+	PrefixMatch = "prefix"
+	// SuffixMatch matches when the name ends with the matching string
+	SuffixMatch = "suffix"
+	// SubstringMatch matches on partial match of the name in the question section of a request
+	SubstringMatch = "substring"
+	// RegexMatch matches when the name in the question section of a request matches a regular expression
+	RegexMatch = "regex"
+)
+
+// Rewrite rewrites the current request based upon exact match of the name
+// in the question section of the request
 func (rule *nameRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
 	if rule.From == r.Question[0].Name {
 		r.Question[0].Name = rule.To
@@ -23,7 +62,100 @@ func (rule *nameRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
 	return RewriteIgnored
 }
 
-// Mode returns the processing mode
+// Rewrite rewrites the current request when the name begins with the matching string
+func (rule *prefixNameRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
+	if strings.HasPrefix(r.Question[0].Name, rule.Prefix) {
+		r.Question[0].Name = rule.Replacement + strings.TrimLeft(r.Question[0].Name, rule.Prefix)
+		return RewriteDone
+	}
+	return RewriteIgnored
+}
+
+// Rewrite rewrites the current request when the name ends with the matching string
+func (rule *suffixNameRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
+	if strings.HasSuffix(r.Question[0].Name, rule.Suffix) {
+		r.Question[0].Name = strings.TrimRight(r.Question[0].Name, rule.Suffix) + rule.Replacement
+		return RewriteDone
+	}
+	return RewriteIgnored
+}
+
+// Rewrite rewrites the current request based upon partial match of the
+// name in the question section of the request
+func (rule *substringNameRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
+	if strings.Contains(r.Question[0].Name, rule.Substring) {
+		r.Question[0].Name = strings.Replace(r.Question[0].Name, rule.Substring, rule.Replacement, -1)
+		return RewriteDone
+	}
+	return RewriteIgnored
+}
+
+// Rewrite rewrites the current request when the name in the question
+// section of the request matches a regular expression
+func (rule *regexNameRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
+	regexGroups := rule.Pattern.FindStringSubmatch(r.Question[0].Name)
+	if len(regexGroups) == 0 {
+		return RewriteIgnored
+	}
+	s := rule.Replacement
+	for groupIndex, groupValue := range regexGroups {
+		groupIndexStr := "{" + strconv.Itoa(groupIndex) + "}"
+		if strings.Contains(s, groupIndexStr) {
+			s = strings.Replace(s, groupIndexStr, groupValue, -1)
+		}
+	}
+	r.Question[0].Name = s
+	return RewriteDone
+}
+
+// newNameRule creates a name matching rule based on exact, partial, or regex match
+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:
+			return &nameRule{nextAction, plugin.Name(args[1]).Normalize(), plugin.Name(args[2]).Normalize()}, nil
+		case PrefixMatch:
+			return &prefixNameRule{nextAction, plugin.Name(args[1]).Normalize(), plugin.Name(args[2]).Normalize()}, nil
+		case SuffixMatch:
+			return &suffixNameRule{nextAction, plugin.Name(args[1]).Normalize(), plugin.Name(args[2]).Normalize()}, nil
+		case SubstringMatch:
+			return &substringNameRule{nextAction, plugin.Name(args[1]).Normalize(), plugin.Name(args[2]).Normalize()}, nil
+		case RegexMatch:
+			regexPattern, err := regexp.Compile(args[1])
+			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
+		default:
+			return nil, fmt.Errorf("A name rule supports only exact, prefix, suffix, substring, and regex name matching")
+		}
+	}
+	return &nameRule{nextAction, plugin.Name(args[0]).Normalize(), plugin.Name(args[1]).Normalize()}, nil
+}
+
+// Mode returns the processing nextAction
 func (rule *nameRule) Mode() string {
-	return Stop
+	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
 }
diff --git a/plugin/rewrite/rewrite.go b/plugin/rewrite/rewrite.go
index 0a0bb831c..4924ea6e5 100644
--- a/plugin/rewrite/rewrite.go
+++ b/plugin/rewrite/rewrite.go
@@ -100,12 +100,12 @@ func newRule(args ...string) (Rule, error) {
 		startArg = 1
 	}
 
-	if ruleType != "edns0" && expectNumArgs != 3 {
+	if ruleType != "edns0" && ruleType != "name" && expectNumArgs != 3 {
 		return nil, fmt.Errorf("%s rules must have exactly two arguments", ruleType)
 	}
 	switch ruleType {
 	case "name":
-		return newNameRule(args[startArg], args[startArg+1])
+		return newNameRule(mode, args[startArg:]...)
 	case "class":
 		return newClassRule(args[startArg], args[startArg+1])
 	case "type":
diff --git a/plugin/rewrite/rewrite_test.go b/plugin/rewrite/rewrite_test.go
index 44b3f948f..6314e5407 100644
--- a/plugin/rewrite/rewrite_test.go
+++ b/plugin/rewrite/rewrite_test.go
@@ -30,6 +30,13 @@ func TestNewRule(t *testing.T) {
 		{[]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", "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{})},
+		{[]string{"name", "regex", "([a])\\.com", "new-{1}.com"}, false, reflect.TypeOf(&regexNameRule{})},
+		{[]string{"name", "regex", "([a]\\.com", "new-{1}.com"}, true, nil},
+		{[]string{"name", "substring", "a.com", "b.com", "c.com"}, true, nil},
 		{[]string{"type"}, true, nil},
 		{[]string{"type", "a"}, true, nil},
 		{[]string{"type", "any", "a", "a"}, true, nil},
@@ -143,7 +150,17 @@ func TestNewRule(t *testing.T) {
 
 func TestRewrite(t *testing.T) {
 	rules := []Rule{}
-	r, _ := newNameRule("from.nl.", "to.nl.")
+	r, _ := newNameRule("stop", "from.nl.", "to.nl.")
+	rules = append(rules, r)
+	r, _ = newNameRule("stop", "exact", "from.exact.nl.", "to.nl.")
+	rules = append(rules, r)
+	r, _ = newNameRule("stop", "prefix", "prefix", "to")
+	rules = append(rules, r)
+	r, _ = newNameRule("stop", "suffix", ".suffix.", ".nl.")
+	rules = append(rules, r)
+	r, _ = newNameRule("stop", "substring", "from.substring", "to")
+	rules = append(rules, r)
+	r, _ = newNameRule("stop", "regex", "(f.*m)\\.regex\\.(nl)", "to.{2}")
 	rules = append(rules, r)
 	r, _ = newClassRule("CH", "IN")
 	rules = append(rules, r)
@@ -170,6 +187,11 @@ func TestRewrite(t *testing.T) {
 		{"a.nl.", dns.TypeANY, dns.ClassINET, "a.nl.", dns.TypeHINFO, dns.ClassINET},
 		// name is rewritten, type is not.
 		{"from.nl.", dns.TypeANY, dns.ClassINET, "to.nl.", dns.TypeANY, dns.ClassINET},
+		{"from.exact.nl.", dns.TypeA, dns.ClassINET, "to.nl.", dns.TypeA, dns.ClassINET},
+		{"prefix.nl.", dns.TypeA, dns.ClassINET, "to.nl.", dns.TypeA, dns.ClassINET},
+		{"to.suffix.", dns.TypeA, dns.ClassINET, "to.nl.", dns.TypeA, dns.ClassINET},
+		{"from.substring.nl.", dns.TypeA, dns.ClassINET, "to.nl.", dns.TypeA, dns.ClassINET},
+		{"from.regex.nl.", dns.TypeA, dns.ClassINET, "to.nl.", dns.TypeA, dns.ClassINET},
 		// name is not, type is, but class is, because class is the 2nd rule.
 		{"a.nl.", dns.TypeANY, dns.ClassCHAOS, "a.nl.", dns.TypeANY, dns.ClassINET},
 	}