test: rate limiter acceptance tests

test: rate limit with all listener types

5 files changed, 151 insertions, 28 deletions

diff --git a/internal/ratelimiter/middleware.go b/internal/ratelimiter/middleware.go
index 4f7a91e9..5db004b7 100644
--- a/internal/ratelimiter/middleware.go
+++ b/internal/ratelimiter/middleware.go
@@ -19,8 +19,7 @@ const (
 	headerXForwardedProto = "X-Forwarded-Proto"
 )
 
-// SourceIPLimiter middleware ensures that the originating
-// rate limit. See -rate-limiter
+// SourceIPLimiter returns middleware for rate-limiting clients based on their IP
 func (rl *RateLimiter) SourceIPLimiter(handler http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		host, sourceIP := rl.getReqDetails(r)
@@ -48,6 +47,7 @@ func (rl *RateLimiter) getReqDetails(r *http.Request) (string, string) {
 		host = r.Host
 	}
 
+	// TODO: consider using X-Real-IP https://gitlab.com/gitlab-org/gitlab-pages/-/issues/644
 	// choose between r.RemoteAddr and X-Forwarded-For. Only uses XFF when proxied
 	remoteAddr := xff.GetRemoteAddrIfAllowed(r, func(sip string) bool {
 		// We enable github.com/gorilla/handlers.ProxyHeaders which sets r.RemoteAddr
@@ -84,6 +84,7 @@ func (rl *RateLimiter) logSourceIP(r *http.Request, host, sourceIP string) {
 		Info("source IP hit rate limit")
 }
 
+// TODO: remove https://gitlab.com/gitlab-org/gitlab-pages/-/issues/629
 func rateLimiterEnabled() bool {
 	return os.Getenv("FF_ENABLE_RATE_LIMITER") == "true"
 }
diff --git a/internal/ratelimiter/middleware_test.go b/internal/ratelimiter/middleware_test.go
index eb431554..a6de0d10 100644
--- a/internal/ratelimiter/middleware_test.go
+++ b/internal/ratelimiter/middleware_test.go
@@ -137,6 +137,8 @@ func TestSourceIPLimiterDenyRequestsAfterBurst(t *testing.T) {
 				rr := httptest.NewRequest(http.MethodGet, tc.host, nil)
 				if tc.enabled {
 					testhelpers.EnableRateLimiter(t)
+				} else {
+					testhelpers.DisableRateLimiter(t)
 				}
 
 				rr.Header.Set(headerXForwardedFor, xForwardedFor)
diff --git a/internal/testhelpers/testhelpers.go b/internal/testhelpers/testhelpers.go
index 04dfc3d2..2f1f7c27 100644
--- a/internal/testhelpers/testhelpers.go
+++ b/internal/testhelpers/testhelpers.go
@@ -13,6 +13,9 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+// TODO: remove https://gitlab.com/gitlab-org/gitlab-pages/-/issues/629
+const ffEnableRateLimiter = "FF_ENABLE_RATE_LIMITER"
+
 // AssertHTTP404 asserts handler returns 404 with provided str body
 func AssertHTTP404(t *testing.T, handler http.HandlerFunc, mode, url string, values url.Values, str interface{}) {
 	w := httptest.NewRecorder()
@@ -82,12 +85,26 @@ func Getwd(t *testing.T) string {
 func EnableRateLimiter(t *testing.T) {
 	t.Helper()
 
-	orig := os.Getenv("FF_ENABLE_RATE_LIMITER")
+	orig := os.Getenv(ffEnableRateLimiter)
+
+	err := os.Setenv(ffEnableRateLimiter, "true")
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		os.Setenv(ffEnableRateLimiter, orig)
+	})
+}
+
+// DisableRateLimiter environment variable
+func DisableRateLimiter(t *testing.T) {
+	t.Helper()
+
+	orig := os.Getenv(ffEnableRateLimiter)
 
-	err := os.Setenv("FF_ENABLE_RATE_LIMITER", "true")
+	err := os.Setenv(ffEnableRateLimiter, "false")
 	require.NoError(t, err)
 
 	t.Cleanup(func() {
-		os.Setenv("FF_ENABLE_RATE_LIMITER", orig)
+		os.Setenv(ffEnableRateLimiter, orig)
 	})
 }
diff --git a/test/acceptance/helpers_test.go b/test/acceptance/helpers_test.go
index e2e1c1d0..e600a9e5 100644
--- a/test/acceptance/helpers_test.go
+++ b/test/acceptance/helpers_test.go
@@ -399,6 +399,22 @@ func GetCompressedPageFromListener(t *testing.T, spec ListenSpec, host, urlsuffi
 	return DoPagesRequest(t, spec, req)
 }
 
+func GetPageFromListenerWithRemoteAddrAndXFF(t *testing.T, spec ListenSpec, host, urlsuffix, xForwardedFor, xForwardedHost string) (*http.Response, error) {
+	t.Helper()
+
+	url := spec.URL(urlsuffix)
+	req, err := http.NewRequest("GET", url, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	req.Host = host
+	req.Header.Set("X-Forwarded-For", xForwardedFor)
+	req.Header.Set("X-Forwarded-Host", xForwardedHost)
+
+	return DoPagesRequest(t, spec, req)
+}
+
 func GetProxiedPageFromListener(t *testing.T, spec ListenSpec, host, xForwardedHost, urlsuffix string) (*http.Response, error) {
 	url := spec.URL(urlsuffix)
 	req, err := http.NewRequest("GET", url, nil)
diff --git a/test/acceptance/ratelimiter_test.go b/test/acceptance/ratelimiter_test.go
index d67b3d04..8f96d74b 100644
--- a/test/acceptance/ratelimiter_test.go
+++ b/test/acceptance/ratelimiter_test.go
@@ -1,6 +1,7 @@
 package acceptance_test
 
 import (
+	"fmt"
 	"net/http"
 	"testing"
 	"time"
@@ -10,30 +11,116 @@ import (
 	"gitlab.com/gitlab-org/gitlab-pages/internal/testhelpers"
 )
 
-func TestRateLimitMiddleware(t *testing.T) {
+func TestSourceIPRateLimitMiddleware(t *testing.T) {
 	testhelpers.EnableRateLimiter(t)
 
-	RunPagesProcess(t,
-		withListeners([]ListenSpec{httpListener}),
-		// 10 = 1 req every 100ms
-		withExtraArgument("rate-limit-source-ip", "1.0"),
-		withExtraArgument("rate-limit-source-ip-burst", "1"),
-	)
-
-	for i := 0; i < 20; i++ {
-		rsp1, err := GetPageFromListener(t, httpListener, "group.gitlab-example.com", "project/")
-		require.NoError(t, err)
-		rsp1.Body.Close()
-
-		// every other request should fail
-		//if i%2 != 0 {
-		//	require.Equal(t, http.StatusTooManyRequests, rsp1.StatusCode, "group.gitlab-example.com request: %d failed", i)
-		//	// wait for another token to become available
-		//	time.Sleep(100 * time.Millisecond)
-		//	continue
-		//}
-
-		require.Equal(t, http.StatusOK, rsp1.StatusCode, "group.gitlab-example.com request: %d failed", i)
-		time.Sleep(time.Millisecond)
+	tcs := map[string]struct {
+		listener       ListenSpec
+		rateLimit      float64
+		rateBurst      string
+		blockedIP      string
+		xForwardedHost string
+		xForwardedFor  string
+		expectFail     bool
+		sleep          time.Duration
+	}{
+		"http_slow_requests_should_not_be_blocked": {
+			listener:  httpListener,
+			rateLimit: 1000,
+			// RunPagesProcess makes one request, so we need to allow a burst of 2
+			// because r.RemoteAddr == 127.0.0.1 and X-Forwarded-For is ignored for non-proxy requests
+			// TODO: consider using X-Real-IP https://gitlab.com/gitlab-org/gitlab-pages/-/issues/644
+			rateBurst: "2",
+			sleep:     10 * time.Millisecond,
+		},
+		"https_slow_requests_should_not_be_blocked": {
+			listener:  httpsListener,
+			rateLimit: 1000,
+			rateBurst: "2",
+			sleep:     10 * time.Millisecond,
+		},
+		"proxy_slow_requests_should_not_be_blocked": {
+			listener:  proxyListener,
+			rateLimit: 1000,
+			// listen-proxy uses X-Forwarded-For
+			rateBurst:      "1",
+			xForwardedFor:  "172.16.123.1",
+			xForwardedHost: "group.gitlab-example.com",
+			sleep:          10 * time.Millisecond,
+		},
+		"proxyv2_slow_requests_should_not_be_blocked": {
+			listener:  httpsProxyv2Listener,
+			rateLimit: 1000,
+			rateBurst: "2",
+			sleep:     10 * time.Millisecond,
+		},
+		"http_fast_requests_blocked_after_burst": {
+			listener:   httpListener,
+			rateLimit:  1,
+			rateBurst:  "2",
+			expectFail: true,
+			blockedIP:  "127.0.0.1",
+		},
+		"https_fast_requests_blocked_after_burst": {
+			listener:   httpsListener,
+			rateLimit:  1,
+			rateBurst:  "2",
+			expectFail: true,
+			blockedIP:  "127.0.0.1",
+		},
+		"proxy_fast_requests_blocked_after_burst": {
+			listener:       proxyListener,
+			rateLimit:      1,
+			rateBurst:      "1",
+			xForwardedFor:  "172.16.123.1",
+			xForwardedHost: "group.gitlab-example.com",
+			expectFail:     true,
+			blockedIP:      "172.16.123.1",
+		},
+		"proxyv2_fast_requests_blocked_after_burst": {
+			listener:   httpsProxyv2Listener,
+			rateLimit:  1,
+			rateBurst:  "2",
+			expectFail: true,
+			// use TestProxyv2Client SourceIP
+			blockedIP: "10.1.1.1",
+		},
 	}
+
+	for tn, tc := range tcs {
+		t.Run(tn, func(t *testing.T) {
+			logBuf := RunPagesProcess(t,
+				withListeners([]ListenSpec{tc.listener}),
+				withExtraArgument("rate-limit-source-ip", fmt.Sprint(tc.rateLimit)),
+				withExtraArgument("rate-limit-source-ip-burst", tc.rateBurst),
+			)
+
+			for i := 0; i < 5; i++ {
+				rsp, err := GetPageFromListenerWithRemoteAddrAndXFF(t, tc.listener, "group.gitlab-example.com", "project/", tc.xForwardedFor, tc.xForwardedHost)
+				require.NoError(t, err)
+				rsp.Body.Close()
+
+				if tc.expectFail && i >= int(tc.rateLimit) {
+					require.Equal(t, http.StatusTooManyRequests, rsp.StatusCode, "group.gitlab-example.com request: %d failed", i)
+					assertLogFound(t, logBuf, []string{"source IP hit rate limit", "\"source_ip\":\"" + tc.blockedIP + "\""})
+					continue
+				}
+
+				require.Equal(t, http.StatusOK, rsp.StatusCode, "request: %d failed", i)
+				time.Sleep(tc.sleep)
+			}
+		})
+	}
+}
+
+func assertLogFound(t *testing.T, logBuf *LogCaptureBuffer, expectedLogs []string) {
+	t.Helper()
+
+	// give the process enough time to write the log message
+	require.Eventually(t, func() bool {
+		for _, e := range expectedLogs {
+			require.Contains(t, logBuf.String(), e, "log mismatch")
+		}
+		return true
+	}, 100*time.Millisecond, 10*time.Millisecond)
 }