Ability to use "X-Forwarded-For" as a source of IP for white list.

whitelist/ip.go

@@ -3,36 +3,45 @@ package whitelist
 import (
 	"fmt"
 	"net"
+	"net/http"
 
 	"github.com/pkg/errors"
 )
 
+const (
+	// XForwardedFor Header name
+	XForwardedFor = "X-Forwarded-For"
+)
+
 // IP allows to check that addresses are in a white list
 type IP struct {
-	whiteListsIPs []*net.IP
-	whiteListsNet []*net.IPNet
-	insecure      bool
+	whiteListsIPs    []*net.IP
+	whiteListsNet    []*net.IPNet
+	insecure         bool
+	useXForwardedFor bool
 }
 
-// NewIP builds a new IP given a list of CIDR-Strings to whitelist
-func NewIP(whitelistStrings []string, insecure bool) (*IP, error) {
-	if len(whitelistStrings) == 0 && !insecure {
+// NewIP builds a new IP given a list of CIDR-Strings to white list
+func NewIP(whiteList []string, insecure bool, useXForwardedFor bool) (*IP, error) {
+	if len(whiteList) == 0 && !insecure {
 		return nil, errors.New("no white list provided")
 	}
 
-	ip := IP{insecure: insecure}
+	ip := IP{
+		insecure:         insecure,
+		useXForwardedFor: useXForwardedFor,
+	}
 
 	if !insecure {
-		for _, whitelistString := range whitelistStrings {
-			ipAddr := net.ParseIP(whitelistString)
-			if ipAddr != nil {
+		for _, ipMask := range whiteList {
+			if ipAddr := net.ParseIP(ipMask); ipAddr != nil {
 				ip.whiteListsIPs = append(ip.whiteListsIPs, &ipAddr)
 			} else {
-				_, whitelist, err := net.ParseCIDR(whitelistString)
+				_, ipAddr, err := net.ParseCIDR(ipMask)
 				if err != nil {
-					return nil, fmt.Errorf("parsing CIDR whitelist %s: %v", whitelist, err)
+					return nil, fmt.Errorf("parsing CIDR white list %s: %v", ipAddr, err)
 				}
-				ip.whiteListsNet = append(ip.whiteListsNet, whitelist)
+				ip.whiteListsNet = append(ip.whiteListsNet, ipAddr)
 			}
 		}
 	}
@@ -40,13 +49,38 @@ func NewIP(whitelistStrings []string, insecure bool) (*IP, error) {
 	return &ip, nil
 }
 
-// Contains checks if provided address is in the white list
-func (ip *IP) Contains(addr string) (bool, net.IP, error) {
+// IsAuthorized checks if provided request is authorized by the white list
+func (ip *IP) IsAuthorized(req *http.Request) (bool, net.IP, error) {
 	if ip.insecure {
 		return true, nil, nil
 	}
 
-	ipAddr, err := ipFromRemoteAddr(addr)
+	if ip.useXForwardedFor {
+		xFFs := req.Header[XForwardedFor]
+		if len(xFFs) > 1 {
+			for _, xFF := range xFFs {
+				ok, i, err := ip.contains(parseHost(xFF))
+				if err != nil {
+					return false, nil, err
+				}
+
+				if ok {
+					return ok, i, nil
+				}
+			}
+		}
+	}
+
+	host, _, err := net.SplitHostPort(req.RemoteAddr)
+	if err != nil {
+		return false, nil, err
+	}
+	return ip.contains(host)
+}
+
+// contains checks if provided address is in the white list
+func (ip *IP) contains(addr string) (bool, net.IP, error) {
+	ipAddr, err := parseIP(addr)
 	if err != nil {
 		return false, nil, fmt.Errorf("unable to parse address: %s: %s", addr, err)
 	}
@@ -76,7 +110,7 @@ func (ip *IP) ContainsIP(addr net.IP) (bool, error) {
 	return false, nil
 }
 
-func ipFromRemoteAddr(addr string) (net.IP, error) {
+func parseIP(addr string) (net.IP, error) {
 	userIP := net.ParseIP(addr)
 	if userIP == nil {
 		return nil, fmt.Errorf("can't parse IP from address %s", addr)
@@ -84,3 +118,11 @@ func ipFromRemoteAddr(addr string) (net.IP, error) {
 
 	return userIP, nil
 }
+
+func parseHost(addr string) string {
+	host, _, err := net.SplitHostPort(addr)
+	if err != nil {
+		return addr
+	}
+	return host
+}