114 lines
2.3 KiB
Go
114 lines
2.3 KiB
Go
|
package middleware
|
||
|
|
|
||
|
|
import (
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/gin-gonic/gin"
|
||
|
|
)
|
||
|
|
|
||
|
|
// TokenBucket implements a token bucket rate limiter per key (e.g. IP or device ID).
|
||
|
|
type TokenBucket struct {
|
||
|
|
mu sync.Mutex
|
||
|
|
buckets map[string]*bucket
|
||
|
|
rate float64 // tokens per second
|
||
|
|
capacity float64
|
||
|
|
}
|
||
|
|
|
||
|
|
type bucket struct {
|
||
|
|
tokens float64
|
||
|
|
lastFill time.Time
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewTokenBucket(rate, capacity int) *TokenBucket {
|
||
|
|
return &TokenBucket{
|
||
|
|
buckets: make(map[string]*bucket),
|
||
|
|
rate: float64(rate),
|
||
|
|
capacity: float64(capacity),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (tb *TokenBucket) Allow(key string) bool {
|
||
|
|
tb.mu.Lock()
|
||
|
|
defer tb.mu.Unlock()
|
||
|
|
|
||
|
|
b, ok := tb.buckets[key]
|
||
|
|
if !ok {
|
||
|
|
b = &bucket{tokens: tb.capacity, lastFill: time.Now()}
|
||
|
|
tb.buckets[key] = b
|
||
|
|
}
|
||
|
|
|
||
|
|
now := time.Now()
|
||
|
|
elapsed := now.Sub(b.lastFill).Seconds()
|
||
|
|
b.tokens += elapsed * tb.rate
|
||
|
|
if b.tokens > tb.capacity {
|
||
|
|
b.tokens = tb.capacity
|
||
|
|
}
|
||
|
|
b.lastFill = now
|
||
|
|
|
||
|
|
if b.tokens >= 1 {
|
||
|
|
b.tokens--
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
|
||
|
|
// Cleanup removes stale buckets older than the given duration.
|
||
|
|
func (tb *TokenBucket) Cleanup(maxAge time.Duration) {
|
||
|
|
tb.mu.Lock()
|
||
|
|
defer tb.mu.Unlock()
|
||
|
|
|
||
|
|
now := time.Now()
|
||
|
|
for k, b := range tb.buckets {
|
||
|
|
if now.Sub(b.lastFill) > maxAge {
|
||
|
|
delete(tb.buckets, k)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// RateLimitMiddleware returns a Gin middleware that rate-limits requests per client IP.
|
||
|
|
func RateLimitMiddleware(tb *TokenBucket) gin.HandlerFunc {
|
||
|
|
return func(c *gin.Context) {
|
||
|
|
ip := c.ClientIP()
|
||
|
|
if !tb.Allow(ip) {
|
||
|
|
c.AbortWithStatusJSON(429, gin.H{"error": "Rate limit exceeded. Please slow down."})
|
||
|
|
return
|
||
|
|
}
|
||
|
|
c.Next()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// GatedSemaphore limits the number of concurrent operations (e.g. downloads).
|
||
|
|
type GatedSemaphore struct {
|
||
|
|
ch chan struct{}
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewGatedSemaphore(maxConcurrent int) *GatedSemaphore {
|
||
|
|
return &GatedSemaphore{ch: make(chan struct{}, maxConcurrent)}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (gs *GatedSemaphore) Acquire() {
|
||
|
|
gs.ch <- struct{}{}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (gs *GatedSemaphore) Release() {
|
||
|
|
<-gs.ch
|
||
|
|
}
|
||
|
|
|
||
|
|
func (gs *GatedSemaphore) Available() int {
|
||
|
|
return cap(gs.ch) - len(gs.ch)
|
||
|
|
}
|
||
|
|
|
||
|
|
// GatedMiddleware only allows requests through if the semaphore has capacity.
|
||
|
|
func GatedMiddleware(gs *GatedSemaphore) gin.HandlerFunc {
|
||
|
|
return func(c *gin.Context) {
|
||
|
|
select {
|
||
|
|
case gs.ch <- struct{}{}:
|
||
|
|
defer func() { <-gs.ch }()
|
||
|
|
c.Next()
|
||
|
|
default:
|
||
|
|
c.AbortWithStatusJSON(503, gin.H{"error": "Server at capacity. Please retry later."})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|