john-okeefe
4b8cb58c84
- Add TestMode, RateLimitEnabled, RequestsPerMinute to Config - Add getEnvBool() and getEnvInt() helper functions - Update rate limiter to support enabled/disabled state - Pass test environment variables through docker-compose - Configure rate limiter dynamically in main.go This allows disabling rate limiting for integration testing while maintaining security in production environments.
124 lines
2.6 KiB
Go
124 lines
2.6 KiB
Go
package middleware
|
|
|
|
import (
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/labstack/echo/v4"
|
|
)
|
|
|
|
// RateLimiterConfig defines rate limiting configuration
|
|
type RateLimiterConfig struct {
|
|
Enabled bool
|
|
RequestsPerMinute int
|
|
CleanupInterval time.Duration
|
|
}
|
|
|
|
// DefaultRateLimiterConfig returns sensible defaults
|
|
func DefaultRateLimiterConfig() RateLimiterConfig {
|
|
return RateLimiterConfig{
|
|
Enabled: true,
|
|
RequestsPerMinute: 10,
|
|
CleanupInterval: 5 * time.Minute,
|
|
}
|
|
}
|
|
|
|
// RateLimiter tracks request counts per IP
|
|
type RateLimiter struct {
|
|
mu sync.RWMutex
|
|
requests map[string][]time.Time
|
|
config RateLimiterConfig
|
|
}
|
|
|
|
// NewRateLimiter creates a new rate limiter
|
|
func NewRateLimiter(config RateLimiterConfig) *RateLimiter {
|
|
rl := &RateLimiter{
|
|
requests: make(map[string][]time.Time),
|
|
config: config,
|
|
}
|
|
|
|
// Start cleanup goroutine
|
|
go rl.cleanup()
|
|
|
|
return rl
|
|
}
|
|
|
|
// cleanup removes old request records
|
|
func (rl *RateLimiter) cleanup() {
|
|
ticker := time.NewTicker(rl.config.CleanupInterval)
|
|
defer ticker.Stop()
|
|
|
|
for range ticker.C {
|
|
rl.mu.Lock()
|
|
now := time.Now()
|
|
for ip, requests := range rl.requests {
|
|
// Remove requests older than 1 minute
|
|
var valid []time.Time
|
|
for _, reqTime := range requests {
|
|
if now.Sub(reqTime) < time.Minute {
|
|
valid = append(valid, reqTime)
|
|
}
|
|
}
|
|
if len(valid) == 0 {
|
|
delete(rl.requests, ip)
|
|
} else {
|
|
rl.requests[ip] = valid
|
|
}
|
|
}
|
|
rl.mu.Unlock()
|
|
}
|
|
}
|
|
|
|
// Allow checks if a request from the given IP should be allowed
|
|
func (rl *RateLimiter) Allow(ip string) bool {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
|
|
now := time.Now()
|
|
|
|
// Clean old requests for this IP
|
|
var valid []time.Time
|
|
for _, reqTime := range rl.requests[ip] {
|
|
if now.Sub(reqTime) < time.Minute {
|
|
valid = append(valid, reqTime)
|
|
}
|
|
}
|
|
rl.requests[ip] = valid
|
|
|
|
// Check if limit exceeded
|
|
if len(rl.requests[ip]) >= rl.config.RequestsPerMinute {
|
|
return false
|
|
}
|
|
|
|
// Add this request
|
|
rl.requests[ip] = append(rl.requests[ip], now)
|
|
return true
|
|
}
|
|
|
|
// RateLimiterMiddleware returns echo middleware for rate limiting
|
|
func RateLimiterMiddleware(rl *RateLimiter) echo.MiddlewareFunc {
|
|
return func(next echo.HandlerFunc) echo.HandlerFunc {
|
|
return func(c echo.Context) error {
|
|
// If rate limiting is disabled, skip checks
|
|
if !rl.config.Enabled {
|
|
return next(c)
|
|
}
|
|
|
|
// Get client IP
|
|
ip := c.RealIP()
|
|
if ip == "" {
|
|
ip = c.Request().RemoteAddr
|
|
}
|
|
|
|
if !rl.Allow(ip) {
|
|
return c.JSON(http.StatusTooManyRequests, map[string]string{
|
|
"error": "too many requests, please try again later",
|
|
})
|
|
}
|
|
|
|
return next(c)
|
|
}
|
|
}
|
|
}
|