93 lines
2.2 KiB
Go
93 lines
2.2 KiB
Go
package security
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// LoginAttempt 登录尝试记录
|
|
type LoginAttempt struct {
|
|
Count int // 尝试次数
|
|
LastTry time.Time // 最后尝试时间
|
|
LockedUntil time.Time // 锁定截止时间
|
|
}
|
|
|
|
// RateLimiter 登录限流器，基于 IP 防止暴力破解
|
|
type RateLimiter struct {
|
|
mu sync.Mutex // 互斥锁
|
|
attempts map[string]*LoginAttempt // 尝试记录映射
|
|
maxAttempts int // 最大尝试次数
|
|
lockoutDuration time.Duration // 锁定时长
|
|
}
|
|
|
|
// NewRateLimiter 创建登录限流器
|
|
// @param maxAttempts 最大尝试次数
|
|
// @param lockoutDuration 锁定时长（分钟）
|
|
// @return *RateLimiter 限流器实例
|
|
// @author sunct
|
|
func NewRateLimiter(maxAttempts int, lockoutDuration int) *RateLimiter {
|
|
return &RateLimiter{
|
|
attempts: make(map[string]*LoginAttempt),
|
|
maxAttempts: maxAttempts,
|
|
lockoutDuration: time.Duration(lockoutDuration) * time.Minute,
|
|
}
|
|
}
|
|
|
|
// Allow 检查 IP 是否允许登录
|
|
// @param ip 客户端IP
|
|
// @return bool 是否允许
|
|
// @return int 剩余尝试次数
|
|
// @author sunct
|
|
func (rl *RateLimiter) Allow(ip string) (bool, int) {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
|
|
now := time.Now()
|
|
|
|
if attempt, exists := rl.attempts[ip]; exists {
|
|
if now.Before(attempt.LockedUntil) {
|
|
return false, rl.maxAttempts - attempt.Count
|
|
}
|
|
|
|
if now.Sub(attempt.LastTry) > 1*time.Hour {
|
|
delete(rl.attempts, ip)
|
|
return true, rl.maxAttempts
|
|
}
|
|
|
|
if attempt.Count >= rl.maxAttempts {
|
|
attempt.LockedUntil = now.Add(rl.lockoutDuration)
|
|
return false, 0
|
|
}
|
|
}
|
|
|
|
return true, rl.maxAttempts
|
|
}
|
|
|
|
// RecordAttempt 记录一次登录尝试
|
|
// @param ip 客户端IP
|
|
// @author sunct
|
|
func (rl *RateLimiter) RecordAttempt(ip string) {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
|
|
now := time.Now()
|
|
|
|
if attempt, exists := rl.attempts[ip]; exists {
|
|
attempt.Count++
|
|
attempt.LastTry = now
|
|
} else {
|
|
rl.attempts[ip] = &LoginAttempt{
|
|
Count: 1,
|
|
LastTry: now,
|
|
}
|
|
}
|
|
}
|
|
|
|
// Reset 重置 IP 的登录尝试记录
|
|
// @param ip 客户端IP
|
|
// @author sunct
|
|
func (rl *RateLimiter) Reset(ip string) {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
delete(rl.attempts, ip)
|
|
} |