110 lines
3.1 KiB
Go
110 lines
3.1 KiB
Go
package middleware
|
|
|
|
import (
|
|
"container/list"
|
|
"errors"
|
|
"math"
|
|
"strconv"
|
|
"sync"
|
|
"time"
|
|
|
|
apierrors "git.yuntuo.hk/k2safe/go-omni/errors"
|
|
"git.yuntuo.hk/k2safe/go-omni/httpx"
|
|
)
|
|
|
|
type RateLimitConfig struct {
|
|
Rate float64 `json:"rate"` // Tokens per second, per key.
|
|
Burst int `json:"burst"`
|
|
MaxKeys int `json:"maxKeys"`
|
|
IdleTTL time.Duration `json:"idleTtl"`
|
|
Key func(*httpx.Context) string `json:"-"` // Default: trusted ClientIP.
|
|
}
|
|
|
|
type bucket struct {
|
|
key string
|
|
tokens float64
|
|
updated time.Time
|
|
}
|
|
type RateLimiter struct {
|
|
mu sync.Mutex
|
|
config RateLimitConfig
|
|
items map[string]*list.Element
|
|
lru *list.List
|
|
now func() time.Time
|
|
}
|
|
|
|
func NewRateLimiter(cfg RateLimitConfig) (*RateLimiter, error) {
|
|
if math.IsNaN(cfg.Rate) || math.IsInf(cfg.Rate, 0) || cfg.Rate <= 0 || cfg.Burst < 1 || cfg.MaxKeys < 0 || cfg.IdleTTL < 0 {
|
|
return nil, errors.New("middleware: invalid rate limit configuration")
|
|
}
|
|
refill := float64(cfg.Burst) / cfg.Rate
|
|
if refill > (24 * time.Hour).Seconds() {
|
|
return nil, errors.New("middleware: bucket refill time must not exceed 24h")
|
|
}
|
|
minimum := time.Duration(math.Ceil(refill * float64(time.Second)))
|
|
if cfg.IdleTTL == 0 {
|
|
cfg.IdleTTL = max(5*time.Minute, minimum)
|
|
}
|
|
if cfg.IdleTTL < minimum {
|
|
return nil, errors.New("middleware: idleTtl must allow a full bucket refill")
|
|
}
|
|
if cfg.MaxKeys == 0 {
|
|
cfg.MaxKeys = 10000
|
|
}
|
|
if cfg.Key == nil {
|
|
cfg.Key = func(c *httpx.Context) string { return c.ClientIP() }
|
|
}
|
|
return &RateLimiter{config: cfg, items: make(map[string]*list.Element), lru: list.New(), now: time.Now}, nil
|
|
}
|
|
|
|
func (l *RateLimiter) Middleware() httpx.Middleware {
|
|
return func(next httpx.Handler) httpx.Handler {
|
|
return func(c *httpx.Context) error {
|
|
key := l.config.Key(c)
|
|
if key == "" || len(key) > 1024 {
|
|
return apierrors.BadRequest("INVALID_RATE_KEY", "invalid rate limit key")
|
|
}
|
|
delay, full := l.allow(key)
|
|
if full {
|
|
return apierrors.New(503, "RATE_LIMIT_CAPACITY", "rate limiter capacity reached")
|
|
}
|
|
if delay > 0 {
|
|
c.SetResponseHeader("Retry-After", strconv.Itoa(max(1, int(math.Ceil(delay.Seconds())))))
|
|
return apierrors.New(429, "RATE_LIMITED", "too many requests")
|
|
}
|
|
return next(c)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (l *RateLimiter) allow(key string) (time.Duration, bool) {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
now := l.now()
|
|
e := l.items[key]
|
|
if e == nil {
|
|
for l.lru.Len() > 0 {
|
|
old := l.lru.Back()
|
|
if now.Sub(old.Value.(*bucket).updated) < l.config.IdleTTL {
|
|
break
|
|
}
|
|
delete(l.items, old.Value.(*bucket).key)
|
|
l.lru.Remove(old)
|
|
}
|
|
if len(l.items) >= l.config.MaxKeys {
|
|
return 0, true
|
|
}
|
|
e = l.lru.PushFront(&bucket{key: key, tokens: float64(l.config.Burst), updated: now})
|
|
l.items[key] = e
|
|
}
|
|
b := e.Value.(*bucket)
|
|
b.tokens = min(float64(l.config.Burst), b.tokens+max(0, now.Sub(b.updated).Seconds())*l.config.Rate)
|
|
b.updated = now
|
|
l.lru.MoveToFront(e)
|
|
if b.tokens >= 1 {
|
|
b.tokens--
|
|
return 0, false
|
|
}
|
|
return time.Duration(math.Ceil((1 - b.tokens) / l.config.Rate * float64(time.Second))), false
|
|
}
|