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

1 files changed, 55 insertions, 0 deletions

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