Merge branch '627-add-rate-limiting-pkg' into 'master'

feat: add ratelimiter package

See merge request gitlab-org/gitlab-pages!587

4 files changed, 210 insertions, 0 deletions

diff --git a/go.mod b/go.mod
index 5f17d253..0ff516ed 100644
--- a/go.mod
+++ b/go.mod
@@ -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
 )
diff --git a/go.sum b/go.sum
index b8c849c4..095e5fe1 100644
--- a/go.sum
+++ b/go.sum
@@ -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()
+}