/
ratelimiter.go
215 lines (193 loc) · 6.07 KB
/
ratelimiter.go
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
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
package m
import (
"net/http"
"sync"
"time"
"github.com/gin-gonic/gin"
"github.com/gotomicro/cetus/e"
"golang.org/x/time/rate"
"github.com/gotomicro/cetus/m/pkg"
)
type (
// RateLimiterStore is the interface to be implemented by custom stores.
RateLimiterStore interface {
// Stores for the rate limiter have to implement the Allow method
Allow(identifier string) (bool, error)
}
)
type (
// RateLimiterConfig defines the configuration for the rate limiter
RateLimiterConfig struct {
Skipper pkg.Skipper
BeforeFunc pkg.BeforeFunc
// IdentifierExtractor uses *gin.Context to extract the identifier for a visitor
IdentifierExtractor Extractor
// Store defines a store for the rate limiter
Store RateLimiterStore
// ErrorHandler provides a handler to be called when IdentifierExtractor returns an error
ErrorHandler func(context *gin.Context, err error) error
// DenyHandler provides a handler to be called when RateLimiter denies access
DenyHandler func(context *gin.Context, identifier string, err error) error
}
// Extractor is used to extract data from *gin.Context
Extractor func(context *gin.Context) (string, error)
)
// errors
var (
// ErrRateLimitExceeded denotes an error raised when rate limit is exceeded
ErrRateLimitExceeded = e.NewHTTPError(http.StatusTooManyRequests, "rate limit exceeded")
// ErrExtractorError denotes an error raised when extractor function is unsuccessful
ErrExtractorError = e.NewHTTPError(http.StatusForbidden, "error while extracting identifier")
)
// DefaultRateLimiterConfig defines default values for RateLimiterConfig
var DefaultRateLimiterConfig = RateLimiterConfig{
Skipper: pkg.DefaultSkipper,
IdentifierExtractor: func(ctx *gin.Context) (string, error) {
id := ctx.ClientIP()
return id, nil
},
ErrorHandler: func(context *gin.Context, err error) error {
return &e.HTTPError{
Code: ErrExtractorError.Code,
Message: ErrExtractorError.Message,
Internal: err,
}
},
DenyHandler: func(context *gin.Context, identifier string, err error) error {
return &e.HTTPError{
Code: ErrRateLimitExceeded.Code,
Message: ErrRateLimitExceeded.Message,
Internal: err,
}
},
}
/*
RateLimiter returns a rate limiting middleware
limiterStore := middleware.NewRateLimiterMemoryStore(20)
e.GET("/rate-limited", func(c *gin.Context) {
c.String(http.StatusOK, "test")
return
}, RateLimiter(limiterStore))
*/
func RateLimiter(store RateLimiterStore) gin.HandlerFunc {
config := DefaultRateLimiterConfig
config.Store = store
return RateLimiterWithConfig(config)
}
func RateLimiterWithConfig(config RateLimiterConfig) gin.HandlerFunc {
if config.Skipper == nil {
config.Skipper = DefaultRateLimiterConfig.Skipper
}
if config.IdentifierExtractor == nil {
config.IdentifierExtractor = DefaultRateLimiterConfig.IdentifierExtractor
}
if config.ErrorHandler == nil {
config.ErrorHandler = DefaultRateLimiterConfig.ErrorHandler
}
if config.DenyHandler == nil {
config.DenyHandler = DefaultRateLimiterConfig.DenyHandler
}
if config.Store == nil {
panic("Store configuration must be provided")
}
return func(c *gin.Context) {
if config.Skipper(c) {
c.Next()
return
}
if config.BeforeFunc != nil {
config.BeforeFunc(c)
}
identifier, err := config.IdentifierExtractor(c)
if err != nil {
c.Error(config.ErrorHandler(c, err))
return
}
if allow, err := config.Store.Allow(identifier); !allow {
c.Error(config.DenyHandler(c, identifier, err))
return
}
c.Next()
return
}
}
type (
// RateLimiterMemoryStore is the built-in store implementation for RateLimiter
RateLimiterMemoryStore struct {
visitors map[string]*Visitor
mutex sync.Mutex
rate rate.Limit
burst int
expiresIn time.Duration
lastCleanup time.Time
}
// Visitor signifies a unique user's limiter details
Visitor struct {
*rate.Limiter
lastSeen time.Time
}
)
func NewRateLimiterMemoryStore(rate rate.Limit) (store *RateLimiterMemoryStore) {
return NewRateLimiterMemoryStoreWithConfig(RateLimiterMemoryStoreConfig{
Rate: rate,
})
}
func NewRateLimiterMemoryStoreWithConfig(config RateLimiterMemoryStoreConfig) (store *RateLimiterMemoryStore) {
store = &RateLimiterMemoryStore{}
store.rate = config.Rate
store.burst = config.Burst
store.expiresIn = config.ExpiresIn
if config.ExpiresIn == 0 {
store.expiresIn = DefaultRateLimiterMemoryStoreConfig.ExpiresIn
}
if config.Burst == 0 {
store.burst = int(config.Rate)
}
store.visitors = make(map[string]*Visitor)
store.lastCleanup = now()
return
}
// RateLimiterMemoryStoreConfig represents configuration for RateLimiterMemoryStore
type RateLimiterMemoryStoreConfig struct {
Rate rate.Limit // Rate of requests allowed to pass as req/s
Burst int // Burst additionally allows a number of requests to pass when rate limit is reached
ExpiresIn time.Duration // ExpiresIn is the duration after that a rate limiter is cleaned up
}
// DefaultRateLimiterMemoryStoreConfig provides default configuration values for RateLimiterMemoryStore
var DefaultRateLimiterMemoryStoreConfig = RateLimiterMemoryStoreConfig{
ExpiresIn: 3 * time.Minute,
}
// Allow implements RateLimiterStore.Allow
func (store *RateLimiterMemoryStore) Allow(identifier string) (bool, error) {
store.mutex.Lock()
limiter, exists := store.visitors[identifier]
if !exists {
limiter = new(Visitor)
limiter.Limiter = rate.NewLimiter(store.rate, store.burst)
store.visitors[identifier] = limiter
}
limiter.lastSeen = now()
if now().Sub(store.lastCleanup) > store.expiresIn {
store.cleanupStaleVisitors()
}
store.mutex.Unlock()
return limiter.AllowN(now(), 1), nil
}
/*
cleanupStaleVisitors helps manage the size of the visitors map by removing stale records
of users who haven't visited again after the configured expiry time has elapsed
*/
func (store *RateLimiterMemoryStore) cleanupStaleVisitors() {
for id, visitor := range store.visitors {
if now().Sub(visitor.lastSeen) > store.expiresIn {
delete(store.visitors, id)
}
}
store.lastCleanup = now()
}
/*
actual time method which is mocked in test file
*/
var now = func() time.Time {
return time.Now()
}