Rewrite RealIP() to avoid returning an empty string

2 files changed, 111 insertions, 0 deletions

diff --git a/http/request/request.go b/http/request/request.go
index da79a3f..1de6f73 100644
--- a/http/request/request.go
+++ b/http/request/request.go
@@ -6,8 +6,10 @@ package request
 
 import (
 	"fmt"
+	"net"
 	"net/http"
 	"strconv"
+	"strings"
 
 	"github.com/gorilla/mux"
 )
@@ -88,3 +90,30 @@ func HasQueryParam(r *http.Request, param string) bool {
 	_, ok := values[param]
 	return ok
 }
+
+// RealIP returns client's real IP address.
+func RealIP(r *http.Request) string {
+	headers := []string{"X-Forwarded-For", "X-Real-Ip"}
+	for _, header := range headers {
+		value := r.Header.Get(header)
+
+		if value != "" {
+			addresses := strings.Split(value, ",")
+			address := strings.TrimSpace(addresses[0])
+
+			if net.ParseIP(address) != nil {
+				return address
+			}
+		}
+	}
+
+	// Fallback to TCP/IP source IP address.
+	var remoteIP string
+	if strings.ContainsRune(r.RemoteAddr, ':') {
+		remoteIP, _, _ = net.SplitHostPort(r.RemoteAddr)
+	} else {
+		remoteIP = r.RemoteAddr
+	}
+
+	return remoteIP
+}
diff --git a/http/request/request_test.go b/http/request/request_test.go
new file mode 100644
index 0000000..a4aaabf
--- /dev/null
+++ b/http/request/request_test.go
@@ -0,0 +1,82 @@
+// Copyright 2018 Frédéric Guillot. All rights reserved.
+// Use of this source code is governed by the Apache 2.0
+// license that can be found in the LICENSE file.
+
+package request
+
+import (
+	"net/http"
+	"testing"
+)
+
+func TestRealIPWithoutHeaders(t *testing.T) {
+	r := &http.Request{RemoteAddr: "192.168.0.1:4242"}
+	if ip := RealIP(r); ip != "192.168.0.1" {
+		t.Fatalf(`Unexpected result, got: %q`, ip)
+	}
+
+	r = &http.Request{RemoteAddr: "192.168.0.1"}
+	if ip := RealIP(r); ip != "192.168.0.1" {
+		t.Fatalf(`Unexpected result, got: %q`, ip)
+	}
+}
+
+func TestRealIPWithXFFHeader(t *testing.T) {
+	// Test with multiple IPv4 addresses.
+	headers := http.Header{}
+	headers.Set("X-Forwarded-For", "203.0.113.195, 70.41.3.18, 150.172.238.178")
+	r := &http.Request{RemoteAddr: "192.168.0.1:4242", Header: headers}
+
+	if ip := RealIP(r); ip != "203.0.113.195" {
+		t.Fatalf(`Unexpected result, got: %q`, ip)
+	}
+
+	// Test with single IPv6 address.
+	headers = http.Header{}
+	headers.Set("X-Forwarded-For", "2001:db8:85a3:8d3:1319:8a2e:370:7348")
+	r = &http.Request{RemoteAddr: "192.168.0.1:4242", Header: headers}
+
+	if ip := RealIP(r); ip != "2001:db8:85a3:8d3:1319:8a2e:370:7348" {
+		t.Fatalf(`Unexpected result, got: %q`, ip)
+	}
+
+	// Test with single IPv4 address.
+	headers = http.Header{}
+	headers.Set("X-Forwarded-For", "70.41.3.18")
+	r = &http.Request{RemoteAddr: "192.168.0.1:4242", Header: headers}
+
+	if ip := RealIP(r); ip != "70.41.3.18" {
+		t.Fatalf(`Unexpected result, got: %q`, ip)
+	}
+
+	// Test with invalid IP address.
+	headers = http.Header{}
+	headers.Set("X-Forwarded-For", "fake IP")
+	r = &http.Request{RemoteAddr: "192.168.0.1:4242", Header: headers}
+
+	if ip := RealIP(r); ip != "192.168.0.1" {
+		t.Fatalf(`Unexpected result, got: %q`, ip)
+	}
+}
+
+func TestRealIPWithXRealIPHeader(t *testing.T) {
+	headers := http.Header{}
+	headers.Set("X-Real-Ip", "192.168.122.1")
+	r := &http.Request{RemoteAddr: "192.168.0.1:4242", Header: headers}
+
+	if ip := RealIP(r); ip != "192.168.122.1" {
+		t.Fatalf(`Unexpected result, got: %q`, ip)
+	}
+}
+
+func TestRealIPWithBothHeaders(t *testing.T) {
+	headers := http.Header{}
+	headers.Set("X-Forwarded-For", "203.0.113.195, 70.41.3.18, 150.172.238.178")
+	headers.Set("X-Real-Ip", "192.168.122.1")
+
+	r := &http.Request{RemoteAddr: "192.168.0.1:4242", Header: headers}
+
+	if ip := RealIP(r); ip != "203.0.113.195" {
+		t.Fatalf(`Unexpected result, got: %q`, ip)
+	}
+}