PLT-7519: Better rate-limiting. (#7365)

2 files changed, 65 insertions, 1 deletions

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))
+}