hobokenchicken
1d6c9bdfe6
- ipLimiter evicts stale entries after 10 min of inactivity via 5 min sweep - DB pool capped at 25 max open, 5 idle, 15 min lifetime
91 lines
2.1 KiB
Go
91 lines
2.1 KiB
Go
package middleware
|
|
|
|
import (
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
// ipLimiter tracks rate limiters per IP address.
|
|
type ipLimiter struct {
|
|
mu sync.Mutex
|
|
limiters map[string]*rate.Limiter
|
|
lastUsed map[string]time.Time
|
|
rate rate.Limit
|
|
burst int
|
|
}
|
|
|
|
func newIPLimiter(r rate.Limit, burst int) *ipLimiter {
|
|
l := &ipLimiter{
|
|
limiters: make(map[string]*rate.Limiter),
|
|
lastUsed: make(map[string]time.Time),
|
|
rate: r,
|
|
burst: burst,
|
|
}
|
|
// ponytail: evict inactive limiters every 5 minutes to prevent unbounded map growth
|
|
go l.cleanup()
|
|
return l
|
|
}
|
|
|
|
func (l *ipLimiter) getLimiter(ip string) *rate.Limiter {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
|
|
lim, exists := l.limiters[ip]
|
|
if !exists {
|
|
lim = rate.NewLimiter(l.rate, l.burst)
|
|
l.limiters[ip] = lim
|
|
}
|
|
l.lastUsed[ip] = time.Now()
|
|
return lim
|
|
}
|
|
|
|
// cleanup runs in a background goroutine and removes limiters that haven't
|
|
// been accessed in 10 minutes.
|
|
func (l *ipLimiter) cleanup() {
|
|
for {
|
|
time.Sleep(5 * time.Minute)
|
|
l.mu.Lock()
|
|
now := time.Now()
|
|
for ip, last := range l.lastUsed {
|
|
if now.Sub(last) > 10*time.Minute {
|
|
delete(l.limiters, ip)
|
|
delete(l.lastUsed, ip)
|
|
}
|
|
}
|
|
l.mu.Unlock()
|
|
}
|
|
}
|
|
|
|
// RateLimit returns middleware that limits requests per IP.
|
|
// requestsPerSecond is the sustained rate, burst is the max burst size.
|
|
func RateLimit(requestsPerSecond float64, burst int) func(http.Handler) http.Handler {
|
|
limiter := newIPLimiter(rate.Limit(requestsPerSecond), burst)
|
|
|
|
return func(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
ip := r.RemoteAddr
|
|
if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
|
|
// Take the first IP in the chain (the original client)
|
|
if idx := strings.IndexByte(xff, ','); idx > 0 {
|
|
ip = strings.TrimSpace(xff[:idx])
|
|
} else {
|
|
ip = strings.TrimSpace(xff)
|
|
}
|
|
} else if xri := r.Header.Get("X-Real-IP"); xri != "" {
|
|
ip = strings.TrimSpace(xri)
|
|
}
|
|
|
|
if !limiter.getLimiter(ip).Allow() {
|
|
http.Error(w, `{"error":"rate limited"}`, http.StatusTooManyRequests)
|
|
return
|
|
}
|
|
|
|
next.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
}
|