diff options
author | Vladimir Shushlin <v.shushlin@gmail.com> | 2021-12-10 15:14:15 +0300 |
---|---|---|
committer | Vladimir Shushlin <v.shushlin@gmail.com> | 2022-01-11 12:21:44 +0300 |
commit | f62bc0dd08e3a9e14b6bf6c35b04dd6bdcec4491 (patch) | |
tree | 24274bdcd200d7490569415479996c9904a5739e /internal/handlers | |
parent | 2abb995aeb613f184ee0adebfcbd5597ee210169 (diff) |
feat: add domain rate-limiter
Changelog: added
Diffstat (limited to 'internal/handlers')
-rw-r--r-- | internal/handlers/ratelimiter.go | 21 | ||||
-rw-r--r-- | internal/handlers/ratelimiter_test.go | 102 |
2 files changed, 121 insertions, 2 deletions
diff --git a/internal/handlers/ratelimiter.go b/internal/handlers/ratelimiter.go index 8263f497..52281f6e 100644 --- a/internal/handlers/ratelimiter.go +++ b/internal/handlers/ratelimiter.go @@ -4,14 +4,16 @@ import ( "net/http" "gitlab.com/gitlab-org/gitlab-pages/internal/config" + "gitlab.com/gitlab-org/gitlab-pages/internal/feature" "gitlab.com/gitlab-org/gitlab-pages/internal/ratelimiter" + "gitlab.com/gitlab-org/gitlab-pages/internal/request" "gitlab.com/gitlab-org/gitlab-pages/metrics" ) // Ratelimiter configures the ratelimiter middleware // TODO: make this unexported once https://gitlab.com/gitlab-org/gitlab-pages/-/issues/670 is done func Ratelimiter(handler http.Handler, config *config.RateLimit) http.Handler { - rl := ratelimiter.New( + sourceIPLimiter := ratelimiter.New( "source_ip", ratelimiter.WithCacheMaxSize(ratelimiter.DefaultSourceIPCacheSize), ratelimiter.WithCachedEntriesMetric(metrics.RateLimitSourceIPCachedEntries), @@ -19,7 +21,22 @@ func Ratelimiter(handler http.Handler, config *config.RateLimit) http.Handler { ratelimiter.WithBlockedCountMetric(metrics.RateLimitSourceIPBlockedCount), ratelimiter.WithLimitPerSecond(config.SourceIPLimitPerSecond), ratelimiter.WithBurstSize(config.SourceIPBurst), + ratelimiter.WithEnforce(feature.EnforceIPRateLimits.Enabled()), ) - return rl.Middleware(handler) + handler = sourceIPLimiter.Middleware(handler) + + domainLimiter := ratelimiter.New( + "domain", + ratelimiter.WithCacheMaxSize(ratelimiter.DefaultDomainCacheSize), + ratelimiter.WithKeyFunc(request.GetHostWithoutPort), + ratelimiter.WithCachedEntriesMetric(metrics.RateLimitDomainCachedEntries), + ratelimiter.WithCachedRequestsMetric(metrics.RateLimitDomainCacheRequests), + ratelimiter.WithBlockedCountMetric(metrics.RateLimitDomainBlockedCount), + ratelimiter.WithLimitPerSecond(config.DomainLimitPerSecond), + ratelimiter.WithBurstSize(config.DomainBurst), + ratelimiter.WithEnforce(feature.EnforceDomainRateLimits.Enabled()), + ) + + return domainLimiter.Middleware(handler) } diff --git a/internal/handlers/ratelimiter_test.go b/internal/handlers/ratelimiter_test.go new file mode 100644 index 00000000..43acfc9a --- /dev/null +++ b/internal/handlers/ratelimiter_test.go @@ -0,0 +1,102 @@ +package handlers + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/require" + + "gitlab.com/gitlab-org/gitlab-pages/internal/config" + "gitlab.com/gitlab-org/gitlab-pages/internal/feature" + "gitlab.com/gitlab-org/gitlab-pages/internal/testhelpers" +) + +var next = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNoContent) +}) + +func TestRatelimiter(t *testing.T) { + tt := map[string]struct { + firstRemoteAddr string + firstTarget string + secondRemoteAddr string + secondTarget string + sourceIPEnforced bool + domainEnforced bool + expectedSecondCode int + }{ + "rejected_by_ip": { + firstRemoteAddr: "10.0.0.1", + firstTarget: "https://domain.gitlab.io", + secondRemoteAddr: "10.0.0.1", + secondTarget: "https://different.gitlab.io", + sourceIPEnforced: true, + domainEnforced: true, + expectedSecondCode: http.StatusTooManyRequests, + }, + "rejected_by_domain": { + firstRemoteAddr: "10.0.0.1", + firstTarget: "https://domain.gitlab.io", + secondRemoteAddr: "10.0.0.2", + secondTarget: "https://domain.gitlab.io", + sourceIPEnforced: true, + domainEnforced: true, + expectedSecondCode: http.StatusTooManyRequests, + }, + "ip_rate_limiter_disabled": { + firstRemoteAddr: "10.0.0.1", + firstTarget: "https://domain.gitlab.io", + secondRemoteAddr: "10.0.0.1", + secondTarget: "https://different.gitlab.io", + sourceIPEnforced: false, + domainEnforced: true, + expectedSecondCode: http.StatusNoContent, + }, + "domain_rate_limiter_disabled": { + firstRemoteAddr: "10.0.0.1", + firstTarget: "https://domain.gitlab.io", + secondRemoteAddr: "10.0.0.2", + secondTarget: "https://domain.gitlab.io", + sourceIPEnforced: true, + domainEnforced: false, + expectedSecondCode: http.StatusNoContent, + }, + "different_ip_and_domain_passes": { + firstRemoteAddr: "10.0.0.1", + firstTarget: "https://domain.gitlab.io", + secondRemoteAddr: "10.0.0.2", + secondTarget: "https://different.gitlab.io", + sourceIPEnforced: true, + domainEnforced: true, + expectedSecondCode: http.StatusNoContent, + }, + } + + for name, tc := range tt { + t.Run(name, func(t *testing.T) { + testhelpers.StubFeatureFlagValue(t, feature.EnforceIPRateLimits.EnvVariable, tc.sourceIPEnforced) + testhelpers.StubFeatureFlagValue(t, feature.EnforceDomainRateLimits.EnvVariable, tc.domainEnforced) + + conf := config.RateLimit{ + SourceIPLimitPerSecond: 0.1, + SourceIPBurst: 1, + DomainLimitPerSecond: 0.1, + DomainBurst: 1, + } + + handler := Ratelimiter(next, &conf) + + r1 := httptest.NewRequest(http.MethodGet, tc.firstTarget, nil) + r1.RemoteAddr = tc.firstRemoteAddr + + firstCode, _ := testhelpers.PerformRequest(t, handler, r1) + require.Equal(t, http.StatusNoContent, firstCode) + + r2 := httptest.NewRequest(http.MethodGet, tc.secondTarget, nil) + r2.RemoteAddr = tc.secondRemoteAddr + secondCode, _ := testhelpers.PerformRequest(t, handler, r2) + require.Equal(t, tc.expectedSecondCode, secondCode) + }) + } +} |