112 lines
2.2 KiB
Go
112 lines
2.2 KiB
Go
package middleware
|
|
|
|
import (
|
|
"lost-and-found/internal/utils"
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
// RateLimiter stores rate limit data
|
|
type RateLimiter struct {
|
|
visitors map[string]*Visitor
|
|
mu sync.RWMutex
|
|
rate int // requests per window
|
|
window time.Duration // time window
|
|
}
|
|
|
|
// Visitor represents a visitor's rate limit data
|
|
type Visitor struct {
|
|
lastSeen time.Time
|
|
count int
|
|
}
|
|
|
|
var limiter *RateLimiter
|
|
|
|
// InitRateLimiter initializes the rate limiter
|
|
func InitRateLimiter(rate int, window time.Duration) {
|
|
limiter = &RateLimiter{
|
|
visitors: make(map[string]*Visitor),
|
|
rate: rate,
|
|
window: window,
|
|
}
|
|
|
|
// Cleanup old visitors every minute
|
|
go limiter.cleanupVisitors()
|
|
}
|
|
|
|
// cleanupVisitors removes old visitor entries
|
|
func (rl *RateLimiter) cleanupVisitors() {
|
|
for {
|
|
time.Sleep(time.Minute)
|
|
rl.mu.Lock()
|
|
for ip, visitor := range rl.visitors {
|
|
if time.Since(visitor.lastSeen) > rl.window {
|
|
delete(rl.visitors, ip)
|
|
}
|
|
}
|
|
rl.mu.Unlock()
|
|
}
|
|
}
|
|
|
|
// getVisitor retrieves or creates a visitor
|
|
func (rl *RateLimiter) getVisitor(ip string) *Visitor {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
|
|
visitor, exists := rl.visitors[ip]
|
|
if !exists {
|
|
visitor = &Visitor{
|
|
lastSeen: time.Now(),
|
|
count: 0,
|
|
}
|
|
rl.visitors[ip] = visitor
|
|
}
|
|
|
|
return visitor
|
|
}
|
|
|
|
// isAllowed checks if request is allowed
|
|
func (rl *RateLimiter) isAllowed(ip string) bool {
|
|
visitor := rl.getVisitor(ip)
|
|
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
|
|
// Reset count if window has passed
|
|
if time.Since(visitor.lastSeen) > rl.window {
|
|
visitor.count = 0
|
|
visitor.lastSeen = time.Now()
|
|
}
|
|
|
|
// Check if limit exceeded
|
|
if visitor.count >= rl.rate {
|
|
return false
|
|
}
|
|
|
|
visitor.count++
|
|
visitor.lastSeen = time.Now()
|
|
return true
|
|
}
|
|
|
|
// RateLimiterMiddleware applies rate limiting
|
|
func RateLimiterMiddleware() gin.HandlerFunc {
|
|
// Initialize rate limiter (100 requests per minute)
|
|
if limiter == nil {
|
|
InitRateLimiter(100, time.Minute)
|
|
}
|
|
|
|
return func(ctx *gin.Context) {
|
|
ip := ctx.ClientIP()
|
|
|
|
if !limiter.isAllowed(ip) {
|
|
utils.ErrorResponse(ctx, http.StatusTooManyRequests, "Rate limit exceeded", "Too many requests, please try again later")
|
|
ctx.Abort()
|
|
return
|
|
}
|
|
|
|
ctx.Next()
|
|
}
|
|
} |