feat(): custom real ip header

#908

middleware/realip.go

@@ -9,10 +9,12 @@ import (
 	"strings"
 )
 
-var cfConnectionIP = http.CanonicalHeaderKey("CF-Connecting-IP")
-var trueClientIP = http.CanonicalHeaderKey("True-Client-IP")
-var xForwardedFor = http.CanonicalHeaderKey("X-Forwarded-For")
-var xRealIP = http.CanonicalHeaderKey("X-Real-IP")
+var DefaultRealIPHeaders = []string{
+	"CF-Connecting-IP", // Cloudflare free plan
+	"True-Client-IP",   // Cloudflare Enterprise plan
+	"X-Real-IP",
+	"X-Forwarded-For",
+}
 
 // RealIP is a middleware that sets a http.Request's RemoteAddr to the results
 // of parsing either the CF-Connecting-IP, True-Client-IP, X-Real-IP or the X-Forwarded-For headers
@@ -31,7 +33,7 @@ var xRealIP = http.CanonicalHeaderKey("X-Real-IP")
 // how you're using RemoteAddr, vulnerable to an attack of some sort).
 func RealIP(h http.Handler) http.Handler {
 	fn := func(w http.ResponseWriter, r *http.Request) {
-		if rip := realIP(r); rip != "" {
+		if rip := getRealIP(r, DefaultRealIPHeaders); rip != "" {
 			r.RemoteAddr = rip
 		}
 		h.ServeHTTP(w, r)
@@ -40,24 +42,34 @@ func RealIP(h http.Handler) http.Handler {
 	return http.HandlerFunc(fn)
 }
 
-func realIP(r *http.Request) string {
-	var ip string
-
-	if cfcip := r.Header.Get(cfConnectionIP); cfcip != "" {
-		ip = cfcip
-	} else if tcip := r.Header.Get(trueClientIP); tcip != "" {
-		ip = tcip
-	} else if xrip := r.Header.Get(xRealIP); xrip != "" {
-		ip = xrip
-	} else if xff := r.Header.Get(xForwardedFor); xff != "" {
-		i := strings.Index(xff, ",")
-		if i == -1 {
-			i = len(xff)
+// RealIPCustomHeader is a middleware that sets a http.Request's RemoteAddr to the results
+// of parsing the custom headers.
+//
+// usage:
+// r.Use(RealIPCustomHeader([]string{"X-CUSTOM-IP"}))
+// r.Use(RealIPCustomHeader(append(DefaultRealIPHeaders, "X-CUSTOM-IP")))
+func RealIPCustomHeader(realIPHeaders []string) func(http.Handler) http.Handler {
+	f := func(h http.Handler) http.Handler {
+		fn := func(w http.ResponseWriter, r *http.Request) {
+			if rip := getRealIP(r, realIPHeaders); rip != "" {
+				r.RemoteAddr = rip
+			}
+			h.ServeHTTP(w, r)
 		}
-		ip = xff[:i]
+		return http.HandlerFunc(fn)
 	}
-	if ip == "" || net.ParseIP(ip) == nil {
-		return ""
+	return f
+}
+
+func getRealIP(r *http.Request, realIPHeaders []string) string {
+	for _, header := range realIPHeaders {
+		if ip := r.Header.Get(header); ip != "" {
+			ips := strings.Split(ip, ",")
+			if ips[0] == "" || net.ParseIP(ips[0]) == nil {
+				continue
+			}
+			return ips[0]
+		}
 	}
-	return ip
+	return ""
 }

middleware/realip_test.go

@@ -113,3 +113,52 @@ func TestInvalidIP(t *testing.T) {
 		t.Fatal("Invalid IP used.")
 	}
 }
+
+func TestCustomIPHeader(t *testing.T) {
+	var customHeaderKey = "X-CUSTOM-IP"
+	req, _ := http.NewRequest("GET", "/", nil)
+	req.Header.Add(customHeaderKey, "100.100.100.100")
+	w := httptest.NewRecorder()
+
+	r := chi.NewRouter()
+	r.Use(RealIPCustomHeader([]string{customHeaderKey}))
+
+	realIP := ""
+	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
+		realIP = r.RemoteAddr
+		w.Write([]byte("Hello World"))
+	})
+	r.ServeHTTP(w, req)
+
+	if w.Code != 200 {
+		t.Fatal("Response Code should be 200")
+	}
+
+	if realIP != "100.100.100.100" {
+		t.Fatal("Test get real IP precedence error.")
+	}
+}
+
+func TestCustomIPHeaderWithoutDefault(t *testing.T) {
+	req, _ := http.NewRequest("GET", "/", nil)
+	req.Header.Add("X-REAL-IP", "100.100.100.100")
+	w := httptest.NewRecorder()
+
+	r := chi.NewRouter()
+	r.Use(RealIPCustomHeader([]string{"X-CUSTOM-IP"}))
+
+	realIP := ""
+	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
+		realIP = r.RemoteAddr
+		w.Write([]byte("Hello World"))
+	})
+	r.ServeHTTP(w, req)
+
+	if w.Code != 200 {
+		t.Fatal("Response Code should be 200")
+	}
+
+	if realIP != "" {
+		t.Fatal("Invalid IP used.")
+	}
+}