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

Resolves: #1981

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

1 files changed, 32 insertions, 13 deletions

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)