diff options
-rw-r--r-- | app.go | 3 | ||||
-rw-r--r-- | internal/ratelimiter/middleware_test.go | 33 | ||||
-rw-r--r-- | internal/ratelimiter/ratelimiter.go | 9 | ||||
-rw-r--r-- | internal/ratelimiter/ratelimiter_test.go | 30 |
4 files changed, 58 insertions, 17 deletions
@@ -265,6 +265,9 @@ func (a *theApp) buildHandlerPipeline() (http.Handler, error) { if a.config.RateLimit.SourceIPLimitPerSecond > 0 { rl := ratelimiter.New( + metrics.RateLimitSourceIPBlockedCount, + metrics.RateLimitSourceIPCachedEntries, + metrics.RateLimitSourceIPCacheRequests, ratelimiter.WithSourceIPLimitPerSecond(a.config.RateLimit.SourceIPLimitPerSecond), ratelimiter.WithSourceIPBurstSize(a.config.RateLimit.SourceIPBurst), ) diff --git a/internal/ratelimiter/middleware_test.go b/internal/ratelimiter/middleware_test.go index b0134b29..2e51fcad 100644 --- a/internal/ratelimiter/middleware_test.go +++ b/internal/ratelimiter/middleware_test.go @@ -6,6 +6,7 @@ import ( "net/http/httptest" "testing" + "github.com/prometheus/client_golang/prometheus/testutil" testlog "github.com/sirupsen/logrus/hooks/test" "github.com/stretchr/testify/require" @@ -23,10 +24,11 @@ var next = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { func TestSourceIPLimiterWithDifferentLimits(t *testing.T) { hook := testlog.NewGlobal() testhelpers.SetEnvironmentVariable(t, testhelpers.FFEnableRateLimiter, "true") + blocked, cachedEntries, cacheReqs := newTestMetrics(t) for tn, tc := range sharedTestCases { t.Run(tn, func(t *testing.T) { - rl := New( + rl := New(blocked, cachedEntries, cacheReqs, WithNow(mockNow), WithSourceIPLimitPerSecond(tc.sourceIPLimit), WithSourceIPBurstSize(tc.sourceIPBurstSize), @@ -63,6 +65,7 @@ func TestSourceIPLimiterWithDifferentLimits(t *testing.T) { func TestSourceIPLimiterDenyRequestsAfterBurst(t *testing.T) { hook := testlog.NewGlobal() + blocked, cachedEntries, cacheReqs := newTestMetrics(t) tcs := map[string]struct { enabled bool @@ -72,23 +75,15 @@ func TestSourceIPLimiterDenyRequestsAfterBurst(t *testing.T) { enabled: false, expectedStatus: http.StatusNoContent, }, - "disabled_rate_limit_https": { - enabled: false, - expectedStatus: http.StatusNoContent, - }, "enabled_rate_limit_http_blocks": { enabled: true, expectedStatus: http.StatusTooManyRequests, }, - "enabled_rate_limit_https_blocks": { - enabled: true, - expectedStatus: http.StatusTooManyRequests, - }, } for tn, tc := range tcs { t.Run(tn, func(t *testing.T) { - rl := New( + rl := New(blocked, cachedEntries, cacheReqs, WithNow(mockNow), WithSourceIPLimitPerSecond(1), WithSourceIPBurstSize(1), @@ -120,6 +115,24 @@ func TestSourceIPLimiterDenyRequestsAfterBurst(t *testing.T) { require.Equal(t, tc.expectedStatus, res.StatusCode) assertSourceIPLog(t, remoteAddr, hook) } + + blockedCount := testutil.ToFloat64(blocked.WithLabelValues("true")) + if tc.enabled { + require.Equal(t, float64(4), blockedCount, "blocked count") + } else { + require.Equal(t, float64(0), blockedCount, "blocked count") + } + blocked.Reset() + + cachedCount := testutil.ToFloat64(cachedEntries.WithLabelValues("source_ip")) + require.Equal(t, float64(1), cachedCount, "cached count") + cachedEntries.Reset() + + cacheReqMiss := testutil.ToFloat64(cacheReqs.WithLabelValues("source_ip", "miss")) + require.Equal(t, float64(1), cacheReqMiss, "miss count") + cacheReqHit := testutil.ToFloat64(cacheReqs.WithLabelValues("source_ip", "hit")) + require.Equal(t, float64(4), cacheReqHit, "hit count") + cacheReqs.Reset() }) } } diff --git a/internal/ratelimiter/ratelimiter.go b/internal/ratelimiter/ratelimiter.go index 1a622c6f..1359b19c 100644 --- a/internal/ratelimiter/ratelimiter.go +++ b/internal/ratelimiter/ratelimiter.go @@ -7,7 +7,6 @@ import ( "golang.org/x/time/rate" "gitlab.com/gitlab-org/gitlab-pages/internal/lru" - "gitlab.com/gitlab-org/gitlab-pages/metrics" ) const ( @@ -43,18 +42,18 @@ type RateLimiter struct { } // New creates a new RateLimiter with default values that can be configured via Option functions -func New(opts ...Option) *RateLimiter { +func New(blockCountMetric, cachedEntriesMetric *prometheus.GaugeVec, cacheRequestsMetric *prometheus.CounterVec, opts ...Option) *RateLimiter { rl := &RateLimiter{ now: time.Now, sourceIPLimitPerSecond: DefaultSourceIPLimitPerSecond, sourceIPBurstSize: DefaultSourceIPBurstSize, - sourceIPBlockedCount: metrics.RateLimitSourceIPBlockedCount, + sourceIPBlockedCount: blockCountMetric, sourceIPCache: lru.New( "source_ip", defaultSourceIPItems, defaultSourceIPExpirationInterval, - metrics.RateLimitSourceIPCachedEntries, - metrics.RateLimitSourceIPCacheRequests, + cachedEntriesMetric, + cacheRequestsMetric, ), } diff --git a/internal/ratelimiter/ratelimiter_test.go b/internal/ratelimiter/ratelimiter_test.go index 77da8e81..03393fb0 100644 --- a/internal/ratelimiter/ratelimiter_test.go +++ b/internal/ratelimiter/ratelimiter_test.go @@ -5,6 +5,7 @@ import ( "testing" "time" + "github.com/prometheus/client_golang/prometheus" "github.com/stretchr/testify/require" ) @@ -47,9 +48,11 @@ var sharedTestCases = map[string]struct { func TestSourceIPAllowed(t *testing.T) { t.Parallel() + blocked, cachedEntries, cacheReqs := newTestMetrics(t) + for tn, tc := range sharedTestCases { t.Run(tn, func(t *testing.T) { - rl := New( + rl := New(blocked, cachedEntries, cacheReqs, WithNow(mockNow), WithSourceIPLimitPerSecond(tc.sourceIPLimit), WithSourceIPBurstSize(tc.sourceIPBurstSize), @@ -71,7 +74,9 @@ func TestSourceIPAllowed(t *testing.T) { func TestSingleRateLimiterWithMultipleSourceIPs(t *testing.T) { rate := 10 * time.Millisecond - rl := New( + blocked, cachedEntries, cacheReqs := newTestMetrics(t) + + rl := New(blocked, cachedEntries, cacheReqs, WithSourceIPLimitPerSecond(float64(1/rate)), WithSourceIPBurstSize(1), ) @@ -104,3 +109,24 @@ func TestSingleRateLimiterWithMultipleSourceIPs(t *testing.T) { wg.Wait() } + +func newTestMetrics(t *testing.T) (*prometheus.GaugeVec, *prometheus.GaugeVec, *prometheus.CounterVec) { + t.Helper() + + blockedGauge := prometheus.NewGaugeVec( + prometheus.GaugeOpts{ + Name: t.Name(), + }, + []string{"enforced"}, + ) + + cachedEntries := prometheus.NewGaugeVec(prometheus.GaugeOpts{ + Name: t.Name(), + }, []string{"op"}) + + cacheReqs := prometheus.NewCounterVec(prometheus.CounterOpts{ + Name: t.Name(), + }, []string{"op", "cache"}) + + return blockedGauge, cachedEntries, cacheReqs +} |