141 lines
2.8 KiB
Go
141 lines
2.8 KiB
Go
package middleware
|
|
|
|
import (
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/wmms/wmms/pkg/response"
|
|
)
|
|
|
|
const (
|
|
maxAttempts = 5
|
|
windowDuration = 15 * time.Minute
|
|
cleanupInterval = 30 * time.Minute
|
|
)
|
|
|
|
// LoginAttempt tracks failed login attempts for an IP
|
|
type LoginAttempt struct {
|
|
Count int
|
|
FirstTime time.Time
|
|
}
|
|
|
|
// RateLimiter manages login attempt tracking per IP
|
|
type RateLimiter struct {
|
|
attempts map[string]*LoginAttempt
|
|
mu sync.RWMutex
|
|
}
|
|
|
|
var loginRateLimiter = &RateLimiter{
|
|
attempts: make(map[string]*LoginAttempt),
|
|
}
|
|
|
|
// StartCleanup runs a background goroutine to clean up expired entries
|
|
func StartCleanup() {
|
|
go func() {
|
|
ticker := time.NewTicker(cleanupInterval)
|
|
for range ticker.C {
|
|
loginRateLimiter.Cleanup()
|
|
}
|
|
}()
|
|
}
|
|
|
|
func init() {
|
|
StartCleanup()
|
|
}
|
|
|
|
// RecordFailedAttempt increments the failed attempt counter for an IP
|
|
func RecordFailedAttempt(ip string) {
|
|
loginRateLimiter.Record(ip)
|
|
}
|
|
|
|
// ResetAttempts clears the attempt counter for an IP (called on successful login)
|
|
func ResetAttempts(ip string) {
|
|
loginRateLimiter.Reset(ip)
|
|
}
|
|
|
|
// IsRateLimited checks if an IP has exceeded the failed login threshold
|
|
func IsRateLimited(ip string) bool {
|
|
return loginRateLimiter.IsLimited(ip)
|
|
}
|
|
|
|
// Record adds a failed attempt for the given IP
|
|
func (rl *RateLimiter) Record(ip string) {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
|
|
now := time.Now()
|
|
attempt, exists := rl.attempts[ip]
|
|
|
|
if !exists {
|
|
rl.attempts[ip] = &LoginAttempt{
|
|
Count: 1,
|
|
FirstTime: now,
|
|
}
|
|
return
|
|
}
|
|
|
|
// Reset if window has passed
|
|
if now.Sub(attempt.FirstTime) >= windowDuration {
|
|
attempt.Count = 1
|
|
attempt.FirstTime = now
|
|
return
|
|
}
|
|
|
|
attempt.Count++
|
|
}
|
|
|
|
// IsLimited returns true if the IP has exceeded max attempts
|
|
func (rl *RateLimiter) IsLimited(ip string) bool {
|
|
rl.mu.RLock()
|
|
defer rl.mu.RUnlock()
|
|
|
|
attempt, exists := rl.attempts[ip]
|
|
if !exists {
|
|
return false
|
|
}
|
|
|
|
// Check if window has passed
|
|
if time.Since(attempt.FirstTime) >= windowDuration {
|
|
return false
|
|
}
|
|
|
|
return attempt.Count >= maxAttempts
|
|
}
|
|
|
|
// Reset clears all attempts for the given IP
|
|
func (rl *RateLimiter) Reset(ip string) {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
delete(rl.attempts, ip)
|
|
}
|
|
|
|
// Cleanup removes expired entries from the attempts map
|
|
func (rl *RateLimiter) Cleanup() {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
|
|
now := time.Now()
|
|
for ip, attempt := range rl.attempts {
|
|
if now.Sub(attempt.FirstTime) >= windowDuration {
|
|
delete(rl.attempts, ip)
|
|
}
|
|
}
|
|
}
|
|
|
|
// RateLimitMiddleware returns a middleware that rate limits login attempts
|
|
func RateLimitMiddleware() gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
ip := c.ClientIP()
|
|
|
|
if IsRateLimited(ip) {
|
|
response.Error(c, http.StatusTooManyRequests, "too many login attempts, please try again later")
|
|
c.Abort()
|
|
return
|
|
}
|
|
|
|
c.Next()
|
|
}
|
|
}
|