65 lines
1.1 KiB
Go
65 lines
1.1 KiB
Go
package middleware
|
|
|
|
import (
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
type RateLimiter struct {
|
|
mu sync.Mutex
|
|
clients map[string]*clientBuckets
|
|
rate int
|
|
burst int
|
|
}
|
|
|
|
type clientBuckets struct {
|
|
tokens int
|
|
lastFill time.Time
|
|
}
|
|
|
|
func NewRateLimiter(rate, burst int) *RateLimiter {
|
|
return &RateLimiter{
|
|
clients: make(map[string]*clientBuckets),
|
|
rate: rate,
|
|
burst: burst,
|
|
}
|
|
}
|
|
|
|
func (rl *RateLimiter) Handler() gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
if !rl.allow(c.ClientIP()) {
|
|
c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{"error": "rate limit exceeded"})
|
|
return
|
|
}
|
|
c.Next()
|
|
}
|
|
}
|
|
|
|
func (rl *RateLimiter) allow(key string) bool {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
|
|
b, ok := rl.clients[key]
|
|
if !ok {
|
|
b = &clientBuckets{tokens: rl.burst, lastFill: time.Now()}
|
|
rl.clients[key] = b
|
|
}
|
|
|
|
now := time.Now()
|
|
elapsed := now.Sub(b.lastFill)
|
|
b.lastFill = now
|
|
b.tokens += int(elapsed.Seconds()) * rl.rate
|
|
if b.tokens > rl.burst {
|
|
b.tokens = rl.burst
|
|
}
|
|
|
|
if b.tokens > 0 {
|
|
b.tokens--
|
|
return true
|
|
}
|
|
return false
|
|
}
|