84 lines
1.8 KiB
Go
84 lines
1.8 KiB
Go
package router
|
|
|
|
import (
|
|
"net/http"
|
|
"sync"
|
|
|
|
"github.com/atlasos/calypso/internal/common/config"
|
|
"github.com/atlasos/calypso/internal/common/logger"
|
|
"github.com/gin-gonic/gin"
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
// rateLimiter manages rate limiting per IP address
|
|
type rateLimiter struct {
|
|
limiters map[string]*rate.Limiter
|
|
mu sync.RWMutex
|
|
config config.RateLimitConfig
|
|
logger *logger.Logger
|
|
}
|
|
|
|
// newRateLimiter creates a new rate limiter
|
|
func newRateLimiter(cfg config.RateLimitConfig, log *logger.Logger) *rateLimiter {
|
|
return &rateLimiter{
|
|
limiters: make(map[string]*rate.Limiter),
|
|
config: cfg,
|
|
logger: log,
|
|
}
|
|
}
|
|
|
|
// getLimiter returns a rate limiter for the given IP address
|
|
func (rl *rateLimiter) getLimiter(ip string) *rate.Limiter {
|
|
rl.mu.RLock()
|
|
limiter, exists := rl.limiters[ip]
|
|
rl.mu.RUnlock()
|
|
|
|
if exists {
|
|
return limiter
|
|
}
|
|
|
|
// Create new limiter for this IP
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
|
|
// Double-check after acquiring write lock
|
|
if limiter, exists := rl.limiters[ip]; exists {
|
|
return limiter
|
|
}
|
|
|
|
// Create limiter with configured rate
|
|
limiter = rate.NewLimiter(rate.Limit(rl.config.RequestsPerSecond), rl.config.BurstSize)
|
|
rl.limiters[ip] = limiter
|
|
|
|
return limiter
|
|
}
|
|
|
|
// rateLimitMiddleware creates rate limiting middleware
|
|
func rateLimitMiddleware(cfg *config.Config, log *logger.Logger) gin.HandlerFunc {
|
|
if !cfg.Security.RateLimit.Enabled {
|
|
// Rate limiting disabled, return no-op middleware
|
|
return func(c *gin.Context) {
|
|
c.Next()
|
|
}
|
|
}
|
|
|
|
limiter := newRateLimiter(cfg.Security.RateLimit, log)
|
|
|
|
return func(c *gin.Context) {
|
|
ip := c.ClientIP()
|
|
limiter := limiter.getLimiter(ip)
|
|
|
|
if !limiter.Allow() {
|
|
log.Warn("Rate limit exceeded", "ip", ip, "path", c.Request.URL.Path)
|
|
c.JSON(http.StatusTooManyRequests, gin.H{
|
|
"error": "rate limit exceeded",
|
|
})
|
|
c.Abort()
|
|
return
|
|
}
|
|
|
|
c.Next()
|
|
}
|
|
}
|
|
|