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>
2025-11-19 18:32:16 -05:00 · 2018-07-02 15:39:50 +01:00
parent 1abecf99d9
commit 6dd2cf8c4b
12 changed files with 200 additions and 311 deletions
--- a/plugin/rewrite/edns0.go
+++ b/plugin/rewrite/edns0.go
@@ -2,7 +2,6 @@
 package rewrite

 import (
-	"encoding/binary"
 	"encoding/hex"
 	"fmt"
 	"net"
@@ -14,7 +13,7 @@ import (
 	"github.com/miekg/dns"
 )

-// edns0LocalRule is a rewrite rule for EDNS0_LOCAL options
+// edns0LocalRule is a rewrite rule for EDNS0_LOCAL options.
 type edns0LocalRule struct {
 	mode   string
 	action string
@@ -22,7 +21,7 @@ type edns0LocalRule struct {
 	data   []byte
 }

-// edns0VariableRule is a rewrite rule for EDNS0_LOCAL options with variable
+// edns0VariableRule is a rewrite rule for EDNS0_LOCAL options with variable.
 type edns0VariableRule struct {
 	mode     string
 	action   string
@@ -30,13 +29,13 @@ type edns0VariableRule struct {
 	variable string
 }

-// ends0NsidRule is a rewrite rule for EDNS0_NSID options
+// ends0NsidRule is a rewrite rule for EDNS0_NSID options.
 type edns0NsidRule struct {
 	mode   string
 	action string
 }

-// setupEdns0Opt will retrieve the EDNS0 OPT or create it if it does not exist
+// setupEdns0Opt will retrieve the EDNS0 OPT or create it if it does not exist.
 func setupEdns0Opt(r *dns.Msg) *dns.OPT {
 	o := r.IsEdns0()
 	if o == nil {
@@ -47,82 +46,62 @@ func setupEdns0Opt(r *dns.Msg) *dns.OPT {
 }

 // Rewrite will alter the request EDNS0 NSID option
-func (rule *edns0NsidRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
-	result := RewriteIgnored
-	o := setupEdns0Opt(r)
-	found := false
-Option:
+func (rule *edns0NsidRule) Rewrite(state request.Request) Result {
+	o := setupEdns0Opt(state.Req)
+
 	for _, s := range o.Option {
-		switch e := s.(type) {
-		case *dns.EDNS0_NSID:
+		if e, ok := s.(*dns.EDNS0_NSID); ok {
 			if rule.action == Replace || rule.action == Set {
 				e.Nsid = "" // make sure it is empty for request
-				result = RewriteDone
+				return RewriteDone
 			}
-			found = true
-			break Option
 		}
 	}

 	// add option if not found
-	if !found && (rule.action == Append || rule.action == Set) {
+	if rule.action == Append || rule.action == Set {
 		o.Option = append(o.Option, &dns.EDNS0_NSID{Code: dns.EDNS0NSID, Nsid: ""})
-		result = RewriteDone
+		return RewriteDone
 	}

-	return result
+	return RewriteIgnored
 }

-// Mode returns the processing mode
-func (rule *edns0NsidRule) Mode() string {
-	return rule.mode
-}
+// Mode returns the processing mode.
+func (rule *edns0NsidRule) Mode() string { return rule.mode }

 // GetResponseRule return a rule to rewrite the response with. Currently not implemented.
-func (rule *edns0NsidRule) GetResponseRule() ResponseRule {
-	return ResponseRule{}
-}
+func (rule *edns0NsidRule) GetResponseRule() ResponseRule { return ResponseRule{} }
+
+// Rewrite will alter the request EDNS0 local options.
+func (rule *edns0LocalRule) Rewrite(state request.Request) Result {
+	o := setupEdns0Opt(state.Req)

-// Rewrite will alter the request EDNS0 local options
-func (rule *edns0LocalRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
-	result := RewriteIgnored
-	o := setupEdns0Opt(r)
-	found := false
 	for _, s := range o.Option {
-		switch e := s.(type) {
-		case *dns.EDNS0_LOCAL:
+		if e, ok := s.(*dns.EDNS0_LOCAL); ok {
 			if rule.code == e.Code {
 				if rule.action == Replace || rule.action == Set {
 					e.Data = rule.data
-					result = RewriteDone
+					return RewriteDone
 				}
-				found = true
-				break
 			}
 		}
 	}

 	// add option if not found
-	if !found && (rule.action == Append || rule.action == Set) {
-		var opt dns.EDNS0_LOCAL
-		opt.Code = rule.code
-		opt.Data = rule.data
-		o.Option = append(o.Option, &opt)
-		result = RewriteDone
+	if rule.action == Append || rule.action == Set {
+		o.Option = append(o.Option, &dns.EDNS0_LOCAL{Code: rule.code, Data: rule.data})
+		return RewriteDone
 	}

-	return result
+	return RewriteIgnored
 }

-// Mode returns the processing mode
-func (rule *edns0LocalRule) Mode() string {
-	return rule.mode
-}
+// Mode returns the processing mode.
+func (rule *edns0LocalRule) Mode() string { return rule.mode }

 // GetResponseRule return a rule to rewrite the response with. Currently not implemented.
-func (rule *edns0LocalRule) GetResponseRule() ResponseRule {
-	return ResponseRule{}
-}
+func (rule *edns0LocalRule) GetResponseRule() ResponseRule { return ResponseRule{} }

 // newEdns0Rule creates an EDNS0 rule of the appropriate type based on the args
 func newEdns0Rule(mode string, args ...string) (Rule, error) {
@@ -145,7 +124,7 @@ func newEdns0Rule(mode string, args ...string) (Rule, error) {
 		if len(args) != 4 {
 			return nil, fmt.Errorf("EDNS0 local rules require exactly three args")
 		}
-		//Check for variable option
+		// Check for variable option.
 		if strings.HasPrefix(args[3], "{") && strings.HasSuffix(args[3], "}") {
 			return newEdns0VariableRule(mode, action, args[2], args[3])
 		}
@@ -194,136 +173,69 @@ func newEdns0VariableRule(mode, action, code, variable string) (*edns0VariableRu
 	return &edns0VariableRule{mode: mode, action: action, code: uint16(c), variable: variable}, nil
 }

-// ipToWire writes IP address to wire/binary format, 4 or 16 bytes depends on IPV4 or IPV6.
-func (rule *edns0VariableRule) ipToWire(family int, ipAddr string) ([]byte, error) {
+// ruleData returns the data specified by the variable.
+func (rule *edns0VariableRule) ruleData(state request.Request) ([]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 (rule *edns0VariableRule) 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 (rule *edns0VariableRule) portToWire(portStr string) ([]byte, error) {
-
-	port, err := strconv.ParseUint(portStr, 10, 16)
-	if err != nil {
-		return nil, err
-	}
-	return rule.uint16ToWire(uint16(port)), nil
-}
-
-// Family returns the family of the transport, 1 for IPv4 and 2 for IPv6.
-func (rule *edns0VariableRule) family(ip net.Addr) int {
-	var a net.IP
-	if i, ok := ip.(*net.UDPAddr); ok {
-		a = i.IP
-	}
-	if i, ok := ip.(*net.TCPAddr); ok {
-		a = i.IP
-	}
-	if a.To4() != nil {
-		return 1
-	}
-	return 2
-}
-
-// ruleData returns the data specified by the variable
-func (rule *edns0VariableRule) ruleData(w dns.ResponseWriter, r *dns.Msg) ([]byte, error) {
-
-	req := request.Request{W: w, Req: r}
 	switch rule.variable {
 	case queryName:
-		//Query name is written as ascii string
-		return []byte(req.QName()), nil
+		return []byte(state.QName()), nil

 	case queryType:
-		return rule.uint16ToWire(req.QType()), nil
+		return uint16ToWire(state.QType()), nil

 	case clientIP:
-		return rule.ipToWire(req.Family(), req.IP())
-
-	case clientPort:
-		return rule.portToWire(req.Port())
-
-	case protocol:
-		// Proto is written as ascii string
-		return []byte(req.Proto()), nil
+		return ipToWire(state.Family(), state.IP())

 	case serverIP:
-		ip, _, err := net.SplitHostPort(w.LocalAddr().String())
-		if err != nil {
-			ip = w.RemoteAddr().String()
-		}
-		return rule.ipToWire(rule.family(w.RemoteAddr()), ip)
+		return ipToWire(state.Family(), state.LocalIP())
+
+	case clientPort:
+		return portToWire(state.Port())

 	case serverPort:
-		_, port, err := net.SplitHostPort(w.LocalAddr().String())
-		if err != nil {
-			port = "0"
-		}
-		return rule.portToWire(port)
+		return portToWire(state.LocalPort())
+
+	case protocol:
+		return []byte(state.Proto()), nil
 	}

 	return nil, fmt.Errorf("unable to extract data for variable %s", rule.variable)
 }

-// Rewrite will alter the request EDNS0 local options with specified variables
-func (rule *edns0VariableRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
-	result := RewriteIgnored
-
-	data, err := rule.ruleData(w, r)
+// Rewrite will alter the request EDNS0 local options with specified variables.
+func (rule *edns0VariableRule) Rewrite(state request.Request) Result {
+	data, err := rule.ruleData(state)
 	if err != nil || data == nil {
-		return result
+		return RewriteIgnored
 	}

-	o := setupEdns0Opt(r)
-	found := false
+	o := setupEdns0Opt(state.Req)
 	for _, s := range o.Option {
-		switch e := s.(type) {
-		case *dns.EDNS0_LOCAL:
+		if e, ok := s.(*dns.EDNS0_LOCAL); ok {
 			if rule.code == e.Code {
 				if rule.action == Replace || rule.action == Set {
 					e.Data = data
-					result = RewriteDone
+					return RewriteDone
 				}
-				found = true
-				break
+				return RewriteIgnored
 			}
 		}
 	}

 	// add option if not found
-	if !found && (rule.action == Append || rule.action == Set) {
-		var opt dns.EDNS0_LOCAL
-		opt.Code = rule.code
-		opt.Data = data
-		o.Option = append(o.Option, &opt)
-		result = RewriteDone
+	if rule.action == Append || rule.action == Set {
+		o.Option = append(o.Option, &dns.EDNS0_LOCAL{Code: rule.code, Data: data})
+		return RewriteDone
 	}

-	return result
+	return RewriteIgnored
 }

-// Mode returns the processing mode
-func (rule *edns0VariableRule) Mode() string {
-	return rule.mode
-}
+// Mode returns the processing mode.
+func (rule *edns0VariableRule) Mode() string { return rule.mode }

 // GetResponseRule return a rule to rewrite the response with. Currently not implemented.
-func (rule *edns0VariableRule) GetResponseRule() ResponseRule {
-	return ResponseRule{}
-}
+func (rule *edns0VariableRule) GetResponseRule() ResponseRule { return ResponseRule{} }

 func isValidVariable(variable string) bool {
 	switch variable {
@@ -353,8 +265,8 @@ func newEdns0SubnetRule(mode, action, v4BitMaskLen, v6BitMaskLen string) (*edns0
 	if err != nil {
 		return nil, err
 	}
-	// Validate V4 length
-	if v4Len > maxV4BitMaskLen {
+	// validate V4 length
+	if v4Len > net.IPv4len*8 {
 		return nil, fmt.Errorf("invalid IPv4 bit mask length %d", v4Len)
 	}

@@ -362,8 +274,8 @@ func newEdns0SubnetRule(mode, action, v4BitMaskLen, v6BitMaskLen string) (*edns0
 	if err != nil {
 		return nil, err
 	}
-	//Validate V6 length
-	if v6Len > maxV6BitMaskLen {
+	// validate V6 length
+	if v6Len > net.IPv6len*8 {
 		return nil, fmt.Errorf("invalid IPv6 bit mask length %d", v6Len)
 	}

@@ -372,10 +284,8 @@ func newEdns0SubnetRule(mode, action, v4BitMaskLen, v6BitMaskLen string) (*edns0
 }

 // fillEcsData sets the subnet data into the ecs option
-func (rule *edns0SubnetRule) fillEcsData(w dns.ResponseWriter, r *dns.Msg, ecs *dns.EDNS0_SUBNET) error {
-
-	req := request.Request{W: w, Req: r}
-	family := req.Family()
+func (rule *edns0SubnetRule) fillEcsData(state request.Request, ecs *dns.EDNS0_SUBNET) error {
+	family := state.Family()
 	if (family != 1) && (family != 2) {
 		return fmt.Errorf("unable to fill data for EDNS0 subnet due to invalid IP family")
 	}
@@ -383,7 +293,7 @@ func (rule *edns0SubnetRule) fillEcsData(w dns.ResponseWriter, r *dns.Msg, ecs *
 	ecs.Family = uint16(family)
 	ecs.SourceScope = 0

-	ipAddr := req.IP()
+	ipAddr := state.IP()
 	switch family {
 	case 1:
 		ipv4Mask := net.CIDRMask(int(rule.v4BitMaskLen), 32)
@@ -399,45 +309,38 @@ func (rule *edns0SubnetRule) fillEcsData(w dns.ResponseWriter, r *dns.Msg, ecs *
 	return nil
 }

-// Rewrite will alter the request EDNS0 subnet option
-func (rule *edns0SubnetRule) Rewrite(w dns.ResponseWriter, r *dns.Msg) Result {
-	result := RewriteIgnored
-	o := setupEdns0Opt(r)
-	found := false
+// Rewrite will alter the request EDNS0 subnet option.
+func (rule *edns0SubnetRule) Rewrite(state request.Request) Result {
+	o := setupEdns0Opt(state.Req)
+
 	for _, s := range o.Option {
-		switch e := s.(type) {
-		case *dns.EDNS0_SUBNET:
+		if e, ok := s.(*dns.EDNS0_SUBNET); ok {
 			if rule.action == Replace || rule.action == Set {
-				if rule.fillEcsData(w, r, e) == nil {
-					result = RewriteDone
+				if rule.fillEcsData(state, e) == nil {
+					return RewriteDone
 				}
 			}
-			found = true
-			break
+			return RewriteIgnored
 		}
 	}

 	// add option if not found
-	if !found && (rule.action == Append || rule.action == Set) {
-		opt := dns.EDNS0_SUBNET{Code: dns.EDNS0SUBNET}
-		if rule.fillEcsData(w, r, &opt) == nil {
-			o.Option = append(o.Option, &opt)
-			result = RewriteDone
+	if rule.action == Append || rule.action == Set {
+		opt := &dns.EDNS0_SUBNET{Code: dns.EDNS0SUBNET}
+		if rule.fillEcsData(state, opt) == nil {
+			o.Option = append(o.Option, opt)
+			return RewriteDone
 		}
 	}

-	return result
+	return RewriteIgnored
 }

 // Mode returns the processing mode
-func (rule *edns0SubnetRule) Mode() string {
-	return rule.mode
-}
+func (rule *edns0SubnetRule) Mode() string { return rule.mode }

 // GetResponseRule return a rule to rewrite the response with. Currently not implemented.
-func (rule *edns0SubnetRule) GetResponseRule() ResponseRule {
-	return ResponseRule{}
-}
+func (rule *edns0SubnetRule) GetResponseRule() ResponseRule { return ResponseRule{} }

 // These are all defined actions.
 const (
@@ -456,9 +359,3 @@ const (
 	serverIP   = "{server_ip}"
 	serverPort = "{server_port}"
 )
-
-// Subnet maximum bit mask length
-const (
-	maxV4BitMaskLen = 32
-	maxV6BitMaskLen = 128
-)