Add rate limiting
This commit is contained in:
135
internal/ratelimit/ratelimit.go
Normal file
135
internal/ratelimit/ratelimit.go
Normal file
@@ -0,0 +1,135 @@
|
||||
package ratelimit
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
// Config defines rate limiting configuration.
|
||||
type Config struct {
|
||||
// RequestsPerSecond is the number of requests allowed per second per IP.
|
||||
RequestsPerSecond float64
|
||||
// Burst is the maximum burst size allowed.
|
||||
Burst int
|
||||
// Enabled controls whether rate limiting is active.
|
||||
Enabled bool
|
||||
}
|
||||
|
||||
// Middleware provides per-IP rate limiting using token bucket algorithm.
|
||||
type Middleware struct {
|
||||
limiters map[string]*rate.Limiter
|
||||
mu sync.RWMutex
|
||||
config Config
|
||||
logger *slog.Logger
|
||||
}
|
||||
|
||||
// New creates a new rate limiting middleware.
|
||||
func New(config Config, logger *slog.Logger) *Middleware {
|
||||
m := &Middleware{
|
||||
limiters: make(map[string]*rate.Limiter),
|
||||
config: config,
|
||||
logger: logger,
|
||||
}
|
||||
|
||||
// Start cleanup goroutine to remove old limiters
|
||||
if config.Enabled {
|
||||
go m.cleanupLimiters()
|
||||
}
|
||||
|
||||
return m
|
||||
}
|
||||
|
||||
// Handler wraps an http.Handler with rate limiting.
|
||||
func (m *Middleware) Handler(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if !m.config.Enabled {
|
||||
next.ServeHTTP(w, r)
|
||||
return
|
||||
}
|
||||
|
||||
// Extract client IP (handle X-Forwarded-For for proxies)
|
||||
ip := m.getClientIP(r)
|
||||
|
||||
limiter := m.getLimiter(ip)
|
||||
if !limiter.Allow() {
|
||||
m.logger.Warn("rate limit exceeded",
|
||||
slog.String("ip", ip),
|
||||
slog.String("path", r.URL.Path),
|
||||
)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.Header().Set("Retry-After", "1")
|
||||
w.WriteHeader(http.StatusTooManyRequests)
|
||||
w.Write([]byte(`{"error":"rate limit exceeded","message":"too many requests"}`))
|
||||
return
|
||||
}
|
||||
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
|
||||
// getLimiter returns the rate limiter for a given IP, creating one if needed.
|
||||
func (m *Middleware) getLimiter(ip string) *rate.Limiter {
|
||||
m.mu.RLock()
|
||||
limiter, exists := m.limiters[ip]
|
||||
m.mu.RUnlock()
|
||||
|
||||
if exists {
|
||||
return limiter
|
||||
}
|
||||
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
// Double-check after acquiring write lock
|
||||
limiter, exists = m.limiters[ip]
|
||||
if exists {
|
||||
return limiter
|
||||
}
|
||||
|
||||
limiter = rate.NewLimiter(rate.Limit(m.config.RequestsPerSecond), m.config.Burst)
|
||||
m.limiters[ip] = limiter
|
||||
return limiter
|
||||
}
|
||||
|
||||
// getClientIP extracts the client IP from the request.
|
||||
func (m *Middleware) getClientIP(r *http.Request) string {
|
||||
// Check X-Forwarded-For header (for proxies/load balancers)
|
||||
xff := r.Header.Get("X-Forwarded-For")
|
||||
if xff != "" {
|
||||
// X-Forwarded-For can be a comma-separated list, use the first IP
|
||||
for idx := 0; idx < len(xff); idx++ {
|
||||
if xff[idx] == ',' {
|
||||
return xff[:idx]
|
||||
}
|
||||
}
|
||||
return xff
|
||||
}
|
||||
|
||||
// Check X-Real-IP header
|
||||
if xri := r.Header.Get("X-Real-IP"); xri != "" {
|
||||
return xri
|
||||
}
|
||||
|
||||
// Fall back to RemoteAddr
|
||||
return r.RemoteAddr
|
||||
}
|
||||
|
||||
// cleanupLimiters periodically removes unused limiters to prevent memory leaks.
|
||||
func (m *Middleware) cleanupLimiters() {
|
||||
ticker := time.NewTicker(5 * time.Minute)
|
||||
defer ticker.Stop()
|
||||
|
||||
for range ticker.C {
|
||||
m.mu.Lock()
|
||||
// Clear all limiters periodically
|
||||
// In production, you might want a more sophisticated LRU cache
|
||||
m.limiters = make(map[string]*rate.Limiter)
|
||||
m.mu.Unlock()
|
||||
|
||||
m.logger.Debug("cleaned up rate limiters")
|
||||
}
|
||||
}
|
||||
175
internal/ratelimit/ratelimit_test.go
Normal file
175
internal/ratelimit/ratelimit_test.go
Normal file
@@ -0,0 +1,175 @@
|
||||
package ratelimit
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestRateLimitMiddleware(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
config Config
|
||||
requestCount int
|
||||
expectedAllowed int
|
||||
expectedRateLimited int
|
||||
}{
|
||||
{
|
||||
name: "disabled rate limiting allows all requests",
|
||||
config: Config{
|
||||
Enabled: false,
|
||||
RequestsPerSecond: 1,
|
||||
Burst: 1,
|
||||
},
|
||||
requestCount: 10,
|
||||
expectedAllowed: 10,
|
||||
expectedRateLimited: 0,
|
||||
},
|
||||
{
|
||||
name: "enabled rate limiting enforces limits",
|
||||
config: Config{
|
||||
Enabled: true,
|
||||
RequestsPerSecond: 1,
|
||||
Burst: 2,
|
||||
},
|
||||
requestCount: 5,
|
||||
expectedAllowed: 2, // Burst allows 2 immediately
|
||||
expectedRateLimited: 3,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
middleware := New(tt.config, logger)
|
||||
|
||||
handler := middleware.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
|
||||
allowed := 0
|
||||
rateLimited := 0
|
||||
|
||||
for i := 0; i < tt.requestCount; i++ {
|
||||
req := httptest.NewRequest("GET", "/test", nil)
|
||||
req.RemoteAddr = "192.168.1.1:1234"
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
handler.ServeHTTP(w, req)
|
||||
|
||||
if w.Code == http.StatusOK {
|
||||
allowed++
|
||||
} else if w.Code == http.StatusTooManyRequests {
|
||||
rateLimited++
|
||||
}
|
||||
}
|
||||
|
||||
if allowed != tt.expectedAllowed {
|
||||
t.Errorf("expected %d allowed requests, got %d", tt.expectedAllowed, allowed)
|
||||
}
|
||||
if rateLimited != tt.expectedRateLimited {
|
||||
t.Errorf("expected %d rate limited requests, got %d", tt.expectedRateLimited, rateLimited)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetClientIP(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
middleware := New(Config{Enabled: false}, logger)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
headers map[string]string
|
||||
remoteAddr string
|
||||
expectedIP string
|
||||
}{
|
||||
{
|
||||
name: "uses X-Forwarded-For if present",
|
||||
headers: map[string]string{"X-Forwarded-For": "203.0.113.1, 198.51.100.1"},
|
||||
remoteAddr: "192.168.1.1:1234",
|
||||
expectedIP: "203.0.113.1",
|
||||
},
|
||||
{
|
||||
name: "uses X-Real-IP if X-Forwarded-For not present",
|
||||
headers: map[string]string{"X-Real-IP": "203.0.113.1"},
|
||||
remoteAddr: "192.168.1.1:1234",
|
||||
expectedIP: "203.0.113.1",
|
||||
},
|
||||
{
|
||||
name: "uses RemoteAddr as fallback",
|
||||
headers: map[string]string{},
|
||||
remoteAddr: "192.168.1.1:1234",
|
||||
expectedIP: "192.168.1.1:1234",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
req := httptest.NewRequest("GET", "/test", nil)
|
||||
req.RemoteAddr = tt.remoteAddr
|
||||
for k, v := range tt.headers {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
|
||||
ip := middleware.getClientIP(req)
|
||||
if ip != tt.expectedIP {
|
||||
t.Errorf("expected IP %q, got %q", tt.expectedIP, ip)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimitRefill(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
|
||||
config := Config{
|
||||
Enabled: true,
|
||||
RequestsPerSecond: 10, // 10 requests per second
|
||||
Burst: 5,
|
||||
}
|
||||
middleware := New(config, logger)
|
||||
|
||||
handler := middleware.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
|
||||
// Use up the burst
|
||||
for i := 0; i < 5; i++ {
|
||||
req := httptest.NewRequest("GET", "/test", nil)
|
||||
req.RemoteAddr = "192.168.1.1:1234"
|
||||
w := httptest.NewRecorder()
|
||||
handler.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Errorf("request %d should be allowed, got status %d", i, w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// Next request should be rate limited
|
||||
req := httptest.NewRequest("GET", "/test", nil)
|
||||
req.RemoteAddr = "192.168.1.1:1234"
|
||||
w := httptest.NewRecorder()
|
||||
handler.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusTooManyRequests {
|
||||
t.Errorf("expected rate limit, got status %d", w.Code)
|
||||
}
|
||||
|
||||
// Wait for tokens to refill (100ms = 1 token at 10/s)
|
||||
time.Sleep(150 * time.Millisecond)
|
||||
|
||||
// Should be allowed now
|
||||
req = httptest.NewRequest("GET", "/test", nil)
|
||||
req.RemoteAddr = "192.168.1.1:1234"
|
||||
w = httptest.NewRecorder()
|
||||
handler.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Errorf("request should be allowed after refill, got status %d", w.Code)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user