78 lines
1.6 KiB
Go
78 lines
1.6 KiB
Go
package middleware
|
|
|
|
import (
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
"users_management/m/utils/common"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
var (
|
|
rateLimiters = make(map[string]*rate.Limiter)
|
|
mu sync.Mutex
|
|
)
|
|
|
|
func getRateLimiter(userID string) *rate.Limiter {
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
|
|
limiter, exists := rateLimiters[userID]
|
|
if !exists {
|
|
limiter = rate.NewLimiter(rate.Every(1*time.Minute), 50) // 1 request per second with a burst of 2 requests
|
|
rateLimiters[userID] = limiter
|
|
}
|
|
|
|
return limiter
|
|
}
|
|
|
|
func getLoginLimiter() *rate.Limiter {
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
|
|
limiter, exists := rateLimiters["login"]
|
|
if !exists {
|
|
limiter = rate.NewLimiter(rate.Every(1*time.Minute), 4) // 5 request per second with a burst of 10 requests
|
|
rateLimiters["login"] = limiter
|
|
}
|
|
|
|
return limiter
|
|
}
|
|
|
|
func RateLimitMiddleware() gin.HandlerFunc{
|
|
return func(c *gin.Context) {
|
|
userID, exists := c.Get("userID")
|
|
if !exists {
|
|
common.ErrorResponses(c, http.StatusUnauthorized, "Unauthorized: No user ID found")
|
|
c.Abort()
|
|
return
|
|
}
|
|
|
|
limiter := getRateLimiter(userID.(string))
|
|
|
|
if !limiter.Allow() {
|
|
common.ErrorResponses(c, http.StatusTooManyRequests, "Too many requests")
|
|
c.Abort()
|
|
return
|
|
}
|
|
|
|
c.Next()
|
|
}
|
|
}
|
|
|
|
func RateLoginMiddleware() gin.HandlerFunc{
|
|
return func(c *gin.Context) {
|
|
limiter := getLoginLimiter()
|
|
|
|
if !limiter.Allow() {
|
|
common.ErrorResponses(c, http.StatusTooManyRequests, "Too many requests")
|
|
c.Abort()
|
|
return
|
|
}
|
|
c.Next()
|
|
}
|
|
}
|
|
|