weijuesen
123 lines
2.8 KiB
Go
123 lines
2.8 KiB
Go
package middleware
|
|
|
|
import (
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
// loginAttempt 登录失败计数（按 IP + 用户名维度）
|
|
type loginAttempt struct {
|
|
failures int
|
|
lockUntil time.Time
|
|
lastFail time.Time
|
|
}
|
|
|
|
type loginLimiter struct {
|
|
mu sync.Mutex
|
|
seen map[string]*loginAttempt
|
|
}
|
|
|
|
const (
|
|
maxFailures = 5 // 连续失败 5 次后锁定
|
|
lockDuration = 15 * time.Minute
|
|
failureWindow = 10 * time.Minute // 失败计数窗口
|
|
cleanupInterval = 5 * time.Minute
|
|
)
|
|
|
|
var defaultLoginLimiter = newLoginLimiter()
|
|
|
|
func newLoginLimiter() *loginLimiter {
|
|
l := &loginLimiter{seen: make(map[string]*loginAttempt)}
|
|
go l.cleanupLoop()
|
|
return l
|
|
}
|
|
|
|
func (l *loginLimiter) cleanupLoop() {
|
|
t := time.NewTicker(cleanupInterval)
|
|
defer t.Stop()
|
|
for range t.C {
|
|
l.mu.Lock()
|
|
now := time.Now()
|
|
for k, v := range l.seen {
|
|
if now.After(v.lockUntil) && now.Sub(v.lastFail) > failureWindow {
|
|
delete(l.seen, k)
|
|
}
|
|
}
|
|
l.mu.Unlock()
|
|
}
|
|
}
|
|
|
|
// key = ip + "|" + username（小写）
|
|
func limiterKey(c *gin.Context, username string) string {
|
|
return c.ClientIP() + "|" + strings.ToLower(strings.TrimSpace(username))
|
|
}
|
|
|
|
// checkLock 返回是否被锁定及剩余锁定时间
|
|
func (l *loginLimiter) checkLock(key string) (bool, time.Duration) {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
a, ok := l.seen[key]
|
|
if !ok {
|
|
return false, 0
|
|
}
|
|
if time.Now().Before(a.lockUntil) {
|
|
return true, time.Until(a.lockUntil)
|
|
}
|
|
return false, 0
|
|
}
|
|
|
|
// recordFailure 记录一次失败，达到阈值则锁定
|
|
func (l *loginLimiter) recordFailure(key string) {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
a, ok := l.seen[key]
|
|
if !ok {
|
|
a = &loginAttempt{}
|
|
l.seen[key] = a
|
|
}
|
|
now := time.Now()
|
|
// 窗口外重置
|
|
if now.Sub(a.lastFail) > failureWindow {
|
|
a.failures = 0
|
|
}
|
|
a.failures++
|
|
a.lastFail = now
|
|
if a.failures >= maxFailures {
|
|
a.lockUntil = now.Add(lockDuration)
|
|
}
|
|
}
|
|
|
|
// recordSuccess 登录成功后清空计数
|
|
func (l *loginLimiter) recordSuccess(key string) {
|
|
l.mu.Lock()
|
|
delete(l.seen, key)
|
|
l.mu.Unlock()
|
|
}
|
|
|
|
// CheckLoginLock 检查是否被锁定，被锁定则写 429 并返回 true（在 handler 解析 body 后调用）
|
|
func CheckLoginLock(c *gin.Context, username string) bool {
|
|
key := limiterKey(c, username)
|
|
if locked, remain := defaultLoginLimiter.checkLock(key); locked {
|
|
c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
|
|
"error": "登录尝试过多，已锁定，请稍后再试",
|
|
"retry": int(remain.Minutes()) + 1,
|
|
})
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
// RecordLoginFail 记录登录失败
|
|
func RecordLoginFail(c *gin.Context, username string) {
|
|
defaultLoginLimiter.recordFailure(limiterKey(c, username))
|
|
}
|
|
|
|
// RecordLoginSuccess 登录成功后清空计数
|
|
func RecordLoginSuccess(c *gin.Context, username string) {
|
|
defaultLoginLimiter.recordSuccess(limiterKey(c, username))
|
|
}
|