73 lines
1.9 KiB
Go
73 lines
1.9 KiB
Go
package auth
|
|
|
|
import (
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/danielgtaylor/huma/v2"
|
|
)
|
|
|
|
// RateLimiter provides per-IP rate limiting for authenticated API endpoints,
|
|
// complementing the login-specific failLimiter with a broader cap on all
|
|
// requests. The window aligns to wall-clock intervals so all IPs share the
|
|
// same boundary.
|
|
type RateLimiter struct {
|
|
mu sync.Mutex
|
|
buckets map[string]*tokenBucket
|
|
limit int
|
|
interval time.Duration
|
|
}
|
|
|
|
type tokenBucket struct {
|
|
count int
|
|
windowEnd time.Time
|
|
}
|
|
|
|
const maxRateLimitKeys = 10000
|
|
|
|
func NewRateLimiter(limit int, interval time.Duration) *RateLimiter {
|
|
return &RateLimiter{
|
|
buckets: map[string]*tokenBucket{},
|
|
limit: limit,
|
|
interval: interval,
|
|
}
|
|
}
|
|
|
|
// Allow reports whether ip may make a request now. Returns false when the
|
|
// limit is exceeded OR the map is full (fail-closed).
|
|
func (l *RateLimiter) Allow(ip string) bool {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
|
|
now := time.Now()
|
|
b, ok := l.buckets[ip]
|
|
if !ok || now.After(b.windowEnd) {
|
|
if len(l.buckets) >= maxRateLimitKeys {
|
|
return false
|
|
}
|
|
l.buckets[ip] = &tokenBucket{count: 1, windowEnd: now.Add(l.interval)}
|
|
return true
|
|
}
|
|
if b.count >= l.limit {
|
|
return false
|
|
}
|
|
b.count++
|
|
return true
|
|
}
|
|
|
|
// RateLimitMiddleware returns Huma middleware that rejects requests exceeding
|
|
// the per-IP rate limit with 429 Too Many Requests. It runs before the RBAC
|
|
// check so abusive IPs are dropped early. The IP is read from the context set
|
|
// by WithClientIP; if absent the request passes through unthrottled.
|
|
func RateLimitMiddleware(api huma.API, rl *RateLimiter) func(huma.Context, func(huma.Context)) {
|
|
return func(ctx huma.Context, next func(huma.Context)) {
|
|
ip := ClientIP(ctx.Context())
|
|
if ip != "" && !rl.Allow(ip) {
|
|
huma.WriteErr(api, ctx, http.StatusTooManyRequests, "rate limit exceeded, try again later")
|
|
return
|
|
}
|
|
next(ctx)
|
|
}
|
|
}
|