diff options
-rw-r--r-- | go.mod | 1 | ||||
-rw-r--r-- | go.sum | 1 | ||||
-rw-r--r-- | internal/ratelimiter/ratelimiter.go | 101 | ||||
-rw-r--r-- | internal/ratelimiter/ratelimiter_test.go | 107 |
4 files changed, 210 insertions, 0 deletions
@@ -25,4 +25,5 @@ require ( golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9 golang.org/x/net v0.0.0-20201202161906-c7110b5ffcbb golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1 + golang.org/x/time v0.0.0-20190308202827-9d24e82272b4 ) @@ -397,6 +397,7 @@ golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk= golang.org/x/text v0.3.3 h1:cokOdA+Jmi5PJGXLlLllQSgYigAEfHXJAERHVMaCc2k= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= +golang.org/x/time v0.0.0-20190308202827-9d24e82272b4 h1:SvFZT6jyqRaOeXpc5h/JSfZenJ2O330aBsf7JfSUXmQ= golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20181221001348-537d06c36207/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= diff --git a/internal/ratelimiter/ratelimiter.go b/internal/ratelimiter/ratelimiter.go new file mode 100644 index 00000000..e1cf076d --- /dev/null +++ b/internal/ratelimiter/ratelimiter.go @@ -0,0 +1,101 @@ +package ratelimiter + +import ( + "time" + + "github.com/prometheus/client_golang/prometheus" + "golang.org/x/time/rate" + + "gitlab.com/gitlab-org/gitlab-pages/internal/lru" +) + +const ( + // DefaultSourceIPLimitPerSecond is the limit per second that rate.Limiter + // needs to generate tokens every second. + // The default value is 20 requests per second. + DefaultSourceIPLimitPerSecond = 20.0 + // DefaultSourceIPBurstSize is the maximum burst allowed per rate limiter. + // E.g. The first 100 requests within 1s will succeed, but the 101st will fail. + DefaultSourceIPBurstSize = 100 + + // based on an avg ~4,000 unique IPs per minute + // https://log.gprd.gitlab.net/app/lens#/edit/f7110d00-2013-11ec-8c8e-ed83b5469915?_g=h@e78830b + defaultSourceIPItems = 5000 + defaultSourceIPExpirationInterval = time.Minute +) + +// Option function to configure a RateLimiter +type Option func(*RateLimiter) + +// RateLimiter holds an LRU cache of elements to be rate limited. Currently, it supports +// a sourceIPCache and each item returns a rate.Limiter. +// It uses "golang.org/x/time/rate" as its Token Bucket rate limiter per source IP entry. +// See example https://www.fatalerrors.org/a/design-and-implementation-of-time-rate-limiter-for-golang-standard-library.html +// It also holds a now function that can be mocked in unit tests. +type RateLimiter struct { + now func() time.Time + sourceIPLimitPerSecond float64 + sourceIPBurstSize int + sourceIPCache *lru.Cache + // TODO: add domainCache https://gitlab.com/gitlab-org/gitlab-pages/-/issues/630 +} + +// New creates a new RateLimiter with default values that can be configured via Option functions +func New(opts ...Option) *RateLimiter { + rl := &RateLimiter{ + now: time.Now, + sourceIPLimitPerSecond: DefaultSourceIPLimitPerSecond, + sourceIPBurstSize: DefaultSourceIPBurstSize, + sourceIPCache: lru.New( + "source_ip", + defaultSourceIPItems, + defaultSourceIPExpirationInterval, + // TODO: @jaime to add proper metrics in subsequent MR + prometheus.NewGaugeVec(prometheus.GaugeOpts{}, []string{"op"}), + prometheus.NewCounterVec(prometheus.CounterOpts{}, []string{"op", "cache"}), + ), + } + + for _, opt := range opts { + opt(rl) + } + + return rl +} + +// WithNow replaces the RateLimiter now function +func WithNow(now func() time.Time) Option { + return func(rl *RateLimiter) { + rl.now = now + } +} + +// WithSourceIPLimitPerSecond allows configuring per source IP limit per second for RateLimiter +func WithSourceIPLimitPerSecond(limit float64) Option { + return func(rl *RateLimiter) { + rl.sourceIPLimitPerSecond = limit + } +} + +// WithSourceIPBurstSize configures burst per source IP for the RateLimiter +func WithSourceIPBurstSize(burst int) Option { + return func(rl *RateLimiter) { + rl.sourceIPBurstSize = burst + } +} + +func (rl *RateLimiter) getSourceIPLimiter(sourceIP string) *rate.Limiter { + limiterI, _ := rl.sourceIPCache.FindOrFetch(sourceIP, sourceIP, func() (interface{}, error) { + return rate.NewLimiter(rate.Limit(rl.sourceIPLimitPerSecond), rl.sourceIPBurstSize), nil + }) + + return limiterI.(*rate.Limiter) +} + +// SourceIPAllowed checks that the real remote IP address is allowed to perform an operation +func (rl *RateLimiter) SourceIPAllowed(sourceIP string) bool { + limiter := rl.getSourceIPLimiter(sourceIP) + + // AllowN allows us to use the rl.now function, so we can test this more easily. + return limiter.AllowN(rl.now(), 1) +} diff --git a/internal/ratelimiter/ratelimiter_test.go b/internal/ratelimiter/ratelimiter_test.go new file mode 100644 index 00000000..cdf12fe6 --- /dev/null +++ b/internal/ratelimiter/ratelimiter_test.go @@ -0,0 +1,107 @@ +package ratelimiter + +import ( + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +var ( + now = "2021-09-13T15:00:00Z" + validTime, _ = time.Parse(time.RFC3339, now) +) + +func mockNow() time.Time { + return validTime +} + +func TestSourceIPAllowed(t *testing.T) { + t.Parallel() + + tcs := map[string]struct { + now string + sourceIPLimit float64 + sourceIPBurstSize int + reqNum int + }{ + "one_request_per_second": { + sourceIPLimit: 1, + sourceIPBurstSize: 1, + reqNum: 2, + }, + "one_request_per_second_but_big_bucket": { + sourceIPLimit: 1, + sourceIPBurstSize: 10, + reqNum: 11, + }, + "three_req_per_second_bucket_size_one": { + sourceIPLimit: 3, + sourceIPBurstSize: 1, // max burst 1 means 1 at a time + reqNum: 3, + }, + "10_requests_per_second": { + sourceIPLimit: 10, + sourceIPBurstSize: 10, + reqNum: 11, + }, + } + + for tn, tc := range tcs { + t.Run(tn, func(t *testing.T) { + rl := New( + WithNow(mockNow), + WithSourceIPLimitPerSecond(tc.sourceIPLimit), + WithSourceIPBurstSize(tc.sourceIPBurstSize), + ) + + for i := 0; i < tc.reqNum; i++ { + got := rl.SourceIPAllowed("172.16.123.1") + if i < tc.sourceIPBurstSize { + require.Truef(t, got, "expected true for request no. %d", i) + } else { + // requests should fail after reaching tc.sourceIPBurstSize because mockNow + // always returns the same time + require.False(t, got, "expected false for request no. %d", i) + } + } + }) + } +} + +func TestSingleRateLimiterWithMultipleSourceIPs(t *testing.T) { + rate := 10 * time.Millisecond + rl := New( + WithSourceIPLimitPerSecond(float64(1/rate)), + WithSourceIPBurstSize(1), + ) + + wg := sync.WaitGroup{} + + testFn := func(domain string) func(t *testing.T) { + return func(t *testing.T) { + wg.Add(1) + go func() { + defer wg.Done() + + for i := 0; i < 5; i++ { + got := rl.SourceIPAllowed(domain) + require.Truef(t, got, "expected true for request no. %d", i) + time.Sleep(rate) + } + }() + } + } + + first := "172.16.123.10" + t.Run(first, testFn(first)) + + second := "172.16.123.20" + t.Run(second, testFn(second)) + + third := "172.16.123.30" + t.Run(third, testFn(third)) + + wg.Wait() +} |