70 lines
1.2 KiB
Go
70 lines
1.2 KiB
Go
|
package middleware
|
||
|
|
|
||
|
|
import (
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/gin-gonic/gin"
|
||
|
|
)
|
||
|
|
|
||
|
|
type RateLimiter struct {
|
||
|
|
requests map[string][]time.Time
|
||
|
|
mutex sync.RWMutex
|
||
|
|
limit int
|
||
|
|
window time.Duration
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewRateLimiter(limit int, window time.Duration) *RateLimiter {
|
||
|
|
return &RateLimiter{
|
||
|
|
requests: make(map[string][]time.Time),
|
||
|
|
limit: limit,
|
||
|
|
window: window,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rl *RateLimiter) Allow(key string) bool {
|
||
|
|
rl.mutex.Lock()
|
||
|
|
defer rl.mutex.Unlock()
|
||
|
|
|
||
|
|
now := time.Now()
|
||
|
|
windowStart := now.Add(-rl.window)
|
||
|
|
|
||
|
|
// Clean old requests
|
||
|
|
if times, exists := rl.requests[key]; exists {
|
||
|
|
var validTimes []time.Time
|
||
|
|
for _, t := range times {
|
||
|
|
if t.After(windowStart) {
|
||
|
|
validTimes = append(validTimes, t)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
rl.requests[key] = validTimes
|
||
|
|
}
|
||
|
|
|
||
|
|
// Check if limit exceeded
|
||
|
|
if len(rl.requests[key]) >= rl.limit {
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
|
||
|
|
// Add current request
|
||
|
|
rl.requests[key] = append(rl.requests[key], now)
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
|
||
|
|
func RateLimit() gin.HandlerFunc {
|
||
|
|
limiter := NewRateLimiter(100, time.Minute) // 100 requests per minute
|
||
|
|
|
||
|
|
return gin.HandlerFunc(func(c *gin.Context) {
|
||
|
|
clientIP := c.ClientIP()
|
||
|
|
|
||
|
|
if !limiter.Allow(clientIP) {
|
||
|
|
c.JSON(429, gin.H{
|
||
|
|
"error": "Rate limit exceeded",
|
||
|
|
})
|
||
|
|
c.Abort()
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
c.Next()
|
||
|
|
})
|
||
|
|
}
|