plugin/rewrite: fix flow control logic for all rule types (#1308)

Resolves: #1307

4 files changed, 34 insertions, 20 deletions

diff --git a/plugin/rewrite/class.go b/plugin/rewrite/class.go
index 8befdf8c2..ca86430d2 100644
--- a/plugin/rewrite/class.go
+++ b/plugin/rewrite/class.go
@@ -8,19 +8,22 @@ import (
 )
 
 type classRule struct {
-	fromClass, toClass uint16
+	fromClass  uint16
+	toClass    uint16
+	NextAction string
 }
 
-func newClassRule(fromS, toS string) (Rule, error) {
+// newClassRule creates a class matching rule
+func newClassRule(nextAction string, args ...string) (Rule, error) {
 	var from, to uint16
 	var ok bool
-	if from, ok = dns.StringToClass[strings.ToUpper(fromS)]; !ok {
-		return nil, fmt.Errorf("invalid class %q", strings.ToUpper(fromS))
+	if from, ok = dns.StringToClass[strings.ToUpper(args[0])]; !ok {
+		return nil, fmt.Errorf("invalid class %q", strings.ToUpper(args[0]))
 	}
-	if to, ok = dns.StringToClass[strings.ToUpper(toS)]; !ok {
-		return nil, fmt.Errorf("invalid class %q", strings.ToUpper(toS))
+	if to, ok = dns.StringToClass[strings.ToUpper(args[1])]; !ok {
+		return nil, fmt.Errorf("invalid class %q", strings.ToUpper(args[1]))
 	}
-	return &classRule{fromClass: from, toClass: to}, nil
+	return &classRule{from, to, nextAction}, nil
 }
 
 // Rewrite rewrites the the current request.
@@ -36,5 +39,5 @@ func (rule *classRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
 
 // Mode returns the processing mode
 func (rule *classRule) Mode() string {
-	return Stop
+	return rule.NextAction
 }
diff --git a/plugin/rewrite/rewrite.go b/plugin/rewrite/rewrite.go
index 4924ea6e5..4e31cb7d5 100644
--- a/plugin/rewrite/rewrite.go
+++ b/plugin/rewrite/rewrite.go
@@ -107,9 +107,9 @@ func newRule(args ...string) (Rule, error) {
 	case "name":
 		return newNameRule(mode, args[startArg:]...)
 	case "class":
-		return newClassRule(args[startArg], args[startArg+1])
+		return newClassRule(mode, args[startArg:]...)
 	case "type":
-		return newTypeRule(args[startArg], args[startArg+1])
+		return newTypeRule(mode, args[startArg:]...)
 	case "edns0":
 		return newEdns0Rule(mode, args[startArg:]...)
 	default:
diff --git a/plugin/rewrite/rewrite_test.go b/plugin/rewrite/rewrite_test.go
index 6314e5407..e2770da0e 100644
--- a/plugin/rewrite/rewrite_test.go
+++ b/plugin/rewrite/rewrite_test.go
@@ -162,9 +162,15 @@ func TestRewrite(t *testing.T) {
 	rules = append(rules, r)
 	r, _ = newNameRule("stop", "regex", "(f.*m)\\.regex\\.(nl)", "to.{2}")
 	rules = append(rules, r)
-	r, _ = newClassRule("CH", "IN")
+	r, _ = newNameRule("continue", "regex", "consul\\.(rocks)", "core.dns.{1}")
 	rules = append(rules, r)
-	r, _ = newTypeRule("ANY", "HINFO")
+	r, _ = newNameRule("stop", "core.dns.rocks", "to.nl.")
+	rules = append(rules, r)
+	r, _ = newClassRule("continue", "HS", "CH")
+	rules = append(rules, r)
+	r, _ = newClassRule("stop", "CH", "IN")
+	rules = append(rules, r)
+	r, _ = newTypeRule("stop", "ANY", "HINFO")
 	rules = append(rules, r)
 
 	rw := Rewrite{
@@ -192,8 +198,11 @@ func TestRewrite(t *testing.T) {
 		{"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},
+		{"consul.rocks.", 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},
+		// 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},
 	}
 
 	ctx := context.TODO()
diff --git a/plugin/rewrite/type.go b/plugin/rewrite/type.go
index a02d08603..22a545675 100644
--- a/plugin/rewrite/type.go
+++ b/plugin/rewrite/type.go
@@ -10,19 +10,21 @@ import (
 
 // typeRule is a type rewrite rule.
 type typeRule struct {
-	fromType, toType uint16
+	fromType   uint16
+	toType     uint16
+	nextAction string
 }
 
-func newTypeRule(fromS, toS string) (Rule, error) {
+func newTypeRule(nextAction string, args ...string) (Rule, error) {
 	var from, to uint16
 	var ok bool
-	if from, ok = dns.StringToType[strings.ToUpper(fromS)]; !ok {
-		return nil, fmt.Errorf("invalid type %q", strings.ToUpper(fromS))
+	if from, ok = dns.StringToType[strings.ToUpper(args[0])]; !ok {
+		return nil, fmt.Errorf("invalid type %q", strings.ToUpper(args[0]))
 	}
-	if to, ok = dns.StringToType[strings.ToUpper(toS)]; !ok {
-		return nil, fmt.Errorf("invalid type %q", strings.ToUpper(toS))
+	if to, ok = dns.StringToType[strings.ToUpper(args[1])]; !ok {
+		return nil, fmt.Errorf("invalid type %q", strings.ToUpper(args[1]))
 	}
-	return &typeRule{fromType: from, toType: to}, nil
+	return &typeRule{from, to, nextAction}, nil
 }
 
 // Rewrite rewrites the the current request.
@@ -38,5 +40,5 @@ func (rule *typeRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
 
 // Mode returns the processing mode
 func (rule *typeRule) Mode() string {
-	return Stop
+	return rule.nextAction
 }