141 lines
2.9 KiB
Go
141 lines
2.9 KiB
Go
package middleware
|
|
|
|
import (
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
type tokenBucket struct {
|
|
mu sync.Mutex
|
|
rate float64 // tokens per second
|
|
burst float64 // max tokens
|
|
tokens float64
|
|
last time.Time
|
|
}
|
|
|
|
func newTokenBucket(max int, per time.Duration, burst int) *tokenBucket {
|
|
if burst <= 0 {
|
|
burst = max
|
|
}
|
|
rate := float64(max) / per.Seconds()
|
|
b := float64(burst)
|
|
now := time.Now()
|
|
return &tokenBucket{rate: rate, burst: b, tokens: b, last: now}
|
|
}
|
|
|
|
func (b *tokenBucket) allow() bool {
|
|
b.mu.Lock()
|
|
defer b.mu.Unlock()
|
|
|
|
now := time.Now()
|
|
delta := now.Sub(b.last).Seconds()
|
|
b.last = now
|
|
|
|
b.tokens += delta * b.rate
|
|
if b.tokens > b.burst {
|
|
b.tokens = b.burst
|
|
}
|
|
|
|
if b.tokens < 1 {
|
|
return false
|
|
}
|
|
b.tokens -= 1
|
|
return true
|
|
}
|
|
|
|
type ipClient struct {
|
|
bucket *tokenBucket
|
|
max int
|
|
burst int
|
|
lastSeen time.Time
|
|
}
|
|
|
|
// RateLimitByIP returns a Gin middleware that rate limits requests per client IP.
|
|
//
|
|
// max: max requests per time window (per)
|
|
// per: the time window duration
|
|
// burst: optional burst capacity (defaults to max if <=0)
|
|
// ttl: how long to keep idle IP buckets around
|
|
func RateLimitByIP(max int, per time.Duration, burst int, ttl time.Duration) gin.HandlerFunc {
|
|
return RateLimitByIPDynamic(func() int { return max }, per, func() int { return burst }, ttl)
|
|
}
|
|
|
|
// RateLimitByIPDynamic is like RateLimitByIP but reads max/burst dynamically.
|
|
// This allows changing limits at runtime (e.g. from an admin config page).
|
|
func RateLimitByIPDynamic(maxFn func() int, per time.Duration, burstFn func() int, ttl time.Duration) gin.HandlerFunc {
|
|
var (
|
|
mu sync.Mutex
|
|
clients = make(map[string]*ipClient)
|
|
)
|
|
|
|
// opportunistic cleanup (runs at most once per minute)
|
|
var (
|
|
cleanupMu sync.Mutex
|
|
lastCleanup time.Time
|
|
)
|
|
|
|
cleanup := func(now time.Time) {
|
|
cleanupMu.Lock()
|
|
defer cleanupMu.Unlock()
|
|
if !lastCleanup.IsZero() && now.Sub(lastCleanup) < time.Minute {
|
|
return
|
|
}
|
|
lastCleanup = now
|
|
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
for ip, c := range clients {
|
|
if now.Sub(c.lastSeen) > ttl {
|
|
delete(clients, ip)
|
|
}
|
|
}
|
|
}
|
|
|
|
getClient := func(ip string, now time.Time, max int, burst int) *ipClient {
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
c, ok := clients[ip]
|
|
if !ok {
|
|
c = &ipClient{bucket: newTokenBucket(max, per, burst), max: max, burst: burst, lastSeen: now}
|
|
clients[ip] = c
|
|
return c
|
|
}
|
|
c.lastSeen = now
|
|
// If settings changed, reset the bucket.
|
|
if c.max != max || c.burst != burst {
|
|
c.bucket = newTokenBucket(max, per, burst)
|
|
c.max = max
|
|
c.burst = burst
|
|
}
|
|
return c
|
|
}
|
|
|
|
return func(c *gin.Context) {
|
|
now := time.Now()
|
|
cleanup(now)
|
|
|
|
ip := c.ClientIP()
|
|
max := maxFn()
|
|
if max <= 0 {
|
|
max = 1
|
|
}
|
|
burst := burstFn()
|
|
if burst <= 0 {
|
|
burst = max
|
|
}
|
|
client := getClient(ip, now, max, burst)
|
|
|
|
if !client.bucket.allow() {
|
|
c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
|
|
"error": "rate limit exceeded",
|
|
})
|
|
return
|
|
}
|
|
|
|
c.Next()
|
|
}
|
|
}
|