Fix whitelist and XFF.

server/header_rewriter.go

@@ -11,20 +11,20 @@ import (
 
 // NewHeaderRewriter Create a header rewriter
 func NewHeaderRewriter(trustedIPs []string, insecure bool) (forward.ReqRewriter, error) {
-	IPs, err := whitelist.NewIP(trustedIPs, insecure, true)
+	ips, err := whitelist.NewIP(trustedIPs, insecure, true)
 	if err != nil {
 		return nil, err
 	}
 
-	h, err := os.Hostname()
+	hostname, err := os.Hostname()
 	if err != nil {
-		h = "localhost"
+		hostname = "localhost"
 	}
 
 	return &headerRewriter{
-		secureRewriter:   &forward.HeaderRewriter{TrustForwardHeader: true, Hostname: h},
-		insecureRewriter: &forward.HeaderRewriter{TrustForwardHeader: false, Hostname: h},
-		ips:              IPs,
+		secureRewriter:   &forward.HeaderRewriter{TrustForwardHeader: false, Hostname: hostname},
+		insecureRewriter: &forward.HeaderRewriter{TrustForwardHeader: true, Hostname: hostname},
+		ips:              ips,
 		insecure:         insecure,
 	}, nil
 }
@@ -37,16 +37,17 @@ type headerRewriter struct {
 }
 
 func (h *headerRewriter) Rewrite(req *http.Request) {
-	authorized, _, err := h.ips.IsAuthorized(req)
+	if h.insecure {
+		h.insecureRewriter.Rewrite(req)
+		return
+	}
+
+	err := h.ips.IsAuthorized(req)
 	if err != nil {
 		log.Error(err)
 		h.secureRewriter.Rewrite(req)
 		return
 	}
 
-	if h.insecure || authorized {
-		h.secureRewriter.Rewrite(req)
-	} else {
-		h.insecureRewriter.Rewrite(req)
-	}
+	h.insecureRewriter.Rewrite(req)
 }

server/header_rewriter_test.go

@@ -0,0 +1,104 @@
+package server
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestHeaderRewriter_Rewrite(t *testing.T) {
+	testCases := []struct {
+		desc       string
+		remoteAddr string
+		trustedIPs []string
+		insecure   bool
+		expected   map[string]string
+	}{
+		{
+			desc:       "Secure & authorized",
+			remoteAddr: "10.10.10.10:80",
+			trustedIPs: []string{"10.10.10.10"},
+			insecure:   false,
+			expected: map[string]string{
+				"X-Foo":           "bar",
+				"X-Forwarded-For": "30.30.30.30",
+			},
+		},
+		{
+			desc:       "Secure & unauthorized",
+			remoteAddr: "50.50.50.50:80",
+			trustedIPs: []string{"10.10.10.10"},
+			insecure:   false,
+			expected: map[string]string{
+				"X-Foo":           "bar",
+				"X-Forwarded-For": "",
+			},
+		},
+		{
+			desc:       "Secure & authorized error",
+			remoteAddr: "10.10.10.10",
+			trustedIPs: []string{"10.10.10.10"},
+			insecure:   false,
+			expected: map[string]string{
+				"X-Foo":           "bar",
+				"X-Forwarded-For": "",
+			},
+		},
+		{
+			desc:       "insecure & authorized",
+			remoteAddr: "10.10.10.10:80",
+			trustedIPs: []string{"10.10.10.10"},
+			insecure:   true,
+			expected: map[string]string{
+				"X-Foo":           "bar",
+				"X-Forwarded-For": "30.30.30.30",
+			},
+		},
+		{
+			desc:       "insecure & unauthorized",
+			remoteAddr: "50.50.50.50:80",
+			trustedIPs: []string{"10.10.10.10"},
+			insecure:   true,
+			expected: map[string]string{
+				"X-Foo":           "bar",
+				"X-Forwarded-For": "30.30.30.30",
+			},
+		},
+		{
+			desc:       "insecure & authorized error",
+			remoteAddr: "10.10.10.10",
+			trustedIPs: []string{"10.10.10.10"},
+			insecure:   true,
+			expected: map[string]string{
+				"X-Foo":           "bar",
+				"X-Forwarded-For": "30.30.30.30",
+			},
+		},
+	}
+
+	for _, test := range testCases {
+		test := test
+		t.Run(test.desc, func(t *testing.T) {
+			t.Parallel()
+
+			rewriter, err := NewHeaderRewriter(test.trustedIPs, test.insecure)
+			require.NoError(t, err)
+
+			req := httptest.NewRequest(http.MethodGet, "http://20.20.20.20/foo", nil)
+			require.NoError(t, err)
+			req.RemoteAddr = test.remoteAddr
+
+			req.Header.Set("X-Foo", "bar")
+			req.Header.Set("X-Forwarded-For", "30.30.30.30")
+
+			rewriter.Rewrite(req)
+
+			for key, value := range test.expected {
+				assert.Equal(t, value, req.Header.Get(key))
+			}
+		})
+	}
+}

server/server.go

@@ -808,7 +808,8 @@ func (s *Server) prepareServer(entryPointName string, entryPoint *configuration.
 				if !ok {
 					return false, fmt.Errorf("type error %v", addr)
 				}
-				return IPs.ContainsIP(ip.IP)
+
+				return IPs.ContainsIP(ip.IP), nil
 			},
 		}
 	}