diff options
Diffstat (limited to 'test/acceptance/ratelimiter_test.go')
-rw-r--r-- | test/acceptance/ratelimiter_test.go | 46 |
1 files changed, 24 insertions, 22 deletions
diff --git a/test/acceptance/ratelimiter_test.go b/test/acceptance/ratelimiter_test.go index 8f96d74b..2986d46b 100644 --- a/test/acceptance/ratelimiter_test.go +++ b/test/acceptance/ratelimiter_test.go @@ -12,24 +12,22 @@ import ( ) func TestSourceIPRateLimitMiddleware(t *testing.T) { - testhelpers.EnableRateLimiter(t) + testhelpers.SetEnvironmentVariable(t, testhelpers.FFEnableRateLimiter, "true") tcs := map[string]struct { - listener ListenSpec - rateLimit float64 - rateBurst string - blockedIP string - xForwardedHost string - xForwardedFor string - expectFail bool - sleep time.Duration + listener ListenSpec + rateLimit float64 + rateBurst string + blockedIP string + header http.Header + 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, }, @@ -43,10 +41,12 @@ func TestSourceIPRateLimitMiddleware(t *testing.T) { 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, + rateBurst: "1", + header: http.Header{ + "X-Forwarded-For": []string{"172.16.123.1"}, + "X-Forwarded-Host": []string{"group.gitlab-example.com"}, + }, + sleep: 10 * time.Millisecond, }, "proxyv2_slow_requests_should_not_be_blocked": { listener: httpsProxyv2Listener, @@ -69,13 +69,15 @@ func TestSourceIPRateLimitMiddleware(t *testing.T) { 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", + listener: proxyListener, + rateLimit: 1, + rateBurst: "1", + header: http.Header{ + "X-Forwarded-For": []string{"172.16.123.1"}, + "X-Forwarded-Host": []string{"group.gitlab-example.com"}, + }, + expectFail: true, + blockedIP: "172.16.123.1", }, "proxyv2_fast_requests_blocked_after_burst": { listener: httpsProxyv2Listener, @@ -96,7 +98,7 @@ func TestSourceIPRateLimitMiddleware(t *testing.T) { ) for i := 0; i < 5; i++ { - rsp, err := GetPageFromListenerWithRemoteAddrAndXFF(t, tc.listener, "group.gitlab-example.com", "project/", tc.xForwardedFor, tc.xForwardedHost) + rsp, err := GetPageFromListenerWithHeaders(t, tc.listener, "group.gitlab-example.com", "project/", tc.header) require.NoError(t, err) rsp.Body.Close() |