plugin/rewrite: add handling of TTL field rewrites (#2048)

Resolves: #1981

Signed-off-by: Paul Greenberg <greenpau@outlook.com>

6 files changed, 453 insertions, 25 deletions

diff --git a/plugin/rewrite/README.md b/plugin/rewrite/README.md
index b432d0ca4..c276fb9b0 100644
--- a/plugin/rewrite/README.md
+++ b/plugin/rewrite/README.md
@@ -13,7 +13,7 @@ Rewrites are invisible to the client. There are simple rewrites (fast) and compl
 
 A simplified/easy to digest syntax for *rewrite* is...
 ~~~
-rewrite [continue|stop] FIELD FROM TO
+rewrite [continue|stop] FIELD [FROM TO|FROM TTL]
 ~~~
 
 * **FIELD** indicates what part of the request/response is being re-written.
@@ -25,9 +25,11 @@ e.g., to rewrite ANY queries to HINFO, use `rewrite type ANY HINFO`.
      name, e.g., `rewrite name example.net example.org`. Other match types are supported, see the **Name Field Rewrites** section below.
    * `answer name` - the query name in the _response_ is rewritten.  This option has special restrictions and requirements, in particular it must always combined with a `name` rewrite.  See below in the **Response Rewrites** section.
    *  `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.
 
-* **FROM** is the name or type to match
+* **FROM** is the name (exact, suffix, prefix, substring, or regex) or type to match
 * **TO** is the destination name or type to rewrite to
+* **TTL** is the number of seconds to set the TTL value to
 
 If you specify multiple rules and an incoming query matches on multiple rules, the rewrite
 will behave as following
@@ -177,6 +179,32 @@ follows:
 rewrite [continue|stop] name regex STRING STRING answer name STRING STRING
 ```
 
+### TTL Field Rewrites
+
+At times, the need for rewriting TTL value could arise. For example, a DNS server
+may prevent caching by setting TTL as low as zero (`0`). An administrator
+may want to increase the TTL to prevent caching, e.g. to 15 seconds.
+
+In the below example, the TTL in the answers for `coredns.rocks` domain are
+being set to `15`:
+
+```
+    rewrite continue {
+        ttl regex (.*)\.coredns\.rocks 15
+    }
+```
+
+By the same token, an administrator may use this feature to force caching by
+setting TTL value really low.
+
+
+The syntax for the TTL rewrite rule is as follows. The meaning of
+`exact|prefix|suffix|substring|regex` is the same as with the name rewrite rules.
+
+```
+rewrite [continue|stop] ttl [exact|prefix|suffix|substring|regex] STRING SECONDS
+```
+
 ## EDNS0 Options
 
 Using FIELD edns0, you can set, append, or replace specific EDNS0 options on the request.
diff --git a/plugin/rewrite/name.go b/plugin/rewrite/name.go
index 23da0b0b5..7c3371b8f 100644
--- a/plugin/rewrite/name.go
+++ b/plugin/rewrite/name.go
@@ -133,7 +133,7 @@ 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(), ResponseRule{}}, nil
+			return &regexNameRule{nextAction, regexPattern, plugin.Name(args[2]).Normalize(), ResponseRule{Type: "name"}}, nil
 		default:
 			return nil, fmt.Errorf("A name rule supports only exact, prefix, suffix, substring, and regex name matching")
 		}
