1 files changed, 13 insertions, 9 deletions
diff --git a/middleware/rewrite/class.go b/middleware/rewrite/class.go
index ba383df81..a2bd00d0f 100644
--- a/middleware/rewrite/class.go
+++ b/middleware/rewrite/class.go
@@ -1,26 +1,30 @@
-// Package rewrite is middleware for rewriting requests internally to something different.
 package rewrite
 
 import (
+	"fmt"
 	"strings"
 
 	"github.com/miekg/dns"
 )
 
-// ClassRule is a class rewrite rule.
-type ClassRule struct {
+type classRule struct {
 	fromClass, toClass uint16
 }
 
-// New initializes rule.
-func (rule ClassRule) New(args ...string) Rule {
-	from, to := args[0], strings.Join(args[1:], " ")
-	return &ClassRule{dns.StringToClass[from], dns.StringToClass[to]}
-
+func newClassRule(fromS, toS 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 to, ok = dns.StringToClass[strings.ToUpper(toS)]; !ok {
+		return nil, fmt.Errorf("invalid class %q", strings.ToUpper(toS))
+	}
+	return &classRule{fromClass: from, toClass: to}, nil
 }
 
 // Rewrite rewrites the the current request.
-func (rule ClassRule) Rewrite(r *dns.Msg) Result {
+func (rule *classRule) Rewrite(r *dns.Msg) Result {
 	if rule.fromClass > 0 && rule.toClass > 0 {
 		if r.Question[0].Qclass == rule.fromClass {
 			r.Question[0].Qclass = rule.toClass