5 files changed,
315 insertions(+),
0 deletions(-)
Author:
Oleksandr Smirnov
olexsmir@gmail.com
Committed at:
2026-06-23 16:45:27 +0300
Authored at:
2026-06-23 16:38:49 +0300
Change ID:
zrtuwswyxnztrtkwtwkxuvvnuruuqqvr
Parent:
da217b4
jump to
| A | ratelimit/go.mod |
| A | ratelimit/go.sum |
| A | ratelimit/ratelimit.go |
| A | ratelimit/ratelimit_test.go |
| M | readme |
A
ratelimit/go.sum
··· 1 +golang.org/x/time v0.14.0 h1:MRx4UaLrDotUKUdCIqzPC48t1Y9hANFKIRpNx+Te8PI= 2 +golang.org/x/time v0.14.0/go.mod h1:eL/Oa2bBBK0TkX57Fyni+NgnyQQN4LitPmob2Hjnqw4=
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 + "golang.org/x/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 +}