Add X-Forwarded-For support to IP Whitelist

zetaab · Jan 17, 2016 · 23fa597 · 23fa597
1 parent 9a8f8f9
commit 23fa597
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 0 deletions.
diff --git a/middleware_ip_whitelist.go b/middleware_ip_whitelist.go
@@ -40,6 +40,15 @@ func (i *IPWhiteListMiddleware) ProcessRequest(w http.ResponseWriter, r *http.Re
 
 		splitIP := strings.Split(r.RemoteAddr, ":")
 		remoteIPString := splitIP[0]
+
+		// If X-Forwarded-For is set, override remoteIPString
+		forwarded := r.Header.Get("X-Forwarded-For")
+		if forwarded != "" {
+			ips := strings.Split(forwarded, ", ")
+			remoteIPString = ips[0]
+			log.Info("X-Forwarded-For set, remote IP: ", remoteIPString)
+		}
+
 		if len(splitIP) > 2 {
 			// Might be an IPv6 address, don't mess with it
 			remoteIPString = r.RemoteAddr

diff --git a/middleware_ip_whitelist_test.go b/middleware_ip_whitelist_test.go
@@ -266,6 +266,65 @@ func TestIpMiddlewareIPPassCIDR(t *testing.T) {
 	}
 }
 
+func TestIPMiddlewareIPFailXForwardedFor(t *testing.T) {
+	spec := MakeIPSampleAPI(ipMiddlewareTestDefinitionEnabledPass)
+	redisStore := RedisStorageManager{KeyPrefix: "apikey-"}
+	healthStore := &RedisStorageManager{KeyPrefix: "apihealth."}
+	orgStore := &RedisStorageManager{KeyPrefix: "orgKey."}
+	spec.Init(&redisStore, &redisStore, healthStore, orgStore)
+	thisSession := createNonThrottledSession()
+	spec.SessionManager.UpdateSession("gfgg1234", thisSession, 60)
+	uri := "/about-lonelycoder/"
+	method := "GET"
+
+	recorder := httptest.NewRecorder()
+	param := make(url.Values)
+	req, err := http.NewRequest(method, uri+param.Encode(), nil)
+	req.RemoteAddr = "10.0.0.1"
+	req.Header.Add("authorization", "gfgg1234")
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	chain := getChain(*spec)
+	chain.ServeHTTP(recorder, req)
+
+	if recorder.Code != 403 {
+		t.Error("Invalid response code, should be 403:  \n", recorder.Code, recorder.Body, req.RemoteAddr)
+	}
+}
+
+func TestIPMiddlewareIPPassXForwardedFor(t *testing.T) {
+	spec := MakeIPSampleAPI(ipMiddlewareTestDefinitionEnabledPass)
+	redisStore := RedisStorageManager{KeyPrefix: "apikey-"}
+	healthStore := &RedisStorageManager{KeyPrefix: "apihealth."}
+	orgStore := &RedisStorageManager{KeyPrefix: "orgKey."}
+	spec.Init(&redisStore, &redisStore, healthStore, orgStore)
+	thisSession := createNonThrottledSession()
+	spec.SessionManager.UpdateSession("gfgg1234", thisSession, 60)
+	uri := "/about-lonelycoder/"
+	method := "GET"
+
+	recorder := httptest.NewRecorder()
+	param := make(url.Values)
+	req, err := http.NewRequest(method, uri+param.Encode(), nil)
+	req.RemoteAddr = "10.0.0.1"
+	req.Header.Add("X-Forwarded-For", "127.0.0.1")
+	req.Header.Add("authorization", "gfgg1234")
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	chain := getChain(*spec)
+	chain.ServeHTTP(recorder, req)
+
+	if recorder.Code != 200 {
+		t.Error("Invalid response code, should be 200:  \n", recorder.Code, recorder.Body, req.RemoteAddr)
+	}
+}
+
 func TestIpMiddlewareIPMissing(t *testing.T) {
 	spec := MakeIPSampleAPI(ipMiddlewareTestDefinitionMissing)
 	redisStore := RedisStorageManager{KeyPrefix: "apikey-"}