@@ -162,6 +162,7 @@ func newNameRule(nextAction string, args ...string) (Rule, error) {
 				plugin.Name(args[2]).Normalize(),
 				ResponseRule{
 					Active:      true,
+					Type:        "name",
 					Pattern:     responseRegexPattern,
 					Replacement: plugin.Name(args[6]).Normalize(),
 				},
diff --git a/plugin/rewrite/reverter.go b/plugin/rewrite/reverter.go
index 63e38708f..570b7d39e 100644
--- a/plugin/rewrite/reverter.go
+++ b/plugin/rewrite/reverter.go
@@ -1,18 +1,19 @@
 package rewrite
 
 import (
+	"github.com/miekg/dns"
 	"regexp"
 	"strconv"
 	"strings"
-
-	"github.com/miekg/dns"
 )
 
 // ResponseRule contains a rule to rewrite a response with.
 type ResponseRule struct {
 	Active      bool
+	Type        string
 	Pattern     *regexp.Regexp
 	Replacement string
+	Ttl         uint32
 }
 
 // ResponseReverter reverses the operations done on the question section of a packet.
@@ -38,22 +39,40 @@ func (r *ResponseReverter) WriteMsg(res *dns.Msg) error {
 	res.Question[0] = r.originalQuestion
 	if r.ResponseRewrite {
 		for _, rr := range res.Answer {
-			name := rr.Header().Name
+			var isNameRewritten bool = false
+			var isTtlRewritten bool = false
+			var name string = rr.Header().Name
+			var ttl uint32 = rr.Header().Ttl
 			for _, rule := range r.ResponseRules {
-				regexGroups := rule.Pattern.FindStringSubmatch(name)
-				if len(regexGroups) == 0 {
-					continue
+				if rule.Type == "" {
+					rule.Type = "name"
 				}
-				s := rule.Replacement
-				for groupIndex, groupValue := range regexGroups {
-					groupIndexStr := "{" + strconv.Itoa(groupIndex) + "}"
-					if strings.Contains(s, groupIndexStr) {
-						s = strings.Replace(s, groupIndexStr, groupValue, -1)
+				switch rule.Type {
+				case "name":
+					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
+					isNameRewritten = true
+				case "ttl":
+					ttl = rule.Ttl
+					isTtlRewritten = true
 				}
-				name = s
 			}
-			rr.Header().Name = name
+			if isNameRewritten == true {
+				rr.Header().Name = name
+			}
+			if isTtlRewritten == true {
+				rr.Header().Ttl = ttl
+			}
 		}
 	}
 	return r.ResponseWriter.WriteMsg(res)
diff --git a/plugin/rewrite/rewrite.go b/plugin/rewrite/rewrite.go
index 643f2d7c9..64283df92 100644
--- a/plugin/rewrite/rewrite.go
+++ b/plugin/rewrite/rewrite.go
@@ -50,7 +50,6 @@ func (rw Rewrite) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg
 				state.Req.Question[0] = wr.originalQuestion
 				return dns.RcodeServerFailure, fmt.Errorf("invalid name after rewrite: %s", x)
 			}
-
 			respRule := rule.GetResponseRule()
 			if respRule.Active == true {
 				wr.ResponseRewrite = true
@@ -111,23 +110,25 @@ func newRule(args ...string) (Rule, error) {
 		startArg = 1
 	}
 
-	if ruleType == "answer" {
-		return nil, fmt.Errorf("response rewrites must begin with a name rule")
-	}
-
-	if ruleType != "edns0" && ruleType != "name" && expectNumArgs != 3 {
-		return nil, fmt.Errorf("%s rules must have exactly two arguments", ruleType)
-	}
-
 	switch ruleType {
+	case "answer":
+		return nil, fmt.Errorf("response rewrites must begin with a name rule")
 	case "name":
 		return newNameRule(mode, args[startArg:]...)
 	case "class":
+		if expectNumArgs != 3 {
+			return nil, fmt.Errorf("%s rules must have exactly two arguments", ruleType)
+		}
 		return newClassRule(mode, args[startArg:]...)
 	case "type":
+		if expectNumArgs != 3 {
+			return nil, fmt.Errorf("%s rules must have exactly two arguments", ruleType)
+		}
 		return newTypeRule(mode, args[startArg:]...)
 	case "edns0":
 		return newEdns0Rule(mode, args[startArg:]...)
+	case "ttl":
+		return newTtlRule(mode, args[startArg:]...)
 	default:
 		return nil, fmt.Errorf("invalid rule type %q", args[0])
 	}
diff --git a/plugin/rewrite/ttl.go b/plugin/rewrite/ttl.go
new file mode 100644
index 000000000..1b06673ff
--- /dev/null
+++ b/plugin/rewrite/ttl.go
@@ -0,0 +1,224 @@
+package rewrite
+
+import (
+	"context"
+	"fmt"
+	"regexp"
+	"strconv"
+	"strings"
+
+	"github.com/coredns/coredns/plugin"
+	"github.com/coredns/coredns/request"
+	//"github.com/miekg/dns"
+)
+
+type exactTtlRule struct {
+	NextAction string
+	From       string
+	ResponseRule
+}
+
+type prefixTtlRule struct {
+	NextAction string
+	Prefix     string
+	ResponseRule
+}
+
+type suffixTtlRule struct {
+	NextAction string
+	Suffix     string
+	ResponseRule
+}
+
+type substringTtlRule struct {
+	NextAction string
+	Substring  string
+	ResponseRule
+}
+
+type regexTtlRule struct {
+	NextAction string
+	Pattern    *regexp.Regexp
+	ResponseRule
+}
+
+// Rewrite rewrites the current request based upon exact match of the name
+// in the question section of the request.
+func (rule *exactTtlRule) Rewrite(ctx context.Context, state request.Request) Result {
+	if rule.From == state.Name() {
+		return RewriteDone
+	}
+	return RewriteIgnored
+}
+
+// Rewrite rewrites the current request when the name begins with the matching string.
+func (rule *prefixTtlRule) Rewrite(ctx context.Context, state request.Request) Result {
+	if strings.HasPrefix(state.Name(), rule.Prefix) {
+		return RewriteDone
+	}
+	return RewriteIgnored
+}
+
+// Rewrite rewrites the current request when the name ends with the matching string.
+func (rule *suffixTtlRule) Rewrite(ctx context.Context, state request.Request) Result {
+	if strings.HasSuffix(state.Name(), rule.Suffix) {
+		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 *substringTtlRule) Rewrite(ctx context.Context, state request.Request) Result {
+	if strings.Contains(state.Name(), rule.Substring) {
+		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 *regexTtlRule) Rewrite(ctx context.Context, state request.Request) Result {
+	regexGroups := rule.Pattern.FindStringSubmatch(state.Name())
+	if len(regexGroups) == 0 {
+		return RewriteIgnored
+	}
+	return RewriteDone
+}
+
+// newTtlRule creates a name matching rule based on exact, partial, or regex match
+func newTtlRule(nextAction string, args ...string) (Rule, error) {
+	if len(args) < 2 {
+		return nil, fmt.Errorf("too few (%d) arguments for a ttl rule", len(args))
+	}
+	var s string
+	if len(args) == 2 {
+		s = args[1]
+	}
+	if len(args) == 3 {
+		s = args[2]
+	}
+	ttl, valid := isValidTtl(s)
+	if valid == false {
+		return nil, fmt.Errorf("invalid TTL '%s' for a ttl rule", s)
+	}
+	if len(args) == 3 {
+		switch strings.ToLower(args[0]) {
+		case ExactMatch:
+			return &exactTtlRule{
+				nextAction,
+				plugin.Name(args[1]).Normalize(),
+				ResponseRule{
+					Active: true,
+					Type:   "ttl",
+					Ttl:    ttl,
+				},
+			}, nil
+		case PrefixMatch:
+			return &prefixTtlRule{
+				nextAction,
+				plugin.Name(args[1]).Normalize(),
+				ResponseRule{
+					Active: true,
+					Type:   "ttl",
+					Ttl:    ttl,
+				},
+			}, nil
+		case SuffixMatch:
+			return &suffixTtlRule{
+				nextAction,
+				plugin.Name(args[1]).Normalize(),
+				ResponseRule{
+					Active: true,
+					Type:   "ttl",
+					Ttl:    ttl,
+				},
+			}, nil
+		case SubstringMatch:
+			return &substringTtlRule{
+				nextAction,
+				plugin.Name(args[1]).Normalize(),
+				ResponseRule{
+					Active: true,
+					Type:   "ttl",
+					Ttl:    ttl,
+				},
+			}, nil
+		case RegexMatch:
+			regexPattern, err := regexp.Compile(args[1])
+			if err != nil {
+				return nil, fmt.Errorf("Invalid regex pattern in a ttl rule: %s", args[1])
+			}
+			return &regexTtlRule{
+				nextAction,
+				regexPattern,
+				ResponseRule{
+					Active: true,
+					Type:   "ttl",
+					Ttl:    ttl,
+				},
+			}, nil
+		default:
+			return nil, fmt.Errorf("A ttl rule supports only exact, prefix, suffix, substring, and regex name matching")
+		}
+	}
+	if len(args) > 3 {
+		return nil, fmt.Errorf("many few arguments for a ttl rule")
+	}
+	return &exactTtlRule{
+		nextAction,
+		plugin.Name(args[0]).Normalize(),
+		ResponseRule{
+			Active: true,
+			Type:   "ttl",
+			Ttl:    ttl,
+		},
+	}, nil
+}
+
+// Mode returns the processing nextAction
+func (rule *exactTtlRule) Mode() string     { return rule.NextAction }
+func (rule *prefixTtlRule) Mode() string    { return rule.NextAction }
+func (rule *suffixTtlRule) Mode() string    { return rule.NextAction }
+func (rule *substringTtlRule) Mode() string { return rule.NextAction }
+func (rule *regexTtlRule) Mode() string     { return rule.NextAction }
+
+// GetResponseRule return a rule to rewrite the response with. Currently not implemented.
+func (rule *exactTtlRule) GetResponseRule() ResponseRule {
+	return rule.ResponseRule
+}
+
+// GetResponseRule return a rule to rewrite the response with. Currently not implemented.
+func (rule *prefixTtlRule) GetResponseRule() ResponseRule {
+	return rule.ResponseRule
+}
+
+// GetResponseRule return a rule to rewrite the response with. Currently not implemented.
+func (rule *suffixTtlRule) GetResponseRule() ResponseRule {
+	return rule.ResponseRule
+}
+
+// GetResponseRule return a rule to rewrite the response with. Currently not implemented.
+func (rule *substringTtlRule) GetResponseRule() ResponseRule {
+	return rule.ResponseRule
+}
+
+// GetResponseRule return a rule to rewrite the response with.
+func (rule *regexTtlRule) GetResponseRule() ResponseRule {
+	return rule.ResponseRule
+}
+
+// validTtl returns true if v is valid TTL value.
+func isValidTtl(v string) (uint32, bool) {
+	i, err := strconv.Atoi(v)
+	if err != nil {
+		return uint32(0), false
+	}
+	if i > 2147483647 {
+		return uint32(0), false
+	}
+	if i < 0 {
+		return uint32(0), false
+	}
+	return uint32(i), true
+}
diff --git a/plugin/rewrite/ttl_test.go b/plugin/rewrite/ttl_test.go
new file mode 100644
index 000000000..5857f776b
--- /dev/null
+++ b/plugin/rewrite/ttl_test.go
@@ -0,0 +1,155 @@
+package rewrite
+
+import (
+	"context"
+	"github.com/coredns/coredns/plugin"
+	"github.com/coredns/coredns/plugin/pkg/dnstest"
+	"github.com/coredns/coredns/plugin/test"
+	"reflect"
+	"testing"
+
+	"github.com/miekg/dns"
+)
+
+func TestNewTtlRule(t *testing.T) {
+	tests := []struct {
+		next         string
+		args         []string
+		expectedFail bool
+	}{
+		{"stop", []string{"srv1.coredns.rocks", "10"}, false},
+		{"stop", []string{"exact", "srv1.coredns.rocks", "15"}, false},
+		{"stop", []string{"prefix", "coredns.rocks", "20"}, false},
+		{"stop", []string{"suffix", "srv1", "25"}, false},
+		{"stop", []string{"substring", "coredns", "30"}, false},
+		{"stop", []string{"regex", `(srv1)\.(coredns)\.(rocks)`, "35"}, false},
+		{"continue", []string{"srv1.coredns.rocks", "10"}, false},
+		{"continue", []string{"exact", "srv1.coredns.rocks", "15"}, false},
+		{"continue", []string{"prefix", "coredns.rocks", "20"}, false},
+		{"continue", []string{"suffix", "srv1", "25"}, false},
+		{"continue", []string{"substring", "coredns", "30"}, false},
+		{"continue", []string{"regex", `(srv1)\.(coredns)\.(rocks)`, "35"}, false},
+		{"stop", []string{"srv1.coredns.rocks", "12345678901234567890"}, true},
+		{"stop", []string{"srv1.coredns.rocks", "coredns.rocks"}, true},
+		{"stop", []string{"srv1.coredns.rocks", "-1"}, true},
+	}
+	for i, tc := range tests {
+		failed := false
+		rule, err := newTtlRule(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
+		}
+		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, "ttl"}, 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)
+	}
+}
+
+func TestTtlRewrite(t *testing.T) {
+	rules := []Rule{}
+	ruleset := []struct {
+		args         []string
+		expectedType reflect.Type
+	}{
+		{[]string{"stop", "ttl", "srv1.coredns.rocks", "1"}, reflect.TypeOf(&exactTtlRule{})},
+		{[]string{"stop", "ttl", "exact", "srv15.coredns.rocks", "15"}, reflect.TypeOf(&exactTtlRule{})},
+		{[]string{"stop", "ttl", "prefix", "srv30", "30"}, reflect.TypeOf(&prefixTtlRule{})},
+		{[]string{"stop", "ttl", "suffix", "45.coredns.rocks", "45"}, reflect.TypeOf(&suffixTtlRule{})},
+		{[]string{"stop", "ttl", "substring", "rv50", "50"}, reflect.TypeOf(&substringTtlRule{})},
+		{[]string{"stop", "ttl", "regex", `(srv10)\.(coredns)\.(rocks)`, "10"}, reflect.TypeOf(&regexTtlRule{})},
+		{[]string{"stop", "ttl", "regex", `(srv20)\.(coredns)\.(rocks)`, "20"}, reflect.TypeOf(&regexTtlRule{})},
+	}
+	for i, r := range ruleset {
+		rule, err := newRule(r.args...)
+		if err != nil {
+			t.Fatalf("Rule %d: FAIL, %s: %s", i, r.args, err)
+			continue
+		}
+		if reflect.TypeOf(rule) != r.expectedType {
+			t.Fatalf("Rule %d: FAIL, %s: rule type mismatch, expected %q, but got %q", i, r.args, r.expectedType, rule)
+		}
+		rules = append(rules, rule)
+	}
+	doTtlTests(rules, t)
+}
+
+func doTtlTests(rules []Rule, t *testing.T) {
+	tests := []struct {
+		from      string
+		fromType  uint16
+		answer    []dns.RR
+		ttl       uint32
+		noRewrite bool
+	}{
+		{"srv1.coredns.rocks.", dns.TypeA, []dns.RR{test.A("srv1.coredns.rocks.  5   IN  A  10.0.0.1")}, 1, false},
+		{"srv15.coredns.rocks.", dns.TypeA, []dns.RR{test.A("srv15.coredns.rocks.  5   IN  A  10.0.0.15")}, 15, false},
+		{"srv30.coredns.rocks.", dns.TypeA, []dns.RR{test.A("srv30.coredns.rocks.  5   IN  A  10.0.0.30")}, 30, false},
+		{"srv45.coredns.rocks.", dns.TypeA, []dns.RR{test.A("srv45.coredns.rocks.  5   IN  A  10.0.0.45")}, 45, false},
+		{"srv50.coredns.rocks.", dns.TypeA, []dns.RR{test.A("srv50.coredns.rocks.  5   IN  A  10.0.0.50")}, 50, false},
+		{"srv10.coredns.rocks.", dns.TypeA, []dns.RR{test.A("srv10.coredns.rocks.  5   IN  A  10.0.0.10")}, 10, false},
+		{"xmpp.coredns.rocks.", dns.TypeSRV, []dns.RR{test.SRV("xmpp.coredns.rocks.  5  IN  SRV 0 100 100 srvxmpp.coredns.rocks.")}, 5, true},
+		{"srv15.coredns.rocks.", dns.TypeHINFO, []dns.RR{test.HINFO("srv15.coredns.rocks.  5  HINFO INTEL-64 \"RHEL 7.5\"")}, 15, false},
+		{"srv20.coredns.rocks.", dns.TypeA, []dns.RR{
+			test.A("srv20.coredns.rocks.  5   IN  A  10.0.0.22"),
+			test.A("srv20.coredns.rocks.  5   IN  A  10.0.0.23"),
+		}, 20, false},
+	}
+	ctx := context.TODO()
+	for i, tc := range tests {
+		failed := false
+		m := new(dns.Msg)
+		m.SetQuestion(tc.from, tc.fromType)
+		m.Question[0].Qclass = dns.ClassINET
+		m.Answer = tc.answer
+		rw := Rewrite{
+			Next:     plugin.HandlerFunc(msgPrinter),
+			Rules:    rules,
+			noRevert: false,
+		}
+		rec := dnstest.NewRecorder(&test.ResponseWriter{})
+		rw.ServeDNS(ctx, rec, m)
+		resp := rec.Msg
+		if len(resp.Answer) == 0 {
+			t.Errorf("Test %d: FAIL %s (%d) Expected valid response but received %q", i, tc.from, tc.fromType, resp)
+			failed = true
+			continue
+		}
+		for _, a := range resp.Answer {
+			if a.Header().Ttl != tc.ttl {
+				t.Errorf("Test %d: FAIL %s (%d) Expected TTL to be %d but was %d", i, tc.from, tc.fromType, tc.ttl, a.Header().Ttl)
+				failed = true
+				break
+			}
+		}
+		if !failed {
+			if tc.noRewrite {
+				t.Logf("Test %d: PASS %s (%d) worked as expected, no rewrite for ttl %d", i, tc.from, tc.fromType, tc.ttl)
+			} else {
+				t.Logf("Test %d: PASS %s (%d) worked as expected, rewrote ttl to %d", i, tc.from, tc.fromType, tc.ttl)
+			}
+		}
+	}
+}