7 files changed,
833 insertions(+),
0 deletions(-)
Author:
Oleksandr Smirnov
olexsmir@gmail.com
Committed at:
2026-06-23 17:52:56 +0300
Authored at:
2026-06-23 16:38:49 +0300
Change ID:
zrtuwswyxnztrtkwtwkxuvvnuruuqqvr
Parent:
da217b4
A
ratelimit/internal/time/rate/LICENSE
··· 1 +Copyright 2009 The Go Authors. 2 + 3 +Redistribution and use in source and binary forms, with or without 4 +modification, are permitted provided that the following conditions are 5 +met: 6 + 7 + * Redistributions of source code must retain the above copyright 8 +notice, this list of conditions and the following disclaimer. 9 + * Redistributions in binary form must reproduce the above 10 +copyright notice, this list of conditions and the following disclaimer 11 +in the documentation and/or other materials provided with the 12 +distribution. 13 + * Neither the name of Google LLC nor the names of its 14 +contributors may be used to endorse or promote products derived from 15 +this software without specific prior written permission. 16 + 17 +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 18 +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 19 +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 20 +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT 21 +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 22 +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT 23 +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 24 +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 25 +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 26 +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
A
ratelimit/internal/time/rate/rate.go
··· 1 +// Copyright 2015 The Go Authors. All rights reserved. 2 +// Use of this source code is governed by a BSD-style 3 +// license that can be found in the LICENSE file. 4 + 5 +// Package rate provides a rate limiter. 6 +package rate 7 + 8 +import ( 9 + "context" 10 + "fmt" 11 + "math" 12 + "sync" 13 + "time" 14 +) 15 + 16 +// Limit defines the maximum frequency of some events. 17 +// Limit is represented as number of events per second. 18 +// A zero Limit allows no events. 19 +type Limit float64 20 + 21 +// Inf is the infinite rate limit; it allows all events (even if burst is zero). 22 +const Inf = Limit(math.MaxFloat64) 23 + 24 +// Every converts a minimum time interval between events to a Limit. 25 +func Every(interval time.Duration) Limit { 26 + if interval <= 0 { 27 + return Inf 28 + } 29 + return 1 / Limit(interval.Seconds()) 30 +} 31 + 32 +// A Limiter controls how frequently events are allowed to happen. 33 +// It implements a "token bucket" of size b, initially full and refilled 34 +// at rate r tokens per second. 35 +// Informally, in any large enough time interval, the Limiter limits the 36 +// rate to r tokens per second, with a maximum burst size of b events. 37 +// As a special case, if r == Inf (the infinite rate), b is ignored. 38 +// See https://en.wikipedia.org/wiki/Token_bucket for more about token buckets. 39 +// 40 +// The zero value is a valid Limiter, but it will reject all events. 41 +// Use NewLimiter to create non-zero Limiters. 42 +// 43 +// Limiter has three main methods, Allow, Reserve, and Wait. 44 +// Most callers should use Wait. 45 +// 46 +// Each of the three methods consumes a single token. 47 +// They differ in their behavior when no token is available. 48 +// If no token is available, Allow returns false. 49 +// If no token is available, Reserve returns a reservation for a future token 50 +// and the amount of time the caller must wait before using it. 51 +// If no token is available, Wait blocks until one can be obtained 52 +// or its associated context.Context is canceled. 53 +// 54 +// The methods AllowN, ReserveN, and WaitN consume n tokens. 55 +// 56 +// Limiter is safe for simultaneous use by multiple goroutines. 57 +type Limiter struct { 58 + mu sync.Mutex 59 + limit Limit 60 + burst int 61 + tokens float64 62 + // last is the last time the limiter's tokens field was updated 63 + last time.Time 64 + // lastEvent is the latest time of a rate-limited event (past or future) 65 + lastEvent time.Time 66 +} 67 + 68 +// Limit returns the maximum overall event rate. 69 +func (lim *Limiter) Limit() Limit { 70 + lim.mu.Lock() 71 + defer lim.mu.Unlock() 72 + return lim.limit 73 +} 74 + 75 +// Burst returns the maximum burst size. Burst is the maximum number of tokens 76 +// that can be consumed in a single call to Allow, Reserve, or Wait, so higher 77 +// Burst values allow more events to happen at once. 78 +// A zero Burst allows no events, unless limit == Inf. 79 +func (lim *Limiter) Burst() int { 80 + lim.mu.Lock() 81 + defer lim.mu.Unlock() 82 + return lim.burst 83 +} 84 + 85 +// TokensAt returns the number of tokens available at time t. 86 +func (lim *Limiter) TokensAt(t time.Time) float64 { 87 + lim.mu.Lock() 88 + tokens := lim.advance(t) // does not mutate lim 89 + lim.mu.Unlock() 90 + return tokens 91 +} 92 + 93 +// Tokens returns the number of tokens available now. 94 +func (lim *Limiter) Tokens() float64 { 95 + return lim.TokensAt(time.Now()) 96 +} 97 + 98 +// NewLimiter returns a new Limiter that allows events up to rate r and permits 99 +// bursts of at most b tokens. 100 +func NewLimiter(r Limit, b int) *Limiter { 101 + return &Limiter{ 102 + limit: r, 103 + burst: b, 104 + tokens: float64(b), 105 + } 106 +} 107 + 108 +// Allow reports whether an event may happen now. 109 +func (lim *Limiter) Allow() bool { 110 + return lim.AllowN(time.Now(), 1) 111 +} 112 + 113 +// AllowN reports whether n events may happen at time t. 114 +// Use this method if you intend to drop / skip events that exceed the rate limit. 115 +// Otherwise use Reserve or Wait. 116 +func (lim *Limiter) AllowN(t time.Time, n int) bool { 117 + return lim.reserveN(t, n, 0).ok 118 +} 119 + 120 +// A Reservation holds information about events that are permitted by a Limiter to happen after a delay. 121 +// A Reservation may be canceled, which may enable the Limiter to permit additional events. 122 +type Reservation struct { 123 + ok bool 124 + lim *Limiter 125 + tokens int 126 + timeToAct time.Time 127 + // This is the Limit at reservation time, it can change later. 128 + limit Limit 129 +} 130 + 131 +// OK returns whether the limiter can provide the requested number of tokens 132 +// within the maximum wait time. If OK is false, Delay returns InfDuration, and 133 +// Cancel does nothing. 134 +func (r *Reservation) OK() bool { 135 + return r.ok 136 +} 137 + 138 +// Delay is shorthand for DelayFrom(time.Now()). 139 +func (r *Reservation) Delay() time.Duration { 140 + return r.DelayFrom(time.Now()) 141 +} 142 + 143 +// InfDuration is the duration returned by Delay when a Reservation is not OK. 144 +const InfDuration = time.Duration(math.MaxInt64) 145 + 146 +// DelayFrom returns the duration for which the reservation holder must wait 147 +// before taking the reserved action. Zero duration means act immediately. 148 +// InfDuration means the limiter cannot grant the tokens requested in this 149 +// Reservation within the maximum wait time. 150 +func (r *Reservation) DelayFrom(t time.Time) time.Duration { 151 + if !r.ok { 152 + return InfDuration 153 + } 154 + delay := r.timeToAct.Sub(t) 155 + if delay < 0 { 156 + return 0 157 + } 158 + return delay 159 +} 160 + 161 +// Cancel is shorthand for CancelAt(time.Now()). 162 +func (r *Reservation) Cancel() { 163 + r.CancelAt(time.Now()) 164 +} 165 + 166 +// CancelAt indicates that the reservation holder will not perform the reserved action 167 +// and reverses the effects of this Reservation on the rate limit as much as possible, 168 +// considering that other reservations may have already been made. 169 +func (r *Reservation) CancelAt(t time.Time) { 170 + if !r.ok { 171 + return 172 + } 173 + 174 + r.lim.mu.Lock() 175 + defer r.lim.mu.Unlock() 176 + 177 + if r.lim.limit == Inf || r.tokens == 0 || r.timeToAct.Before(t) { 178 + return 179 + } 180 + 181 + // calculate tokens to restore 182 + // The duration between lim.lastEvent and r.timeToAct tells us how many tokens were reserved 183 + // after r was obtained. These tokens should not be restored. 184 + restoreTokens := float64(r.tokens) - r.limit.tokensFromDuration(r.lim.lastEvent.Sub(r.timeToAct)) 185 + if restoreTokens <= 0 { 186 + return 187 + } 188 + // advance time to now 189 + tokens := r.lim.advance(t) 190 + // calculate new number of tokens 191 + tokens += restoreTokens 192 + if burst := float64(r.lim.burst); tokens > burst { 193 + tokens = burst 194 + } 195 + // update state 196 + r.lim.last = t 197 + r.lim.tokens = tokens 198 + if r.timeToAct.Equal(r.lim.lastEvent) { 199 + prevEvent := r.timeToAct.Add(r.limit.durationFromTokens(float64(-r.tokens))) 200 + if !prevEvent.Before(t) { 201 + r.lim.lastEvent = prevEvent 202 + } 203 + } 204 +} 205 + 206 +// Reserve is shorthand for ReserveN(time.Now(), 1). 207 +func (lim *Limiter) Reserve() *Reservation { 208 + return lim.ReserveN(time.Now(), 1) 209 +} 210 + 211 +// ReserveN returns a Reservation that indicates how long the caller must wait before n events happen. 212 +// The Limiter takes this Reservation into account when allowing future events. 213 +// The returned Reservation’s OK() method returns false if n exceeds the Limiter's burst size. 214 +// Usage example: 215 +// 216 +// r := lim.ReserveN(time.Now(), 1) 217 +// if !r.OK() { 218 +// // Not allowed to act! Did you remember to set lim.burst to be > 0 ? 219 +// return 220 +// } 221 +// time.Sleep(r.Delay()) 222 +// Act() 223 +// 224 +// Use this method if you wish to wait and slow down in accordance with the rate limit without dropping events. 225 +// If you need to respect a deadline or cancel the delay, use Wait instead. 226 +// To drop or skip events exceeding rate limit, use Allow instead. 227 +func (lim *Limiter) ReserveN(t time.Time, n int) *Reservation { 228 + r := lim.reserveN(t, n, InfDuration) 229 + return &r 230 +} 231 + 232 +// Wait is shorthand for WaitN(ctx, 1). 233 +func (lim *Limiter) Wait(ctx context.Context) (err error) { 234 + return lim.WaitN(ctx, 1) 235 +} 236 + 237 +// WaitN blocks until lim permits n events to happen. 238 +// It returns an error if n exceeds the Limiter's burst size, the Context is 239 +// canceled, or the expected wait time exceeds the Context's Deadline. 240 +// The burst limit is ignored if the rate limit is Inf. 241 +func (lim *Limiter) WaitN(ctx context.Context, n int) (err error) { 242 + // The test code calls lim.wait with a fake timer generator. 243 + // This is the real timer generator. 244 + newTimer := func(d time.Duration) (<-chan time.Time, func() bool, func()) { 245 + timer := time.NewTimer(d) 246 + return timer.C, timer.Stop, func() {} 247 + } 248 + 249 + return lim.wait(ctx, n, time.Now(), newTimer) 250 +} 251 + 252 +// wait is the internal implementation of WaitN. 253 +func (lim *Limiter) wait(ctx context.Context, n int, t time.Time, newTimer func(d time.Duration) (<-chan time.Time, func() bool, func())) error { 254 + lim.mu.Lock() 255 + burst := lim.burst 256 + limit := lim.limit 257 + lim.mu.Unlock() 258 + 259 + if n > burst && limit != Inf { 260 + return fmt.Errorf("rate: Wait(n=%d) exceeds limiter's burst %d", n, burst) 261 + } 262 + // Check if ctx is already cancelled 263 + select { 264 + case <-ctx.Done(): 265 + return ctx.Err() 266 + default: 267 + } 268 + // Determine wait limit 269 + waitLimit := InfDuration 270 + if deadline, ok := ctx.Deadline(); ok { 271 + waitLimit = deadline.Sub(t) 272 + } 273 + // Reserve 274 + r := lim.reserveN(t, n, waitLimit) 275 + if !r.ok { 276 + return fmt.Errorf("rate: Wait(n=%d) would exceed context deadline", n) 277 + } 278 + // Wait if necessary 279 + delay := r.DelayFrom(t) 280 + if delay == 0 { 281 + return nil 282 + } 283 + ch, stop, advance := newTimer(delay) 284 + defer stop() 285 + advance() // only has an effect when testing 286 + select { 287 + case <-ch: 288 + // We can proceed. 289 + return nil 290 + case <-ctx.Done(): 291 + // Context was canceled before we could proceed. Cancel the 292 + // reservation, which may permit other events to proceed sooner. 293 + r.Cancel() 294 + return ctx.Err() 295 + } 296 +} 297 + 298 +// SetLimit is shorthand for SetLimitAt(time.Now(), newLimit). 299 +func (lim *Limiter) SetLimit(newLimit Limit) { 300 + lim.SetLimitAt(time.Now(), newLimit) 301 +} 302 + 303 +// SetLimitAt sets a new Limit for the limiter. The new Limit, and Burst, may be violated 304 +// or underutilized by those which reserved (using Reserve or Wait) but did not yet act 305 +// before SetLimitAt was called. 306 +func (lim *Limiter) SetLimitAt(t time.Time, newLimit Limit) { 307 + lim.mu.Lock() 308 + defer lim.mu.Unlock() 309 + 310 + tokens := lim.advance(t) 311 + 312 + lim.last = t 313 + lim.tokens = tokens 314 + lim.limit = newLimit 315 +} 316 + 317 +// SetBurst is shorthand for SetBurstAt(time.Now(), newBurst). 318 +func (lim *Limiter) SetBurst(newBurst int) { 319 + lim.SetBurstAt(time.Now(), newBurst) 320 +} 321 + 322 +// SetBurstAt sets a new burst size for the limiter. 323 +func (lim *Limiter) SetBurstAt(t time.Time, newBurst int) { 324 + lim.mu.Lock() 325 + defer lim.mu.Unlock() 326 + 327 + tokens := lim.advance(t) 328 + 329 + lim.last = t 330 + lim.tokens = tokens 331 + lim.burst = newBurst 332 +} 333 + 334 +// reserveN is a helper method for AllowN, ReserveN, and WaitN. 335 +// maxFutureReserve specifies the maximum reservation wait duration allowed. 336 +// reserveN returns Reservation, not *Reservation, to avoid allocation in AllowN and WaitN. 337 +func (lim *Limiter) reserveN(t time.Time, n int, maxFutureReserve time.Duration) Reservation { 338 + lim.mu.Lock() 339 + defer lim.mu.Unlock() 340 + 341 + if lim.limit == Inf { 342 + return Reservation{ 343 + ok: true, 344 + lim: lim, 345 + tokens: n, 346 + timeToAct: t, 347 + } 348 + } 349 + 350 + tokens := lim.advance(t) 351 + 352 + // Calculate the remaining number of tokens resulting from the request. 353 + tokens -= float64(n) 354 + 355 + // Calculate the wait duration 356 + var waitDuration time.Duration 357 + if tokens < 0 { 358 + waitDuration = lim.limit.durationFromTokens(-tokens) 359 + } 360 + 361 + // Decide result 362 + ok := n <= lim.burst && waitDuration <= maxFutureReserve 363 + 364 + // Prepare reservation 365 + r := Reservation{ 366 + ok: ok, 367 + lim: lim, 368 + limit: lim.limit, 369 + } 370 + if ok { 371 + r.tokens = n 372 + r.timeToAct = t.Add(waitDuration) 373 + 374 + // Update state 375 + lim.last = t 376 + lim.tokens = tokens 377 + lim.lastEvent = r.timeToAct 378 + } 379 + 380 + return r 381 +} 382 + 383 +// advance calculates and returns an updated number of tokens for lim 384 +// resulting from the passage of time. 385 +// lim is not changed. 386 +// advance requires that lim.mu is held. 387 +func (lim *Limiter) advance(t time.Time) (newTokens float64) { 388 + last := lim.last 389 + if t.Before(last) { 390 + last = t 391 + } 392 + 393 + // Calculate the new number of tokens, due to time that passed. 394 + elapsed := t.Sub(last) 395 + delta := lim.limit.tokensFromDuration(elapsed) 396 + tokens := lim.tokens + delta 397 + if burst := float64(lim.burst); tokens > burst { 398 + tokens = burst 399 + } 400 + return tokens 401 +} 402 + 403 +// durationFromTokens is a unit conversion function from the number of tokens to the duration 404 +// of time it takes to accumulate them at a rate of limit tokens per second. 405 +func (limit Limit) durationFromTokens(tokens float64) time.Duration { 406 + if limit <= 0 { 407 + return InfDuration 408 + } 409 + 410 + duration := (tokens / float64(limit)) * float64(time.Second) 411 + 412 + // Cap the duration to the maximum representable int64 value, to avoid overflow. 413 + if duration > float64(math.MaxInt64) { 414 + return InfDuration 415 + } 416 + 417 + return time.Duration(duration) 418 +} 419 + 420 +// tokensFromDuration is a unit conversion function from a time duration to the number of tokens 421 +// which could be accumulated during that duration at a rate of limit tokens per second. 422 +func (limit Limit) tokensFromDuration(d time.Duration) float64 { 423 + if limit <= 0 { 424 + return 0 425 + } 426 + return d.Seconds() * float64(limit) 427 +}
A
ratelimit/internal/time/rate/sometimes.go
··· 1 +// Copyright 2022 The Go Authors. All rights reserved. 2 +// Use of this source code is governed by a BSD-style 3 +// license that can be found in the LICENSE file. 4 + 5 +package rate 6 + 7 +import ( 8 + "sync" 9 + "time" 10 +) 11 + 12 +// Sometimes will perform an action occasionally. The First, Every, and 13 +// Interval fields govern the behavior of Do, which performs the action. 14 +// A zero Sometimes value will perform an action exactly once. 15 +// 16 +// # Example: logging with rate limiting 17 +// 18 +// var sometimes = rate.Sometimes{First: 3, Interval: 10*time.Second} 19 +// func Spammy() { 20 +// sometimes.Do(func() { log.Info("here I am!") }) 21 +// } 22 +type Sometimes struct { 23 + First int // if non-zero, the first N calls to Do will run f. 24 + Every int // if non-zero, every Nth call to Do will run f. 25 + Interval time.Duration // if non-zero and Interval has elapsed since f's last run, Do will run f. 26 + 27 + mu sync.Mutex 28 + count int // number of Do calls 29 + last time.Time // last time f was run 30 +} 31 + 32 +// Do runs the function f as allowed by First, Every, and Interval. 33 +// 34 +// The model is a union (not intersection) of filters. The first call to Do 35 +// always runs f. Subsequent calls to Do run f if allowed by First or Every or 36 +// Interval. 37 +// 38 +// A non-zero First:N causes the first N Do(f) calls to run f. 39 +// 40 +// A non-zero Every:M causes every Mth Do(f) call, starting with the first, to 41 +// run f. 42 +// 43 +// A non-zero Interval causes Do(f) to run f if Interval has elapsed since 44 +// Do last ran f. 45 +// 46 +// Specifying multiple filters produces the union of these execution streams. 47 +// For example, specifying both First:N and Every:M causes the first N Do(f) 48 +// calls and every Mth Do(f) call, starting with the first, to run f. See 49 +// Examples for more. 50 +// 51 +// If Do is called multiple times simultaneously, the calls will block and run 52 +// serially. Therefore, Do is intended for lightweight operations. 53 +// 54 +// Because a call to Do may block until f returns, if f causes Do to be called, 55 +// it will deadlock. 56 +func (s *Sometimes) Do(f func()) { 57 + s.mu.Lock() 58 + defer s.mu.Unlock() 59 + if s.count == 0 || 60 + (s.First > 0 && s.count < s.First) || 61 + (s.Every > 0 && s.count%s.Every == 0) || 62 + (s.Interval > 0 && time.Since(s.last) >= s.Interval) { 63 + f() 64 + if s.Interval > 0 { 65 + s.last = time.Now() 66 + } 67 + } 68 + s.count++ 69 +}
A
ratelimit/ratelimit.go
··· 1 +// Package ratelimit provides an HTTP middleware that rate-limits requests 2 +// on a per-client basis using the token bucket algorithm. 3 +// 4 +// Ecample: 5 +// 6 +// limiter := ratelimit.MiddlewareWithConfig(ratelimit.Config{ 7 +// RPS: 10, 8 +// Burst: 20, 9 +// TTL: time.Minute, 10 +// }) 11 +// http.Handle("/api", limiter(yourHandler)) 12 +// 13 +package ratelimit 14 + 15 +import ( 16 + "net" 17 + "net/http" 18 + "strings" 19 + "sync" 20 + "time" 21 + 22 + "olexsmir.xyz/x/ratelimit/internal/time/rate" 23 +) 24 + 25 +type ( 26 + rateLimiter struct { 27 + mu sync.RWMutex 28 + 29 + visitors map[visitorIP]*visitor 30 + 31 + // limit is the maximum number of requests per second 32 + limit rate.Limit 33 + 34 + // ttl is the time after which a visitor is forgotten 35 + ttl time.Duration 36 + 37 + // burst is the maximum number of requests allowed in a short burst 38 + burst int 39 + } 40 + 41 + visitorIP string 42 + visitor struct { 43 + limiter *rate.Limiter 44 + lastSeen time.Time 45 + } 46 +) 47 + 48 +func newLimiter(rps, burst int, ttl time.Duration) *rateLimiter { 49 + return &rateLimiter{ 50 + visitors: make(map[visitorIP]*visitor), 51 + limit: rate.Limit(rps), 52 + burst: burst, 53 + ttl: ttl, 54 + } 55 +} 56 + 57 +// getVisitor retrieves the rate limiter for the given IP, or creates one. 58 +func (r *rateLimiter) getVisitor(ip visitorIP) *rate.Limiter { 59 + r.mu.RLock() 60 + v, exists := r.visitors[ip] 61 + r.mu.RUnlock() 62 + 63 + if !exists { 64 + limit := rate.NewLimiter(r.limit, r.burst) 65 + 66 + r.mu.Lock() 67 + r.visitors[ip] = &visitor{ 68 + limiter: limit, 69 + lastSeen: time.Now(), 70 + } 71 + r.mu.Unlock() 72 + 73 + return limit 74 + } 75 + 76 + r.mu.Lock() 77 + v.lastSeen = time.Now() 78 + r.mu.Unlock() 79 + 80 + return v.limiter 81 +} 82 + 83 +// cleanupVisitors removes visitors that haven't been seen for longer than TTL. 84 +func (r *rateLimiter) cleanupVisitors() { 85 + r.mu.Lock() 86 + defer r.mu.Unlock() 87 + 88 + for ip, v := range r.visitors { 89 + if time.Since(v.lastSeen) > r.ttl { 90 + delete(r.visitors, ip) 91 + } 92 + } 93 +} 94 + 95 +// cleanupVisitorsLoop runs cleanupVisitors every minute. 96 +func (r *rateLimiter) cleanupVisitorsLoop() { 97 + for { 98 + time.Sleep(time.Minute) 99 + r.cleanupVisitors() 100 + } 101 +} 102 + 103 +// Config configures the rate limiter middleware. 104 +type Config struct { 105 + // RPS is the maximum number of requests per second per visitor. 106 + RPS int 107 + 108 + // TTL is the duration after which an idle visitor's limiter is cleaned up. 109 + TTL time.Duration 110 + 111 + // Burst is the maximum number of requests that can be made in a short burst. 112 + Burst int 113 +} 114 + 115 +// Middleware returns a rate-limiting HTTP middleware. 116 +// 117 +// It limits requests per client IP using the token bucket algorithm. 118 +// When the limit is exceeded it responds with 429 Too Many Requests. 119 +func Middleware(c Config) func(http.Handler) http.Handler { 120 + lmt := newLimiter(c.RPS, c.Burst, c.TTL) 121 + go lmt.cleanupVisitorsLoop() 122 + 123 + return func(next http.Handler) http.Handler { 124 + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 125 + ip := getIP(r) 126 + limiter := lmt.getVisitor(ip) 127 + if limiter == nil { 128 + http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError) 129 + return 130 + } 131 + 132 + if !limiter.Allow() { 133 + http.Error(w, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests) 134 + return 135 + } 136 + 137 + next.ServeHTTP(w, r) 138 + }) 139 + } 140 +} 141 + 142 +func getIP(r *http.Request) visitorIP { 143 + if xff := r.Header.Get("X-Forwarded-For"); xff != "" { 144 + if i := strings.IndexByte(xff, ','); i != -1 { 145 + xff = xff[:i] 146 + } 147 + if ip := strings.TrimSpace(xff); ip != "" { 148 + return visitorIP(ip) 149 + } 150 + } 151 + 152 + if xri := r.Header.Get("X-Real-IP"); xri != "" { 153 + return visitorIP(xri) 154 + } 155 + 156 + if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil { 157 + return visitorIP(host) 158 + } 159 + 160 + return visitorIP(r.RemoteAddr) 161 +}
A
ratelimit/ratelimit_test.go
··· 1 +package ratelimit 2 + 3 +import ( 4 + "net/http" 5 + "net/http/httptest" 6 + "testing" 7 + "time" 8 +) 9 + 10 +func TestMiddleware(t *testing.T) { 11 + tests := map[string]struct { 12 + config Config 13 + requests int 14 + expectedCode int 15 + }{ 16 + "allows requests within limit": { 17 + config: Config{ 18 + RPS: 2, 19 + Burst: 2, 20 + TTL: time.Minute, 21 + }, 22 + requests: 1, 23 + expectedCode: http.StatusOK, 24 + }, 25 + "blocks requests over limit": { 26 + config: Config{ 27 + RPS: 1, 28 + Burst: 1, 29 + TTL: time.Minute, 30 + }, 31 + requests: 2, 32 + expectedCode: http.StatusTooManyRequests, 33 + }, 34 + "allows burst requests": { 35 + config: Config{ 36 + RPS: 1, 37 + Burst: 3, 38 + TTL: time.Minute, 39 + }, 40 + requests: 3, 41 + expectedCode: http.StatusOK, 42 + }, 43 + } 44 + for name, tt := range tests { 45 + t.Run(name, func(t *testing.T) { 46 + handler := Middleware(tt.config) 47 + nextHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 48 + w.WriteHeader(http.StatusOK) 49 + }) 50 + var lastCode int 51 + 52 + for range tt.requests { 53 + w := httptest.NewRecorder() 54 + r := httptest.NewRequest(http.MethodGet, "/", nil) 55 + 56 + handler(nextHandler).ServeHTTP(w, r) 57 + lastCode = w.Code 58 + } 59 + 60 + if lastCode != tt.expectedCode { 61 + t.Errorf("expected status %d, got %d", tt.expectedCode, lastCode) 62 + } 63 + }) 64 + } 65 +} 66 + 67 +func TestRateLimiter_getVisitor(t *testing.T) { 68 + limiter := newLimiter(10, 20, time.Second) 69 + ip := visitorIP("127.0.0.1") 70 + 71 + v := limiter.getVisitor(ip) 72 + if v == nil { 73 + t.Fatal("expected non-nil limiter") 74 + } 75 + 76 + vAgain := limiter.getVisitor(ip) 77 + if v != vAgain { 78 + t.Fatal("expected the same limiter for the same IP") 79 + } 80 + 81 + if want := 1; len(limiter.visitors) != want { 82 + t.Fatalf("expected %d visitor, got %d", want, len(limiter.visitors)) 83 + } 84 +} 85 + 86 +func TestRateLimiter_cleanupVisitors(t *testing.T) { 87 + limiter := newLimiter(10, 20, time.Millisecond) 88 + limiter.getVisitor("192.168.9.1") 89 + if want := 1; len(limiter.visitors) != want { 90 + t.Fatalf("expected %d visitor, got %d", want, len(limiter.visitors)) 91 + } 92 + 93 + time.Sleep(5 * time.Millisecond) 94 + limiter.cleanupVisitors() 95 + if want := 0; len(limiter.visitors) != want { 96 + t.Fatalf("expected %d visitors after cleanup, got %d", want, len(limiter.visitors)) 97 + } 98 +} 99 + 100 +func TestRateLimiter_differentIPs(t *testing.T) { 101 + limiter := newLimiter(10, 20, time.Second) 102 + 103 + ip1 := limiter.getVisitor("1.1.1.1") 104 + ip2 := limiter.getVisitor("2.2.2.2") 105 + if ip1 == ip2 { 106 + t.Fatal("expected different limiters for different IPs") 107 + } 108 + 109 + if want := 2; len(limiter.visitors) != want { 110 + t.Fatalf("expected %d visitors, got %d", want, len(limiter.visitors)) 111 + } 112 +} 113 + 114 +func TestGetIP(t *testing.T) { 115 + t.Run("uses X-Forwarded-For", func(t *testing.T) { 116 + r := httptest.NewRequest(http.MethodGet, "/", nil) 117 + r.Header.Set("X-Forwarded-For", "203.0.113.1") 118 + if got := getIP(r); got != "203.0.113.1" { 119 + t.Errorf("expected %q, got %q", "203.0.113.1", got) 120 + } 121 + }) 122 + 123 + t.Run("uses first IP from X-Forwarded-For list", func(t *testing.T) { 124 + r := httptest.NewRequest(http.MethodGet, "/", nil) 125 + r.Header.Set("X-Forwarded-For", "203.0.113.1, 198.51.100.2, 192.0.2.3") 126 + if got := getIP(r); got != "203.0.113.1" { 127 + t.Errorf("expected %q, got %q", "203.0.113.1", got) 128 + } 129 + }) 130 + 131 + t.Run("uses X-Real-IP when no X-Forwarded-For", func(t *testing.T) { 132 + r := httptest.NewRequest(http.MethodGet, "/", nil) 133 + r.Header.Set("X-Real-IP", "10.0.0.1") 134 + if got := getIP(r); got != "10.0.0.1" { 135 + t.Errorf("expected %q, got %q", "10.0.0.1", got) 136 + } 137 + }) 138 + 139 + t.Run("falls back to RemoteAddr", func(t *testing.T) { 140 + r := httptest.NewRequest(http.MethodGet, "/", nil) 141 + r.RemoteAddr = "192.168.1.1:12345" 142 + if got := getIP(r); got != "192.168.1.1" { 143 + t.Errorf("expected %q, got %q", "192.168.1.1", got) 144 + } 145 + }) 146 +}