From 7405f66036537095b52c277d9b56969df33bfa57 Mon Sep 17 00:00:00 2001 From: George Goldberg Date: Tue, 5 Sep 2017 20:34:17 +0100 Subject: PLT-7519: Better rate-limiting. (#7365) --- utils/utils.go | 11 ++++++++++- utils/utils_test.go | 55 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 1 deletion(-) (limited to 'utils') diff --git a/utils/utils.go b/utils/utils.go index 89ea4377e..dd8bdb2a8 100644 --- a/utils/utils.go +++ b/utils/utils.go @@ -8,6 +8,7 @@ import ( "net/http" "net/url" "os" + "strings" "github.com/mattermost/platform/model" ) @@ -55,7 +56,15 @@ func RemoveDuplicatesFromStringArray(arr []string) []string { } func GetIpAddress(r *http.Request) string { - address := r.Header.Get(model.HEADER_FORWARDED) + address := "" + + header := r.Header.Get(model.HEADER_FORWARDED) + if len(header) > 0 { + addresses := strings.Fields(header) + if len(addresses) > 0 { + address = strings.TrimRight(addresses[0], ",") + } + } if len(address) == 0 { address = r.Header.Get(model.HEADER_REAL_IP) diff --git a/utils/utils_test.go b/utils/utils_test.go index b80247867..b18566786 100644 --- a/utils/utils_test.go +++ b/utils/utils_test.go @@ -4,7 +4,10 @@ package utils import ( + "net/http" "testing" + + "github.com/stretchr/testify/assert" ) func TestStringArrayIntersection(t *testing.T) { @@ -44,3 +47,55 @@ func TestRemoveDuplicatesFromStringArray(t *testing.T) { t.Fatal("should be 3") } } + +func TestGetIpAddress(t *testing.T) { + // Test with a single IP in the X-Forwarded-For + httpRequest1 := http.Request{ + Header: http.Header{ + "X-Forwarded-For": []string{"10.0.0.1"}, + "X-Real-Ip": []string{"10.1.0.1"}, + }, + RemoteAddr: "10.2.0.1:12345", + } + + assert.Equal(t, "10.0.0.1", GetIpAddress(&httpRequest1)) + + // Test with multiple IPs in the X-Forwarded-For + httpRequest2 := http.Request{ + Header: http.Header{ + "X-Forwarded-For": []string{"10.0.0.1, 10.0.0.2, 10.0.0.3"}, + "X-Real-Ip": []string{"10.1.0.1"}, + }, + RemoteAddr: "10.2.0.1:12345", + } + + assert.Equal(t, "10.0.0.1", GetIpAddress(&httpRequest2)) + + // Test with an empty X-Forwarded-For + httpRequest3 := http.Request{ + Header: http.Header{ + "X-Forwarded-For": []string{""}, + "X-Real-Ip": []string{"10.1.0.1"}, + }, + RemoteAddr: "10.2.0.1:12345", + } + + assert.Equal(t, "10.1.0.1", GetIpAddress(&httpRequest3)) + + // Test without an X-Fowarded-For + httpRequest4 := http.Request{ + Header: http.Header{ + "X-Real-Ip": []string{"10.1.0.1"}, + }, + RemoteAddr: "10.2.0.1:12345", + } + + assert.Equal(t, "10.1.0.1", GetIpAddress(&httpRequest4)) + + // Test without any headers + httpRequest5 := http.Request{ + RemoteAddr: "10.2.0.1:12345", + } + + assert.Equal(t, "10.2.0.1", GetIpAddress(&httpRequest5)) +} -- cgit v1.2.3-1-g7c22