139 lines
3.1 KiB
Go
139 lines
3.1 KiB
Go
package httpserver
|
|
|
|
import (
|
|
"math"
|
|
"net/http"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type rateLimitStore struct {
|
|
mu sync.Mutex
|
|
ratePerSec float64
|
|
burst float64
|
|
lastCleanup time.Time
|
|
entries map[string]*rateLimitEntry
|
|
}
|
|
|
|
type rateLimitEntry struct {
|
|
tokens float64
|
|
lastSeen time.Time
|
|
}
|
|
|
|
func newRateLimitStore(perMinute, burst int) *rateLimitStore {
|
|
return &rateLimitStore{
|
|
ratePerSec: float64(perMinute) / 60,
|
|
burst: float64(burst),
|
|
entries: make(map[string]*rateLimitEntry),
|
|
}
|
|
}
|
|
|
|
func (s *rateLimitStore) Allow(key string, now time.Time) (bool, time.Duration) {
|
|
key = strings.TrimSpace(key)
|
|
if key == "" {
|
|
key = "global"
|
|
}
|
|
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
if entry, ok := s.entries[key]; ok {
|
|
elapsed := now.Sub(entry.lastSeen).Seconds()
|
|
entry.tokens = math.Min(s.burst, entry.tokens+(elapsed*s.ratePerSec))
|
|
entry.lastSeen = now
|
|
if entry.tokens >= 1 {
|
|
entry.tokens--
|
|
s.cleanup(now)
|
|
return true, 0
|
|
}
|
|
|
|
s.cleanup(now)
|
|
return false, retryAfter(entry.tokens, s.ratePerSec)
|
|
}
|
|
|
|
s.entries[key] = &rateLimitEntry{
|
|
tokens: s.burst - 1,
|
|
lastSeen: now,
|
|
}
|
|
s.cleanup(now)
|
|
return true, 0
|
|
}
|
|
|
|
func (s *rateLimitStore) cleanup(now time.Time) {
|
|
if !s.lastCleanup.IsZero() && now.Sub(s.lastCleanup) < 5*time.Minute {
|
|
return
|
|
}
|
|
|
|
for key, entry := range s.entries {
|
|
if now.Sub(entry.lastSeen) > 15*time.Minute {
|
|
delete(s.entries, key)
|
|
}
|
|
}
|
|
|
|
s.lastCleanup = now
|
|
}
|
|
|
|
func retryAfter(tokens, ratePerSec float64) time.Duration {
|
|
if ratePerSec <= 0 {
|
|
return time.Minute
|
|
}
|
|
|
|
if tokens < 0 {
|
|
tokens = 0
|
|
}
|
|
|
|
seconds := (1 - tokens) / ratePerSec
|
|
if seconds < 1 {
|
|
seconds = 1
|
|
}
|
|
|
|
return time.Duration(math.Ceil(seconds * float64(time.Second)))
|
|
}
|
|
|
|
func retryAfterHeader(delay time.Duration) string {
|
|
seconds := int(math.Ceil(delay.Seconds()))
|
|
if seconds < 1 {
|
|
seconds = 1
|
|
}
|
|
|
|
return strconv.Itoa(seconds)
|
|
}
|
|
|
|
func (h *handler) loginRateLimit(next http.Handler) http.Handler {
|
|
return h.rateLimit(h.loginRateLimiter, func(r *http.Request) string {
|
|
return clientIP(r)
|
|
}, func(w http.ResponseWriter, r *http.Request, delay time.Duration) {
|
|
w.Header().Set("Retry-After", retryAfterHeader(delay))
|
|
http.Error(w, "too many login attempts", http.StatusTooManyRequests)
|
|
})(next)
|
|
}
|
|
|
|
func (h *handler) clientAPIRateLimit(next http.Handler) http.Handler {
|
|
return h.rateLimit(h.clientRateLimiter, func(r *http.Request) string {
|
|
return clientIP(r)
|
|
}, func(w http.ResponseWriter, r *http.Request, delay time.Duration) {
|
|
w.Header().Set("Retry-After", retryAfterHeader(delay))
|
|
writeJSON(w, http.StatusTooManyRequests, map[string]any{"error": "rate limit exceeded"})
|
|
})(next)
|
|
}
|
|
|
|
func (h *handler) rateLimit(store *rateLimitStore, keyFn func(*http.Request) string, reject func(http.ResponseWriter, *http.Request, time.Duration)) func(http.Handler) http.Handler {
|
|
return func(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
key := clientIP(r)
|
|
if keyFn != nil {
|
|
key = keyFn(r)
|
|
}
|
|
|
|
allowed, delay := store.Allow(key, time.Now().UTC())
|
|
if !allowed {
|
|
reject(w, r, delay)
|
|
return
|
|
}
|
|
|
|
next.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
}
|