1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
|
package ratelimiter
import (
"crypto/tls"
"net"
"net/http"
"time"
"github.com/prometheus/client_golang/prometheus"
"golang.org/x/time/rate"
"gitlab.com/gitlab-org/gitlab-pages/internal/lru"
"gitlab.com/gitlab-org/gitlab-pages/internal/request"
)
const (
// 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
DefaultSourceIPCacheSize = 5000
// we have less than 4000 different hosts per minute
// https://log.gprd.gitlab.net/app/dashboards#/view/d52ab740-61a4-11ec-b20d-65f14d890d9b?_a=(viewMode:edit)&_g=h@42b0d52
DefaultDomainCacheSize = 4000
)
// Option function to configure a RateLimiter
type Option func(*RateLimiter)
// KeyFunc returns unique identifier for the subject of rate limit(e.g. client IP or domain)
type KeyFunc func(*http.Request) string
// TLSKeyFunc is used by GetCertificateMiddleware to identify the subject of rate limit (client IP or SNI servername)
type TLSKeyFunc func(*tls.ClientHelloInfo) string
// RateLimiter holds an LRU cache of elements to be rate limited.
// 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 {
name string
now func() time.Time
keyFunc KeyFunc
tlsKeyFunc TLSKeyFunc
limitPerSecond float64
burstSize int
blockedCount *prometheus.GaugeVec
cache *lru.Cache
cacheOptions []lru.Option
}
// New creates a new RateLimiter with default values that can be configured via Option functions
func New(name string, opts ...Option) *RateLimiter {
rl := &RateLimiter{
name: name,
now: time.Now,
keyFunc: request.GetRemoteAddrWithoutPort,
}
for _, opt := range opts {
opt(rl)
}
if rl.limitPerSecond > 0.0 {
rl.cache = lru.New(name, rl.cacheOptions...)
}
return rl
}
// WithNow replaces the RateLimiter now function
func WithNow(now func() time.Time) Option {
return func(rl *RateLimiter) {
rl.now = now
}
}
// WithLimitPerSecond allows configuring limit per second for RateLimiter
func WithLimitPerSecond(limit float64) Option {
return func(rl *RateLimiter) {
rl.limitPerSecond = limit
}
}
// WithBurstSize configures burst per keyFunc value for the RateLimiter
func WithBurstSize(burst int) Option {
return func(rl *RateLimiter) {
rl.burstSize = burst
}
}
// WithBlockedCountMetric configures metric reporting how many requests were blocked
func WithBlockedCountMetric(m *prometheus.GaugeVec) Option {
return func(rl *RateLimiter) {
rl.blockedCount = m
}
}
// WithCacheMaxSize configures cache size for ratelimiter
func WithCacheMaxSize(size int64) Option {
return func(rl *RateLimiter) {
rl.cacheOptions = append(rl.cacheOptions, lru.WithMaxSize(size))
}
}
// WithCachedEntriesMetric configures metric reporting how many keys are currently stored in
// the rate-limiter cache
func WithCachedEntriesMetric(m *prometheus.GaugeVec) Option {
return func(rl *RateLimiter) {
rl.cacheOptions = append(rl.cacheOptions, lru.WithCachedEntriesMetric(m))
}
}
// WithCachedRequestsMetric configures metric for how many times we access cache
func WithCachedRequestsMetric(m *prometheus.CounterVec) Option {
return func(rl *RateLimiter) {
rl.cacheOptions = append(rl.cacheOptions, lru.WithCachedRequestsMetric(m))
}
}
// WithKeyFunc configures keyFunc
func WithKeyFunc(f KeyFunc) Option {
return func(rl *RateLimiter) {
rl.keyFunc = f
}
}
func TLSHostnameKey(info *tls.ClientHelloInfo) string {
return info.ServerName
}
func TLSClientIPKey(info *tls.ClientHelloInfo) string {
remoteAddr := info.Conn.RemoteAddr().String()
remoteAddr, _, err := net.SplitHostPort(remoteAddr)
if err != nil {
return remoteAddr
}
return remoteAddr
}
func WithTLSKeyFunc(keyFunc TLSKeyFunc) Option {
return func(rl *RateLimiter) {
rl.tlsKeyFunc = keyFunc
}
}
func (rl *RateLimiter) limiter(key string) *rate.Limiter {
limiterI, _ := rl.cache.FindOrFetch(key, key, func() (interface{}, error) {
return rate.NewLimiter(rate.Limit(rl.limitPerSecond), rl.burstSize), nil
})
return limiterI.(*rate.Limiter)
}
// requestAllowed checks if request is within the rate-limit
func (rl *RateLimiter) requestAllowed(r *http.Request) bool {
rateLimitedKey := rl.keyFunc(r)
return rl.allowed(rateLimitedKey)
}
func (rl *RateLimiter) allowed(rateLimitedKey string) bool {
limiter := rl.limiter(rateLimitedKey)
// AllowN allows us to use the rl.now function, so we can test this more easily.
return limiter.AllowN(rl.now(), 1)
}
|
