Add Pages rate-limiting for IPv6 based on /64 prefix

6 files changed, 120 insertions, 8 deletions

diff --git a/internal/ratelimiter/middleware.go b/internal/ratelimiter/middleware.go
index fa56f5e4..3462ef8a 100644
--- a/internal/ratelimiter/middleware.go
+++ b/internal/ratelimiter/middleware.go
@@ -43,7 +43,7 @@ func (rl *RateLimiter) logRateLimitedRequest(r *http.Request) {
 		"rate_limiter_name":             rl.name,
 		"scheme":                        r.URL.Scheme,
 		"remote_addr":                   r.RemoteAddr,
-		"source_ip":                     request.GetRemoteAddrWithoutPort(r),
+		"source_ip":                     request.GetIPV4orIPV6PrefixWithoutPort(r),
 		"x_forwarded_proto":             r.Header.Get(headerXForwardedProto),
 		"x_forwarded_for":               r.Header.Get(headerXForwardedFor),
 		"gitlab_real_ip":                r.Header.Get(headerGitLabRealIP),
diff --git a/internal/ratelimiter/middleware_test.go b/internal/ratelimiter/middleware_test.go
index fb557462..7937dee6 100644
--- a/internal/ratelimiter/middleware_test.go
+++ b/internal/ratelimiter/middleware_test.go
@@ -122,7 +122,7 @@ func TestKeyFunc(t *testing.T) {
 		expectedSecondCode int
 	}{
 		"rejected_by_ip": {
-			keyFunc:            request.GetRemoteAddrWithoutPort,
+			keyFunc:            request.GetIPV4orIPV6PrefixWithoutPort,
 			firstRemoteAddr:    "10.0.0.1",
 			firstTarget:        "https://domain.gitlab.io",
 			secondRemoteAddr:   "10.0.0.1",
@@ -130,7 +130,7 @@ func TestKeyFunc(t *testing.T) {
 			expectedSecondCode: http.StatusTooManyRequests,
 		},
 		"rejected_by_ip_with_different_port": {
-			keyFunc:            request.GetRemoteAddrWithoutPort,
+			keyFunc:            request.GetIPV4orIPV6PrefixWithoutPort,
 			firstRemoteAddr:    "10.0.0.1:41000",
 			firstTarget:        "https://domain.gitlab.io",
 			secondRemoteAddr:   "10.0.0.1:41001",
@@ -162,7 +162,7 @@ func TestKeyFunc(t *testing.T) {
 			expectedSecondCode: http.StatusNoContent,
 		},
 		"ip_limiter_allows_same_domain": {
-			keyFunc:            request.GetRemoteAddrWithoutPort,
+			keyFunc:            request.GetIPV4orIPV6PrefixWithoutPort,
 			firstRemoteAddr:    "10.0.0.1",
 			firstTarget:        "https://domain.gitlab.io",
 			secondRemoteAddr:   "10.0.0.2",
diff --git a/internal/ratelimiter/ratelimiter.go b/internal/ratelimiter/ratelimiter.go
index bf49cc95..fd74ba0d 100644
--- a/internal/ratelimiter/ratelimiter.go
+++ b/internal/ratelimiter/ratelimiter.go
@@ -54,7 +54,7 @@ func New(name string, opts ...Option) *RateLimiter {
 	rl := &RateLimiter{
 		name:    name,
 		now:     time.Now,
-		keyFunc: request.GetRemoteAddrWithoutPort,
+		keyFunc: request.GetIPV4orIPV6PrefixWithoutPort,
 	}
 
 	for _, opt := range opts {
@@ -135,8 +135,7 @@ func TLSClientIPKey(info *tls.ClientHelloInfo) string {
 	if err != nil {
 		return remoteAddr
 	}
-
-	return remoteAddr
+	return request.GetIPV4orIPV6Prefix(remoteAddr)
 }
 
 func WithTLSKeyFunc(keyFunc TLSKeyFunc) Option {
diff --git a/internal/ratelimiter/tls_test.go b/internal/ratelimiter/tls_test.go
index f198df47..8a18e940 100644
--- a/internal/ratelimiter/tls_test.go
+++ b/internal/ratelimiter/tls_test.go
@@ -30,7 +30,7 @@ func TestTLSClientIPKey(t *testing.T) {
 		},
 		{
 			"[2001:db8:3333:4444:5555:6666:7777:8888]:1234",
-			"2001:db8:3333:4444:5555:6666:7777:8888",
+			"2001:db8:3333:4444::/64",
 		},
 	}
 
diff --git a/internal/request/request.go b/internal/request/request.go
index f98b0819..14ee612b 100644
--- a/internal/request/request.go
+++ b/internal/request/request.go
@@ -3,6 +3,7 @@ package request
 import (
 	"net"
 	"net/http"
+	"net/netip"
 )
 
 const (
@@ -10,6 +11,8 @@ const (
 	SchemeHTTP = "http"
 	// SchemeHTTPS name for the HTTPS scheme
 	SchemeHTTPS = "https"
+	// IPV6PrefixLength is the length of the IPv6 prefix
+	IPV6PrefixLength = 64
 )
 
 // IsHTTPS checks whether the request originated from HTTP or HTTPS.
@@ -38,3 +41,36 @@ func GetRemoteAddrWithoutPort(r *http.Request) string {
 
 	return remoteAddr
 }
+
+// GetIPV4orIPV6PrefixWithoutPort strips the port from the r.RemoteAddr
+// and returns IPV4 address or IPV6 Prefix.
+func GetIPV4orIPV6PrefixWithoutPort(r *http.Request) string {
+	return GetIPV4orIPV6Prefix(r.RemoteAddr)
+}
+
+// GetIPV4orIPV6Prefix returns either the full IPv4 address or the /64 prefix
+// of the IPv6 address from the provided remote address, without the port.
+// For IPv4 it returns the full address. For IPv6 it returns the /64 prefix.
+func GetIPV4orIPV6Prefix(remoteAddr string) string {
+	remoteIP, _, err := net.SplitHostPort(remoteAddr)
+	if err != nil {
+		remoteIP = remoteAddr
+	}
+
+	addr, err := netip.ParseAddr(remoteIP)
+	if err != nil {
+		return remoteIP
+	}
+
+	if addr.Is4() {
+		return remoteIP
+	} else if addr.Is6() {
+		ipv6Prefix, err := addr.Prefix(IPV6PrefixLength)
+		if err != nil {
+			return remoteIP
+		}
+		return ipv6Prefix.String()
+	}
+
+	return remoteIP
+}
diff --git a/internal/request/request_test.go b/internal/request/request_test.go
index 9e71db37..e380c312 100644
--- a/internal/request/request_test.go
+++ b/internal/request/request_test.go
@@ -87,3 +87,80 @@ func TestGetRemoteAddrWithoutPort(t *testing.T) {
 		})
 	}
 }
+
+func TestGetIPV4orIPV6PrefixWithoutPort(t *testing.T) {
+	tests := map[string]struct {
+		u          string
+		remoteAddr string
+		expected   string
+	}{
+		"when IPv4 and port component is provided": {
+			u:          "https://example.com:443",
+			remoteAddr: "127.0.0.1:1000",
+			expected:   "127.0.0.1",
+		},
+		"when IPv4 and port component is not provided": {
+			u:          "http://example.com",
+			remoteAddr: "127.0.0.1",
+			expected:   "127.0.0.1",
+		},
+		"when IPv6 and port component is provided": {
+			u:          "https://[2001:db8:3333:4444:5555:6666:7777:8888]:1234",
+			remoteAddr: "[2001:db8:3333:4444:5555:6666:7777:8888]:1234",
+			expected:   "2001:db8:3333:4444::/64",
+		},
+		"when IPv6 and port component is not provided": {
+			u:          "https://2001:db8:3333:4444:5555:6666:7777:8888",
+			remoteAddr: "2001:db8:3333:4444:5555:6666:7777:8888",
+			expected:   "2001:db8:3333:4444::/64",
+		},
+		"when empty remoteAddr is provided": {
+			u:          "http://example.com",
+			remoteAddr: "",
+			expected:   "",
+		},
+	}
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			req := httptest.NewRequest(http.MethodGet, test.u, nil)
+			req.RemoteAddr = test.remoteAddr
+
+			addr := GetIPV4orIPV6PrefixWithoutPort(req)
+			require.Equal(t, test.expected, addr)
+		})
+	}
+}
+
+func TestGetIPV4orIPV6Prefix(t *testing.T) {
+	tests := map[string]struct {
+		remoteAddr string
+		expected   string
+	}{
+		"when IPv4 and port component is provided": {
+			remoteAddr: "127.0.0.1:1000",
+			expected:   "127.0.0.1",
+		},
+		"when IPv4 and port component is not provided": {
+			remoteAddr: "127.0.0.1",
+			expected:   "127.0.0.1",
+		},
+		"when IPv6 and port component is provided": {
+			remoteAddr: "[2001:db8:3333:4444:5555:6666:7777:8888]:1234",
+			expected:   "2001:db8:3333:4444::/64",
+		},
+		"when IPv6 and port component is not provided": {
+			remoteAddr: "2001:db8:3333:4444:5555:6666:7777:8888",
+			expected:   "2001:db8:3333:4444::/64",
+		},
+		"when empty remoteAddr is provided": {
+			remoteAddr: "",
+			expected:   "",
+		},
+	}
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			addr := GetIPV4orIPV6Prefix(test.remoteAddr)
+			require.Equal(t, test.expected, addr)
+		})
+	}
+}