plugin/rewrite: use request.Request and other cleanups (#1920)

This was done anyway, but only deep in the functions, just do this
everywhere; allows for shorter code and request.Request allows for
caching as well.

Cleanups, make it more Go like.
* remove unneeded switches
* remove testdir (why was this there??)
* simplify the logic
* remove unneeded variables
* put short functions on a single line
* fix documentation.
* spin off wire funcs in wire.go, make them functions.
Signed-off-by: Miek Gieben <miek@miek.nl>

plugin/rewrite/class.go

@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/coredns/coredns/request"
+
 	"github.com/miekg/dns"
 )
 
@@ -27,22 +29,18 @@ func newClassRule(nextAction string, args ...string) (Rule, error) {
 }
 
 // Rewrite rewrites the the current request.
-func (rule *classRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
+func (rule *classRule) Rewrite(state request.Request) Result {
 	if rule.fromClass > 0 && rule.toClass > 0 {
-		if r.Question[0].Qclass == rule.fromClass {
-			r.Question[0].Qclass = rule.toClass
+		if state.Req.Question[0].Qclass == rule.fromClass {
+			state.Req.Question[0].Qclass = rule.toClass
 			return RewriteDone
 		}
 	}
 	return RewriteIgnored
 }
 
-// Mode returns the processing mode
-func (rule *classRule) Mode() string {
-	return rule.NextAction
-}
+// Mode returns the processing mode.
+func (rule *classRule) Mode() string { return rule.NextAction }
 
 // GetResponseRule return a rule to rewrite the response with. Currently not implemented.
-func (rule *classRule) GetResponseRule() ResponseRule {
-	return ResponseRule{}
-}
+func (rule *classRule) GetResponseRule() ResponseRule { return ResponseRule{} }

plugin/rewrite/condition.go

@@ -10,7 +10,7 @@ import (
 	"github.com/miekg/dns"
 )
 
-// Operators
+// Operators that are defined.
 const (
 	Is         = "is"
 	Not        = "not"
@@ -22,13 +22,7 @@ const (
 	NotMatch   = "not_match"
 )
 
-func operatorError(operator string) error {
-	return fmt.Errorf("invalid operator %v", operator)
-}
-
-func newReplacer(r *dns.Msg) replacer.Replacer {
-	return replacer.New(r, nil, "")
-}
+func newReplacer(r *dns.Msg) replacer.Replacer { return replacer.New(r, nil, "") }
 
 // condition is a rewrite condition.
 type condition func(string, string) bool
@@ -44,53 +38,35 @@ var conditions = map[string]condition{
 	NotMatch:   notMatchFunc,
 }
 
-// isFunc is condition for Is operator.
-// It checks for equality.
-func isFunc(a, b string) bool {
-	return a == b
-}
+// isFunc is condition for Is operator. It checks for equality.
+func isFunc(a, b string) bool { return a == b }
 
-// notFunc is condition for Not operator.
-// It checks for inequality.
-func notFunc(a, b string) bool {
-	return a != b
-}
+// notFunc is condition for Not operator. It checks for inequality.
+func notFunc(a, b string) bool { return a != b }
 
-// hasFunc is condition for Has operator.
-// It checks if b is a substring of a.
-func hasFunc(a, b string) bool {
-	return strings.Contains(a, b)
-}
+// hasFunc is condition for Has operator. It checks if b is a substring of a.
+func hasFunc(a, b string) bool { return strings.Contains(a, b) }
 
-// notHasFunc is condition for NotHas operator.
-// It checks if b is not a substring of a.
-func notHasFunc(a, b string) bool {
-	return !strings.Contains(a, b)
-}
+// notHasFunc is condition for NotHas operator. It checks if b is not a substring of a.
+func notHasFunc(a, b string) bool { return !strings.Contains(a, b) }
 
-// startsWithFunc is condition for StartsWith operator.
-// It checks if b is a prefix of a.
-func startsWithFunc(a, b string) bool {
-	return strings.HasPrefix(a, b)
-}
+// startsWithFunc is condition for StartsWith operator. It checks if b is a prefix of a.
+func startsWithFunc(a, b string) bool { return strings.HasPrefix(a, b) }
 
-// endsWithFunc is condition for EndsWith operator.
-// It checks if b is a suffix of a.
+// endsWithFunc is condition for EndsWith operator. It checks if b is a suffix of a.
 func endsWithFunc(a, b string) bool {
 	// TODO(miek): IsSubDomain
 	return strings.HasSuffix(a, b)
 }
 
-// matchFunc is condition for Match operator.
-// It does regexp matching of a against pattern in b
+// matchFunc is condition for Match operator. It does regexp matching of a against pattern in b
 // and returns if they match.
 func matchFunc(a, b string) bool {
 	matched, _ := regexp.MatchString(b, a)
 	return matched
 }
 
-// notMatchFunc is condition for NotMatch operator.
-// It does regexp matching of a against pattern in b
+// notMatchFunc is condition for NotMatch operator. It does regexp matching of a against pattern in b
 // and returns if they do not match.
 func notMatchFunc(a, b string) bool {
 	matched, _ := regexp.MatchString(b, a)
@@ -122,7 +98,7 @@ func (i If) True(r *dns.Msg) bool {
 // NewIf creates a new If condition.
 func NewIf(a, operator, b string) (If, error) {
 	if _, ok := conditions[operator]; !ok {
-		return If{}, operatorError(operator)
+		return If{}, fmt.Errorf("invalid operator %v", operator)
 	}
 	return If{
 		A:        a,

plugin/rewrite/name.go

@@ -7,8 +7,7 @@ import (
 	"strings"
 
 	"github.com/coredns/coredns/plugin"
-
-	"github.com/miekg/dns"
+	"github.com/coredns/coredns/request"
 )
 
 type nameRule struct {
@@ -56,47 +55,47 @@ const (
 )
 
 // Rewrite rewrites the current request based upon exact match of the name
-// in the question section of the request
-func (rule *nameRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
-	if rule.From == r.Question[0].Name {
-		r.Question[0].Name = rule.To
+// in the question section of the request.
+func (rule *nameRule) Rewrite(state request.Request) Result {
+	if rule.From == state.Name() {
+		state.Req.Question[0].Name = rule.To
 		return RewriteDone
 	}
 	return RewriteIgnored
 }
 
-// Rewrite rewrites the current request when the name begins with the matching string
-func (rule *prefixNameRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
-	if strings.HasPrefix(r.Question[0].Name, rule.Prefix) {
-		r.Question[0].Name = rule.Replacement + strings.TrimLeft(r.Question[0].Name, rule.Prefix)
+// Rewrite rewrites the current request when the name begins with the matching string.
+func (rule *prefixNameRule) Rewrite(state request.Request) Result {
+	if strings.HasPrefix(state.Name(), rule.Prefix) {
+		state.Req.Question[0].Name = rule.Replacement + strings.TrimLeft(state.Name(), rule.Prefix)
 		return RewriteDone
 	}
 	return RewriteIgnored
 }
 
-// Rewrite rewrites the current request when the name ends with the matching string
-func (rule *suffixNameRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
-	if strings.HasSuffix(r.Question[0].Name, rule.Suffix) {
-		r.Question[0].Name = strings.TrimRight(r.Question[0].Name, rule.Suffix) + rule.Replacement
+// Rewrite rewrites the current request when the name ends with the matching string.
+func (rule *suffixNameRule) Rewrite(state request.Request) Result {
+	if strings.HasSuffix(state.Name(), rule.Suffix) {
+		state.Req.Question[0].Name = strings.TrimRight(state.Name(), rule.Suffix) + rule.Replacement
 		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 *substringNameRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
-	if strings.Contains(r.Question[0].Name, rule.Substring) {
-		r.Question[0].Name = strings.Replace(r.Question[0].Name, rule.Substring, rule.Replacement, -1)
+// name in the question section of the request.
+func (rule *substringNameRule) Rewrite(state request.Request) Result {
+	if strings.Contains(state.Name(), rule.Substring) {
+		state.Req.Question[0].Name = strings.Replace(state.Name(), rule.Substring, rule.Replacement, -1)
 		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 *regexNameRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
-	regexGroups := rule.Pattern.FindStringSubmatch(r.Question[0].Name)
+// section of the request matches a regular expression.
+func (rule *regexNameRule) Rewrite(state request.Request) Result {
+	regexGroups := rule.Pattern.FindStringSubmatch(state.Name())
 	if len(regexGroups) == 0 {
 		return RewriteIgnored
 	}
@@ -107,7 +106,7 @@ func (rule *regexNameRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
 			s = strings.Replace(s, groupIndexStr, groupValue, -1)
 		}
 	}
-	r.Question[0].Name = s
+	state.Req.Question[0].Name = s
 	return RewriteDone
 }
 
@@ -174,47 +173,23 @@ func newNameRule(nextAction string, args ...string) (Rule, error) {
 }
 
 // Mode returns the processing nextAction
-func (rule *nameRule) Mode() string {
-	return rule.NextAction
-}
-
-func (rule *prefixNameRule) Mode() string {
-	return rule.NextAction
-}
-
-func (rule *suffixNameRule) Mode() string {
-	return rule.NextAction
-}
-
-func (rule *substringNameRule) Mode() string {
-	return rule.NextAction
-}
-
-func (rule *regexNameRule) Mode() string {
-	return rule.NextAction
-}
+func (rule *nameRule) Mode() string          { return rule.NextAction }
+func (rule *prefixNameRule) Mode() string    { return rule.NextAction }
+func (rule *suffixNameRule) Mode() string    { return rule.NextAction }
+func (rule *substringNameRule) Mode() string { return rule.NextAction }
+func (rule *regexNameRule) Mode() string     { return rule.NextAction }
 
 // GetResponseRule return a rule to rewrite the response with. Currently not implemented.
-func (rule *nameRule) GetResponseRule() ResponseRule {
-	return ResponseRule{}
-}
+func (rule *nameRule) GetResponseRule() ResponseRule { return ResponseRule{} }
 
 // GetResponseRule return a rule to rewrite the response with. Currently not implemented.
-func (rule *prefixNameRule) GetResponseRule() ResponseRule {
-	return ResponseRule{}
-}
+func (rule *prefixNameRule) GetResponseRule() ResponseRule { return ResponseRule{} }
 
 // GetResponseRule return a rule to rewrite the response with. Currently not implemented.
-func (rule *suffixNameRule) GetResponseRule() ResponseRule {
-	return ResponseRule{}
-}
+func (rule *suffixNameRule) GetResponseRule() ResponseRule { return ResponseRule{} }
 
 // GetResponseRule return a rule to rewrite the response with. Currently not implemented.
-func (rule *substringNameRule) GetResponseRule() ResponseRule {
-	return ResponseRule{}
-}
+func (rule *substringNameRule) GetResponseRule() ResponseRule { return ResponseRule{} }
 
 // GetResponseRule return a rule to rewrite the response with.
-func (rule *regexNameRule) GetResponseRule() ResponseRule {
-	return rule.ResponseRule
-}
+func (rule *regexNameRule) GetResponseRule() ResponseRule { return rule.ResponseRule }

plugin/rewrite/reverter.go

@@ -33,8 +33,7 @@ func NewResponseReverter(w dns.ResponseWriter, r *dns.Msg) *ResponseReverter {
 	}
 }
 
-// WriteMsg records the status code and calls the
-// underlying ResponseWriter's WriteMsg method.
+// WriteMsg records the status code and calls the underlying ResponseWriter's WriteMsg method.
 func (r *ResponseReverter) WriteMsg(res *dns.Msg) error {
 	res.Question[0] = r.originalQuestion
 	if r.ResponseRewrite {

plugin/rewrite/rewrite.go

@@ -6,6 +6,7 @@ import (
 	"strings"
 
 	"github.com/coredns/coredns/plugin"
+	"github.com/coredns/coredns/request"
 
 	"github.com/miekg/dns"
 )
@@ -38,8 +39,10 @@ type Rewrite struct {
 // ServeDNS implements the plugin.Handler interface.
 func (rw Rewrite) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
 	wr := NewResponseReverter(w, r)
+	state := request.Request{W: w, Req: r}
+
 	for _, rule := range rw.Rules {
-		switch result := rule.Rewrite(w, r); result {
+		switch result := rule.Rewrite(state); result {
 		case RewriteDone:
 			respRule := rule.GetResponseRule()
 			if respRule.Active == true {
@@ -68,7 +71,7 @@ func (rw Rewrite) Name() string { return "rewrite" }
 // Rule describes a rewrite rule.
 type Rule interface {
 	// Rewrite rewrites the current request.
-	Rewrite(dns.ResponseWriter, *dns.Msg) Result
+	Rewrite(state request.Request) Result
 	// Mode returns the processing mode stop or continue.
 	Mode() string
 	// GetResponseRule returns the rule to rewrite response with, if any.

plugin/rewrite/rewrite_test.go

@@ -482,14 +482,14 @@ func TestRewriteEDNS0LocalVariable(t *testing.T) {
 		},
 		{
 			[]dns.EDNS0{},
-			[]string{"local", "set", "0xffee", "{server_ip}"},
-			[]dns.EDNS0{&dns.EDNS0_LOCAL{Code: 0xffee, Data: []byte{0x7F, 0x00, 0x00, 0x01}}},
+			[]string{"local", "set", "0xffee", "{server_port}"},
+			[]dns.EDNS0{&dns.EDNS0_LOCAL{Code: 0xffee, Data: []byte{0x00, 0x35}}},
 			true,
 		},
 		{
 			[]dns.EDNS0{},
-			[]string{"local", "set", "0xffee", "{server_port}"},
-			[]dns.EDNS0{&dns.EDNS0_LOCAL{Code: 0xffee, Data: []byte{0x00, 0x35}}},
+			[]string{"local", "set", "0xffee", "{server_ip}"},
+			[]dns.EDNS0{&dns.EDNS0_LOCAL{Code: 0xffee, Data: []byte{0x7F, 0x00, 0x00, 0x01}}},
 			true,
 		},
 	}
@@ -498,7 +498,6 @@ func TestRewriteEDNS0LocalVariable(t *testing.T) {
 	for i, tc := range tests {
 		m := new(dns.Msg)
 		m.SetQuestion("example.com.", dns.TypeA)
-		m.Question[0].Qclass = dns.ClassINET
 
 		r, err := newEdns0Rule("stop", tc.args...)
 		if err != nil {
@@ -620,7 +619,6 @@ func TestRewriteEDNS0Subnet(t *testing.T) {
 	for i, tc := range tests {
 		m := new(dns.Msg)
 		m.SetQuestion("example.com.", dns.TypeA)
-		m.Question[0].Qclass = dns.ClassINET
 
 		r, err := newEdns0Rule("stop", tc.args...)
 		if err != nil {

plugin/rewrite/testdata/testfile

-empty
\ No newline at end of file

plugin/rewrite/type.go

-// Package rewrite is plugin for rewriting requests internally to something different.
+// Package rewrite is a plugin for rewriting requests internally to something different.
 package rewrite
 
 import (
 	"fmt"
 	"strings"
 
+	"github.com/coredns/coredns/request"
+
 	"github.com/miekg/dns"
 )
 
@@ -28,22 +30,18 @@ func newTypeRule(nextAction string, args ...string) (Rule, error) {
 }
 
 // Rewrite rewrites the the current request.
-func (rule *typeRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
+func (rule *typeRule) Rewrite(state request.Request) Result {
 	if rule.fromType > 0 && rule.toType > 0 {
-		if r.Question[0].Qtype == rule.fromType {
-			r.Question[0].Qtype = rule.toType
+		if state.QType() == rule.fromType {
+			state.Req.Question[0].Qtype = rule.toType
 			return RewriteDone
 		}
 	}
 	return RewriteIgnored
 }
 
-// Mode returns the processing mode
-func (rule *typeRule) Mode() string {
-	return rule.nextAction
-}
+// Mode returns the processing mode.
+func (rule *typeRule) Mode() string { return rule.nextAction }
 
 // GetResponseRule return a rule to rewrite the response with. Currently not implemented.
-func (rule *typeRule) GetResponseRule() ResponseRule {
-	return ResponseRule{}
-}
+func (rule *typeRule) GetResponseRule() ResponseRule { return ResponseRule{} }

plugin/rewrite/wire.go

+package rewrite
+
+import (
+	"encoding/binary"
+	"fmt"
+	"net"
+	"strconv"
+)
+
+// ipToWire writes IP address to wire/binary format, 4 or 16 bytes depends on IPV4 or IPV6.
+func ipToWire(family int, ipAddr string) ([]byte, error) {
+
+	switch family {
+	case 1:
+		return net.ParseIP(ipAddr).To4(), nil
+	case 2:
+		return net.ParseIP(ipAddr).To16(), nil
+	}
+	return nil, fmt.Errorf("invalid IP address family (i.e. version) %d", family)
+}
+
+// uint16ToWire writes unit16 to wire/binary format
+func uint16ToWire(data uint16) []byte {
+	buf := make([]byte, 2)
+	binary.BigEndian.PutUint16(buf, uint16(data))
+	return buf
+}
+
+// portToWire writes port to wire/binary format, 2 bytes
+func portToWire(portStr string) ([]byte, error) {
+
+	port, err := strconv.ParseUint(portStr, 10, 16)
+	if err != nil {
+		return nil, err
+	}
+	return uint16ToWire(uint16(port)), nil
+}

request/request.go

@@ -59,6 +59,15 @@ func (r *Request) IP() string {
 	return r.ip
 }
 
+// LocalIP gets the (local) IP address of server handling the request.
+func (r *Request) LocalIP() string {
+	ip, _, err := net.SplitHostPort(r.W.LocalAddr().String())
+	if err != nil {
+		return r.W.LocalAddr().String()
+	}
+	return ip
+}
+
 // Port gets the (remote) port of the client making the request.
 func (r *Request) Port() string {
 	if r.port != "" {