105 lines
2.7 KiB
Go
105 lines
2.7 KiB
Go
package service
|
||
|
||
import (
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
// RateLimiter 内存速率限制器
|
||
type RateLimiter struct {
|
||
mu sync.Mutex
|
||
records map[string][]time.Time
|
||
limits map[string]int // key -> 每分钟最大请求数
|
||
cooldown map[string]time.Duration
|
||
}
|
||
|
||
func NewRateLimiter() *RateLimiter {
|
||
return &RateLimiter{
|
||
records: make(map[string][]time.Time),
|
||
limits: make(map[string]int),
|
||
cooldown: make(map[string]time.Duration),
|
||
}
|
||
}
|
||
|
||
// SetLimit 设置某类请求的每分钟限制
|
||
func (r *RateLimiter) SetLimit(key string, perMinute int) {
|
||
r.mu.Lock()
|
||
defer r.mu.Unlock()
|
||
r.limits[key] = perMinute
|
||
}
|
||
|
||
// Allow 检查是否允许请求
|
||
func (r *RateLimiter) Allow(key string) bool {
|
||
r.mu.Lock()
|
||
defer r.mu.Unlock()
|
||
|
||
limit, ok := r.limits[strings.SplitN(key, ":", 2)[0]]
|
||
if !ok {
|
||
return true
|
||
}
|
||
|
||
now := time.Now()
|
||
if len(r.records) > 10000 {
|
||
for k, v := range r.records {
|
||
if len(v) == 0 || now.Sub(v[len(v)-1]) > time.Minute {
|
||
delete(r.records, k)
|
||
}
|
||
}
|
||
if len(r.records) > 20000 {
|
||
return false
|
||
}
|
||
}
|
||
windowStart := now.Add(-1 * time.Minute)
|
||
|
||
// 清理过期记录
|
||
records := r.records[key]
|
||
valid := records[:0]
|
||
for _, t := range records {
|
||
if t.After(windowStart) {
|
||
valid = append(valid, t)
|
||
}
|
||
}
|
||
|
||
if len(valid) >= limit {
|
||
r.records[key] = valid
|
||
return false
|
||
}
|
||
|
||
valid = append(valid, now)
|
||
r.records[key] = valid
|
||
return true
|
||
}
|
||
|
||
// 速率限制类型常量
|
||
const (
|
||
RateLogin = "login"
|
||
RateRegister = "register"
|
||
RatePost = "post"
|
||
RateComment = "comment"
|
||
RateChat = "chat" // 群聊发消息
|
||
RateUpload = "upload" // 帖子插图等上传
|
||
RateInteract = "interact" // 投票/抽奖/解锁等互动
|
||
RateHidePassword = "hide_password" // 密码隐藏块尝试
|
||
RateTimelineGit = "timeline_git" // Git 提交导入 10/分钟
|
||
RateAdSubmit = "ad_submit" // 广告申购 5/分钟（按 IP）
|
||
RateCaptcha = "captcha" // 验证码刷新 30/分钟
|
||
)
|
||
|
||
// DefaultRateLimiter 创建默认速率限制器
|
||
func DefaultRateLimiter() *RateLimiter {
|
||
rl := NewRateLimiter()
|
||
rl.SetLimit(RateLogin, 20) // 登录 20/分钟
|
||
rl.SetLimit(RateRegister, 10) // 注册 10/分钟
|
||
rl.SetLimit(RatePost, 10) // 发帖 10/分钟
|
||
rl.SetLimit(RateComment, 30) // 评论 30/分钟
|
||
rl.SetLimit(RateChat, 30) // 群聊消息 30/分钟
|
||
rl.SetLimit(RateUpload, 20) // 图片上传 20/分钟
|
||
rl.SetLimit(RateInteract, 40) // 互动 40/分钟
|
||
rl.SetLimit(RateHidePassword, 20) // 密码尝试 20/分钟（按 IP）
|
||
rl.SetLimit(RateTimelineGit, 10) // 时间线 Git 导入 10/分钟
|
||
rl.SetLimit(RateAdSubmit, 5) // 广告申购 5/分钟
|
||
rl.SetLimit(RateCaptcha, 30) // 验证码 30/分钟
|
||
return rl
|
||
}
|