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

Resolves: #1307

1 files changed, 11 insertions, 8 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
 }