mirror of
https://gitee.com/samwaf/SamWaf.git
synced 2026-10-30 17:51:44 +08:00
fix:cc protect rate/window mode
#IBVCH3
This commit is contained in:
samwaf
@@ -109,6 +109,7 @@ func (w *WafAntiCCApi) RemoveCCBanIPApi(c *gin.Context) {
|
||||
ccCacheKey := enums.CACHE_CCVISITBAN_PRE + req.Ip
|
||||
if global.GCACHE_WAFCACHE.IsKeyExist(ccCacheKey) {
|
||||
global.GCACHE_WAFCACHE.Remove(ccCacheKey)
|
||||
global.GWAF_CHAN_CLEAR_CC_IP <- req.Ip
|
||||
response.OkWithMessage(req.Ip+" 移除成功", c)
|
||||
} else {
|
||||
response.FailWithMessage("键值未找到或以过期", c)
|
||||
|
||||
@@ -16,5 +16,6 @@ const (
|
||||
TASK_SSL_EXPIRE_CHECK = "task_ssl_expire_check" //SSL证书到期检测
|
||||
TASK_NOTICE = "task_notice" //通知信息
|
||||
TASK_HEALTH = "task_health" //健康检测
|
||||
TASK_CLEAR_CC_WINDOWS = "task_clear_cc_windows" //清除ccWindows记录
|
||||
|
||||
)
|
||||
|
||||
@@ -88,6 +88,9 @@ var (
|
||||
GWAF_CHAN_SSL_EXPIRE_CHECK = make(chan int, 10) //SSL证书到期检测
|
||||
GWAF_CHAN_SYNC_HOST_TO_SSL_EXPIRE = make(chan int, 10) //同步已存在主机到SSL证书检测任务里
|
||||
GWAF_CHAN_TASK = make(chan string, 10) //手工执行任务
|
||||
GWAF_CHAN_CLEAR_CC_WINDOWS = make(chan int, 10) //清除cc缓存信息
|
||||
GWAF_CHAN_CLEAR_CC_IP = make(chan string, 10) //清除cc缓存信息IP
|
||||
|
||||
/*****CACHE相关*********/
|
||||
GCACHE_WAFCACHE *cache.WafCache //cache
|
||||
GCACHE_WECHAT_ACCESS string = "" //微信访问密钥
|
||||
|
||||
@@ -301,6 +301,7 @@ func (m *wafSystenService) run() {
|
||||
globalobj.GWAF_RUNTIME_OBJ_WAF_TaskRegistry.RegisterTask(enums.TASK_SSL_EXPIRE_CHECK, waftask.SSLExpireCheck)
|
||||
globalobj.GWAF_RUNTIME_OBJ_WAF_TaskRegistry.RegisterTask(enums.TASK_NOTICE, waftask.TaskStatusNotify)
|
||||
globalobj.GWAF_RUNTIME_OBJ_WAF_TaskRegistry.RegisterTask(enums.TASK_HEALTH, waftask.TaskHealth)
|
||||
globalobj.GWAF_RUNTIME_OBJ_WAF_TaskRegistry.RegisterTask(enums.TASK_CLEAR_CC_WINDOWS, waftask.TaskCC)
|
||||
go waftask.TaskShareDbInfo()
|
||||
|
||||
globalobj.GWAF_RUNTIME_OBJ_WAF_TaskScheduler = waftask.NewTaskScheduler(globalobj.GWAF_RUNTIME_OBJ_WAF_TaskRegistry)
|
||||
@@ -586,6 +587,14 @@ func (m *wafSystenService) run() {
|
||||
zlog.Debug("需要执行的方法", taskMethod)
|
||||
globalobj.GWAF_RUNTIME_OBJ_WAF_TaskScheduler.RunManual(taskMethod)
|
||||
break
|
||||
case clearCcWindows := <-global.GWAF_CHAN_CLEAR_CC_WINDOWS:
|
||||
zlog.Debug("定时清空CCwindows", clearCcWindows)
|
||||
globalobj.GWAF_RUNTIME_OBJ_WAF_ENGINE.ClearCcWindows()
|
||||
break
|
||||
case clearCcWindowsIp := <-global.GWAF_CHAN_CLEAR_CC_IP:
|
||||
zlog.Debug("定时清空CCip", clearCcWindowsIp)
|
||||
globalobj.GWAF_RUNTIME_OBJ_WAF_ENGINE.ClearCcWindowsForIP(clearCcWindowsIp)
|
||||
break
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
+15
-9
@@ -1,15 +1,21 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"SamWaf/model/baseorm"
|
||||
)
|
||||
import "SamWaf/model/baseorm"
|
||||
|
||||
/**
|
||||
|
||||
- 平均速率模式 ：将请求平均分配到每一秒，适合防止突发攻击
|
||||
- 滑动窗口模式 ：严格按照"N秒内最多M次"的规则限流，更符合直观理解
|
||||
**/
|
||||
|
||||
type AntiCC struct {
|
||||
baseorm.BaseOrm
|
||||
HostCode string `json:"host_code"` //网站唯一码（主要键）
|
||||
Rate int `json:"rate"` //速率
|
||||
Limit int `json:"limit"` //限制
|
||||
LockIPMinutes int `json:"lock_minutes"` //封禁分钟
|
||||
Url string `json:"url"` //保护的url
|
||||
Remarks string `json:"remarks"` //备注
|
||||
Id string `json:"id" gorm:"column:id;primary_key"`
|
||||
HostCode string `json:"host_code" gorm:"column:host_code"`
|
||||
Rate int `json:"rate" gorm:"column:rate"`
|
||||
Limit int `json:"limit" gorm:"column:limit"`
|
||||
LockIPMinutes int `json:"lock_ip_minutes" gorm:"column:lock_ip_minutes"`
|
||||
LimitMode string `json:"limit_mode" gorm:"column:limit_mode"` // "rate" 或 "window"
|
||||
Url string `json:"url"` //保护的url
|
||||
Remarks string `json:"remarks"` //备注
|
||||
}
|
||||
|
||||
@@ -3,12 +3,13 @@ package request
|
||||
import "SamWaf/model/common/request"
|
||||
|
||||
type WafAntiCCAddReq struct {
|
||||
HostCode string `json:"host_code" form:"host_code"` //网站唯一码（主要键）
|
||||
Url string `json:"url" form:"url"` //白名单url
|
||||
Rate int `json:"rate" form:"rate"` //速率
|
||||
Limit int `json:"limit" form:"limit"` //限制
|
||||
LockIPMinutes int `json:"lock_minutes"` //封禁分钟
|
||||
Remarks string `json:"remarks" form:"remarks"` //备注
|
||||
HostCode string `json:"host_code" form:"host_code"` //网站唯一码（主要键）
|
||||
Url string `json:"url" form:"url"` //白名单url
|
||||
Rate int `json:"rate" form:"rate"` //速率
|
||||
Limit int `json:"limit" form:"limit"` //限制
|
||||
LockIPMinutes int `json:"lock_ip_minutes" form:"lock_ip_minutes"` //封禁分钟
|
||||
LimitMode string `json:"limit_mode" form:"limit_mode"` // "rate" 或 "window"
|
||||
Remarks string `json:"remarks" form:"remarks"` //备注
|
||||
}
|
||||
type WafAntiCCSearchReq struct {
|
||||
HostCode string `json:"host_code" ` //主机码
|
||||
@@ -21,13 +22,14 @@ type WafAntiCCDetailReq struct {
|
||||
Id string `json:"id" form:"id"` //白名单Url唯一键
|
||||
}
|
||||
type WafAntiCCEditReq struct {
|
||||
Id string `json:"id"` //白名单url唯一键
|
||||
HostCode string `json:"host_code" form:"host_code"` //网站唯一码（主要键）
|
||||
Url string `json:"url" form:"url"` //白名单url
|
||||
Rate int `json:"rate" form:"rate"` //速率
|
||||
Limit int `json:"limit" form:"limit"` //限制
|
||||
LockIPMinutes int `json:"lock_minutes"` //封禁分钟
|
||||
Remarks string `json:"remarks" form:"remarks"` //备注
|
||||
Id string `json:"id"` //白名单url唯一键
|
||||
HostCode string `json:"host_code" form:"host_code"` //网站唯一码（主要键）
|
||||
Url string `json:"url" form:"url"` //白名单url
|
||||
Rate int `json:"rate" form:"rate"` //速率
|
||||
Limit int `json:"limit" form:"limit"` //限制
|
||||
LockIPMinutes int `json:"lock_ip_minutes" form:"lock_ip_minutes"` //封禁分钟
|
||||
LimitMode string `json:"limit_mode" form:"limit_mode"` // "rate" 或 "window"
|
||||
Remarks string `json:"remarks" form:"remarks"` //备注
|
||||
}
|
||||
|
||||
// WafAntiCCRemoveBanIpReq 移除封禁ip
|
||||
|
||||
@@ -30,6 +30,7 @@ func (receiver *WafAntiCCService) AddApi(req request.WafAntiCCAddReq) error {
|
||||
LockIPMinutes: req.LockIPMinutes,
|
||||
Url: req.Url,
|
||||
Remarks: req.Remarks,
|
||||
LimitMode: req.LimitMode,
|
||||
}
|
||||
global.GWAF_LOCAL_DB.Create(bean)
|
||||
return nil
|
||||
@@ -54,6 +55,7 @@ func (receiver *WafAntiCCService) ModifyApi(req request.WafAntiCCEditReq) error
|
||||
"LockIPMinutes": req.LockIPMinutes,
|
||||
"Remarks": req.Remarks,
|
||||
"UPDATE_TIME": customtype.JsonTime(time.Now()),
|
||||
"LimitMode": req.LimitMode,
|
||||
}
|
||||
err := global.GWAF_LOCAL_DB.Model(model.AntiCC{}).Where("id = ?", req.Id).Updates(ipWhiteMap).Error
|
||||
|
||||
|
||||
@@ -137,6 +137,13 @@ func pathCoreSql(db *gorm.DB) {
|
||||
} else {
|
||||
zlog.Info("db", "hosts :insecure_skip_verify init successfully")
|
||||
}
|
||||
//20250328 cc模式 默认还是现有的
|
||||
err = db.Exec("UPDATE anti_ccs SET limit_mode='window' WHERE limit_mode IS NULL ").Error
|
||||
if err != nil {
|
||||
panic("failed to anti_ccs :limit_mode " + err.Error())
|
||||
} else {
|
||||
zlog.Info("db", "anti_ccs :limit_mode init successfully")
|
||||
}
|
||||
// 记录结束时间并计算耗时
|
||||
duration := time.Since(startTime)
|
||||
zlog.Info("create core index completely", "duration", duration.String())
|
||||
|
||||
@@ -24,13 +24,13 @@ func (waf *WafEngine) CheckCC(r *http.Request, weblogbean *innerbean.WebLog, for
|
||||
}
|
||||
// cc 防护 (局部检测 )
|
||||
if hostTarget.PluginIpRateLimiter != nil {
|
||||
limiter := hostTarget.PluginIpRateLimiter.GetLimiter(weblogbean.SRC_IP)
|
||||
if !limiter.Allow() {
|
||||
|
||||
if !hostTarget.PluginIpRateLimiter.Allow(weblogbean.NetSrcIp) {
|
||||
weblogbean.RISK_LEVEL = 1
|
||||
result.IsBlock = true
|
||||
result.Title = "触发IP频次访问限制1"
|
||||
result.Content = "您的访问被阻止超量了1"
|
||||
cacheKey := enums.CACHE_CCVISITBAN_PRE + weblogbean.SRC_IP
|
||||
result.Title = "【局部】触发IP频次访问限制"
|
||||
result.Content = "您的访问被阻止超量了"
|
||||
cacheKey := enums.CACHE_CCVISITBAN_PRE + weblogbean.NetSrcIp
|
||||
//将该IP添加到封禁里
|
||||
global.GCACHE_WAFCACHE.SetWithTTl(cacheKey, 1, time.Duration(hostTarget.AntiCCBean.LockIPMinutes)*time.Minute)
|
||||
return result
|
||||
@@ -38,13 +38,12 @@ func (waf *WafEngine) CheckCC(r *http.Request, weblogbean *innerbean.WebLog, for
|
||||
}
|
||||
// cc 防护 （全局检测 ）
|
||||
if waf.HostTarget[global.GWAF_GLOBAL_HOST_NAME].Host.GUARD_STATUS == 1 && waf.HostTarget[global.GWAF_GLOBAL_HOST_NAME].PluginIpRateLimiter != nil {
|
||||
limiter := waf.HostTarget[global.GWAF_GLOBAL_HOST_NAME].PluginIpRateLimiter.GetLimiter(weblogbean.SRC_IP)
|
||||
if !limiter.Allow() {
|
||||
if !waf.HostTarget[global.GWAF_GLOBAL_HOST_NAME].PluginIpRateLimiter.Allow(weblogbean.NetSrcIp) {
|
||||
weblogbean.RISK_LEVEL = 1
|
||||
result.IsBlock = true
|
||||
result.Title = "【全局】触发IP频次访问限制"
|
||||
result.Content = "您的访问被阻止超量了"
|
||||
cacheKey := enums.CACHE_CCVISITBAN_PRE + weblogbean.SRC_IP
|
||||
cacheKey := enums.CACHE_CCVISITBAN_PRE + weblogbean.NetSrcIp
|
||||
//将该IP添加到封禁里
|
||||
global.GCACHE_WAFCACHE.SetWithTTl(cacheKey, 1, time.Duration(waf.HostTarget[global.GWAF_GLOBAL_HOST_NAME].AntiCCBean.LockIPMinutes)*time.Minute)
|
||||
return result
|
||||
|
||||
@@ -228,9 +228,9 @@ func (waf *WafEngine) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
|
||||
// 检测是否已经被CC封禁
|
||||
ccCacheKey := enums.CACHE_CCVISITBAN_PRE + clientIP
|
||||
ccCacheKey := enums.CACHE_CCVISITBAN_PRE + weblogbean.NetSrcIp
|
||||
if global.GCACHE_WAFCACHE.IsKeyExist(ccCacheKey) {
|
||||
visitIPError := fmt.Sprintf("当前IP已经被CC封禁，IP:%s 归属地区：%s", clientIP, region)
|
||||
visitIPError := fmt.Sprintf("当前IP已经被CC封禁，IP:%s 归属地区：%s", weblogbean.NetSrcIp, region)
|
||||
global.GQEQUE_MESSAGE_DB.Enqueue(innerbean.OperatorMessageInfo{
|
||||
BaseMessageInfo: innerbean.BaseMessageInfo{OperaType: "CC封禁提醒"},
|
||||
OperaCnt: visitIPError,
|
||||
@@ -1025,3 +1025,44 @@ func (waf *WafEngine) StopProxyServer(v innerbean.ServerRunTime) {
|
||||
v.Svr.Close()
|
||||
}
|
||||
}
|
||||
func (waf *WafEngine) ClearCcWindows() {
|
||||
// 清理所有主机的IP限流器记录
|
||||
for _, hostSafe := range waf.HostTarget {
|
||||
if hostSafe.PluginIpRateLimiter != nil {
|
||||
hostSafe.PluginIpRateLimiter.CleanupOldRecords()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ClearCcWindowsForIP 清理特定IP的CC限流记录
|
||||
func (waf *WafEngine) ClearCcWindowsForIP(ip string) {
|
||||
// 遍历所有主机，清理指定IP的限流记录
|
||||
hostCount := 0
|
||||
for hostKey, hostSafe := range waf.HostTarget {
|
||||
if hostSafe.PluginIpRateLimiter != nil {
|
||||
// 获取清理前的请求计数
|
||||
var countBefore int
|
||||
if hostSafe.PluginIpRateLimiter.GetRequestCount != nil {
|
||||
countBefore = hostSafe.PluginIpRateLimiter.GetRequestCount(ip)
|
||||
}
|
||||
|
||||
// 清理该IP的限流记录
|
||||
hostSafe.PluginIpRateLimiter.ClearWindowForIP(ip)
|
||||
hostCount++
|
||||
|
||||
// 获取清理后的请求计数
|
||||
var countAfter int
|
||||
if hostSafe.PluginIpRateLimiter.GetRequestCount != nil {
|
||||
countAfter = hostSafe.PluginIpRateLimiter.GetRequestCount(ip)
|
||||
}
|
||||
|
||||
// 打印清理前后的计数信息
|
||||
zlog.Debug(fmt.Sprintf("主机 %s 的 IP %s 限流记录: 清理前=%d, 清理后=%d",
|
||||
hostKey, ip, countBefore, countAfter))
|
||||
}
|
||||
}
|
||||
|
||||
zlog.Debug(fmt.Sprintf("已清理 IP %s 在 %d 个主机上的 CC 限流记录", ip, hostCount),
|
||||
zap.String("ip", ip),
|
||||
zap.Int("hostCount", hostCount))
|
||||
}
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"SamWaf/webplugin"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
goahocorasick "github.com/samwafgo/ahocorasick"
|
||||
"golang.org/x/time/rate"
|
||||
"net/http"
|
||||
@@ -136,7 +137,19 @@ func (waf *WafEngine) LoadHost(inHost model.Hosts) []innerbean.ServerRunTime {
|
||||
//初始化插件-ip计数器
|
||||
var pluginIpRateLimiter *webplugin.IPRateLimiter
|
||||
if anticcBean.Id != "" {
|
||||
pluginIpRateLimiter = webplugin.NewIPRateLimiter(rate.Limit(anticcBean.Rate), anticcBean.Limit)
|
||||
// 根据配置选择限流模式
|
||||
if anticcBean.LimitMode == "window" {
|
||||
// 使用滑动窗口模式
|
||||
pluginIpRateLimiter = webplugin.NewWindowIPRateLimiter(anticcBean.Rate, anticcBean.Limit)
|
||||
zlog.Debug(fmt.Sprintf("初始化CC防护(滑动窗口模式) 主机%v 时间窗口(秒)%v 最大请求数%v",
|
||||
inHost.Host, anticcBean.Rate, anticcBean.Limit))
|
||||
} else {
|
||||
// 使用平均速率模式(默认)
|
||||
ratePerSecond := rate.Limit(float64(anticcBean.Limit) / float64(anticcBean.Rate))
|
||||
pluginIpRateLimiter = webplugin.NewIPRateLimiter(ratePerSecond, anticcBean.Limit)
|
||||
zlog.Debug(fmt.Sprintf("初始化CC防护(平均速率模式) 主机%v 时间窗口(秒)%v 最大请求数%v 每秒速率%v",
|
||||
inHost.Host, anticcBean.Rate, anticcBean.Limit, float64(anticcBean.Limit)/float64(anticcBean.Rate)))
|
||||
}
|
||||
}
|
||||
|
||||
//查询ip白名单
|
||||
|
||||
+9
-2
@@ -111,12 +111,19 @@ func InitTaskDb() []model.Task {
|
||||
TaskMethod: enums.TASK_NOTICE,
|
||||
})
|
||||
syncTaskToDb(model.Task{
|
||||
TaskName: "每天5s进行健康度检测",
|
||||
TaskName: "每30s进行健康度检测",
|
||||
TaskUnit: enums.TASK_SECOND,
|
||||
TaskValue: 5,
|
||||
TaskValue: 30,
|
||||
TaskAt: "",
|
||||
TaskMethod: enums.TASK_HEALTH,
|
||||
})
|
||||
syncTaskToDb(model.Task{
|
||||
TaskName: "每5分钟进行CCWindows旧信息清除",
|
||||
TaskUnit: enums.TASK_MIN,
|
||||
TaskValue: 5,
|
||||
TaskAt: "",
|
||||
TaskMethod: enums.TASK_CLEAR_CC_WINDOWS,
|
||||
})
|
||||
list, _, _ := wafTaskService.GetList()
|
||||
return list
|
||||
}
|
||||
|
||||
@@ -0,0 +1,9 @@
|
||||
package waftask
|
||||
|
||||
import (
|
||||
"SamWaf/global"
|
||||
)
|
||||
|
||||
func TaskCC() {
|
||||
global.GWAF_CHAN_CLEAR_CC_WINDOWS <- 1
|
||||
}
|
||||
+155
-9
@@ -2,25 +2,60 @@ package webplugin
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
// LimitMode 限流模式
|
||||
type LimitMode int
|
||||
|
||||
const (
|
||||
// RateMode 平均速率模式 - 每秒固定速率
|
||||
RateMode LimitMode = iota
|
||||
// WindowMode 滑动窗口模式 - N秒内最多M次
|
||||
WindowMode
|
||||
)
|
||||
|
||||
// IPRateLimiter .
|
||||
type IPRateLimiter struct {
|
||||
ips map[string]*rate.Limiter
|
||||
mu *sync.RWMutex
|
||||
r rate.Limit
|
||||
b int
|
||||
ips map[string]*rate.Limiter
|
||||
mu *sync.RWMutex
|
||||
r rate.Limit
|
||||
b int
|
||||
mode LimitMode
|
||||
window int // 时间窗口大小(秒)
|
||||
requests map[string][]time.Time // 用于滑动窗口模式记录请求时间
|
||||
}
|
||||
|
||||
// NewIPRateLimiter .
|
||||
// NewIPRateLimiter 创建一个新的IP限流器
|
||||
// r: 每秒请求速率
|
||||
// b: 突发请求数量
|
||||
func NewIPRateLimiter(r rate.Limit, b int) *IPRateLimiter {
|
||||
i := &IPRateLimiter{
|
||||
ips: make(map[string]*rate.Limiter),
|
||||
mu: &sync.RWMutex{},
|
||||
r: r,
|
||||
b: b,
|
||||
ips: make(map[string]*rate.Limiter),
|
||||
mu: &sync.RWMutex{},
|
||||
r: r,
|
||||
b: b,
|
||||
mode: RateMode, // 默认使用平均速率模式，保持向后兼容
|
||||
requests: make(map[string][]time.Time),
|
||||
}
|
||||
|
||||
return i
|
||||
}
|
||||
|
||||
// NewWindowIPRateLimiter 创建一个基于滑动窗口的IP限流器
|
||||
// window: 时间窗口大小(秒)
|
||||
// maxRequests: 窗口内最大请求数
|
||||
func NewWindowIPRateLimiter(window, maxRequests int) *IPRateLimiter {
|
||||
i := &IPRateLimiter{
|
||||
ips: make(map[string]*rate.Limiter),
|
||||
mu: &sync.RWMutex{},
|
||||
r: rate.Limit(float64(maxRequests) / float64(window)), // 保持兼容
|
||||
b: maxRequests,
|
||||
mode: WindowMode,
|
||||
window: window,
|
||||
requests: make(map[string][]time.Time),
|
||||
}
|
||||
|
||||
return i
|
||||
@@ -35,6 +70,7 @@ func (i *IPRateLimiter) AddIP(ip string) *rate.Limiter {
|
||||
limiter := rate.NewLimiter(i.r, i.b)
|
||||
|
||||
i.ips[ip] = limiter
|
||||
i.requests[ip] = []time.Time{}
|
||||
|
||||
return limiter
|
||||
}
|
||||
@@ -54,3 +90,113 @@ func (i *IPRateLimiter) GetLimiter(ip string) *rate.Limiter {
|
||||
|
||||
return limiter
|
||||
}
|
||||
|
||||
// Allow 检查是否允许请求通过
|
||||
// 根据模式使用不同的限流策略
|
||||
func (i *IPRateLimiter) Allow(ip string) bool {
|
||||
if i.mode == RateMode {
|
||||
// 使用令牌桶算法限流
|
||||
return i.GetLimiter(ip).Allow()
|
||||
} else {
|
||||
// 使用滑动窗口算法限流
|
||||
return i.allowByWindow(ip)
|
||||
}
|
||||
}
|
||||
|
||||
// allowByWindow 使用滑动窗口算法检查是否允许请求通过
|
||||
func (i *IPRateLimiter) allowByWindow(ip string) bool {
|
||||
i.mu.Lock()
|
||||
defer i.mu.Unlock()
|
||||
|
||||
now := time.Now()
|
||||
windowStart := now.Add(-time.Duration(i.window) * time.Second)
|
||||
|
||||
// 如果IP不存在，初始化
|
||||
if _, exists := i.requests[ip]; !exists {
|
||||
i.requests[ip] = []time.Time{}
|
||||
}
|
||||
|
||||
// 清理过期的请求记录
|
||||
var validRequests []time.Time
|
||||
for _, t := range i.requests[ip] {
|
||||
if t.After(windowStart) { // 修改：添加Equal条件，确保边界值也被包含|| t.Equal(windowStart)
|
||||
validRequests = append(validRequests, t)
|
||||
}
|
||||
}
|
||||
|
||||
// 检查是否超过限制
|
||||
if len(validRequests) >= i.b {
|
||||
i.requests[ip] = validRequests
|
||||
return false
|
||||
}
|
||||
|
||||
// 记录新请求
|
||||
i.requests[ip] = append(validRequests, now)
|
||||
return true
|
||||
}
|
||||
|
||||
// CleanupOldRecords 清理过期的请求记录
|
||||
func (i *IPRateLimiter) CleanupOldRecords() {
|
||||
i.mu.Lock()
|
||||
defer i.mu.Unlock()
|
||||
|
||||
now := time.Now()
|
||||
windowStart := now.Add(-time.Duration(i.window) * time.Second)
|
||||
|
||||
for ip, times := range i.requests {
|
||||
var validRequests []time.Time
|
||||
for _, t := range times {
|
||||
if t.After(windowStart) {
|
||||
validRequests = append(validRequests, t)
|
||||
}
|
||||
}
|
||||
|
||||
if len(validRequests) == 0 {
|
||||
delete(i.requests, ip)
|
||||
delete(i.ips, ip)
|
||||
} else {
|
||||
i.requests[ip] = validRequests
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ClearWindowForIP 清空指定IP的滑动窗口记录
|
||||
// 用于手动重置某个IP的限流状态
|
||||
func (i *IPRateLimiter) ClearWindowForIP(ip string) {
|
||||
i.mu.Lock()
|
||||
defer i.mu.Unlock()
|
||||
|
||||
// 清空滑动窗口记录
|
||||
if _, exists := i.requests[ip]; exists {
|
||||
i.requests[ip] = []time.Time{}
|
||||
}
|
||||
|
||||
// 重置令牌桶（用于平均速率模式）
|
||||
if _, exists := i.ips[ip]; exists {
|
||||
// 创建一个新的限流器替换旧的
|
||||
newLimiter := rate.NewLimiter(i.r, i.b)
|
||||
i.ips[ip] = newLimiter
|
||||
}
|
||||
}
|
||||
|
||||
// GetRequestCount 获取指定IP在当前窗口内的请求数量
|
||||
func (i *IPRateLimiter) GetRequestCount(ip string) int {
|
||||
i.mu.RLock()
|
||||
defer i.mu.RUnlock()
|
||||
|
||||
now := time.Now()
|
||||
windowStart := now.Add(-time.Duration(i.window) * time.Second)
|
||||
|
||||
if _, exists := i.requests[ip]; !exists {
|
||||
return 0
|
||||
}
|
||||
|
||||
count := 0
|
||||
for _, t := range i.requests[ip] {
|
||||
if t.After(windowStart) || t.Equal(windowStart) {
|
||||
count++
|
||||
}
|
||||
}
|
||||
|
||||
return count
|
||||
}
|
||||
|
||||
@@ -0,0 +1,695 @@
|
||||
package webplugin
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"golang.org/x/time/rate"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestIPRateLimiter(t *testing.T) {
|
||||
// 创建一个限流器，设置为每秒2个请求，最大突发10个请求
|
||||
// 相当于5秒内最多允许10个请求
|
||||
limiter := NewIPRateLimiter(rate.Limit(2), 10)
|
||||
|
||||
// 测试单个IP的限流
|
||||
t.Run("单IP限流测试", func(t *testing.T) {
|
||||
ip := "192.168.1.1"
|
||||
ipLimiter := limiter.GetLimiter(ip)
|
||||
|
||||
// 模拟11个连续请求
|
||||
allowedCount := 0
|
||||
for i := 0; i < 11; i++ {
|
||||
if ipLimiter.Allow() {
|
||||
allowedCount++
|
||||
}
|
||||
}
|
||||
|
||||
// 由于突发限制为10，所以应该只有10个请求被允许
|
||||
if allowedCount != 10 {
|
||||
t.Errorf("预期允许10个请求，实际允许了%d个请求", allowedCount)
|
||||
}
|
||||
|
||||
// 等待一段时间后，应该可以继续发送请求
|
||||
time.Sleep(1 * time.Second)
|
||||
if !ipLimiter.Allow() {
|
||||
t.Error("等待1秒后应该允许新的请求")
|
||||
}
|
||||
})
|
||||
|
||||
// 测试多个IP是否相互独立
|
||||
t.Run("多IP独立限流测试", func(t *testing.T) {
|
||||
ip1 := "192.168.1.2"
|
||||
ip2 := "192.168.1.3"
|
||||
|
||||
ipLimiter1 := limiter.GetLimiter(ip1)
|
||||
ipLimiter2 := limiter.GetLimiter(ip2)
|
||||
|
||||
// IP1发送10个请求
|
||||
for i := 0; i < 10; i++ {
|
||||
ipLimiter1.Allow()
|
||||
}
|
||||
|
||||
// IP1应该被限流
|
||||
if ipLimiter1.Allow() {
|
||||
t.Error("IP1应该被限流")
|
||||
}
|
||||
|
||||
// IP2不应该受IP1的影响
|
||||
if !ipLimiter2.Allow() {
|
||||
t.Error("IP2不应该受IP1的限流影响")
|
||||
}
|
||||
})
|
||||
|
||||
// 测试并发场景
|
||||
t.Run("并发请求测试", func(t *testing.T) {
|
||||
ip := "192.168.1.4"
|
||||
var wg sync.WaitGroup
|
||||
successCount := 0
|
||||
var mu sync.Mutex
|
||||
|
||||
// 创建20个并发请求
|
||||
for i := 0; i < 20; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if limiter.GetLimiter(ip).Allow() {
|
||||
mu.Lock()
|
||||
successCount++
|
||||
mu.Unlock()
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
// 由于突发限制为10，所以应该只有10个请求被允许
|
||||
if successCount > 10 {
|
||||
t.Errorf("并发场景下预期最多允许10个请求，实际允许了%d个请求", successCount)
|
||||
}
|
||||
})
|
||||
|
||||
// 测试实际场景：5秒内最多10个请求
|
||||
t.Run("实际场景测试：5秒内最多10个请求", func(t *testing.T) {
|
||||
// 创建一个新的限流器，设置为每秒2个请求（5秒10个）
|
||||
scenarioLimiter := NewIPRateLimiter(rate.Limit(2), 10)
|
||||
ip := "192.168.1.5"
|
||||
ipLimiter := scenarioLimiter.GetLimiter(ip)
|
||||
|
||||
fmt.Println("开始测试5秒内最多10个请求的场景...")
|
||||
|
||||
// 第一秒：尝试发送4个请求，应该允许前3个（初始令牌桶有2个，每秒产生2个）
|
||||
allowedInFirstSecond := 0
|
||||
for i := 0; i < 4; i++ {
|
||||
if ipLimiter.Allow() {
|
||||
allowedInFirstSecond++
|
||||
}
|
||||
}
|
||||
fmt.Printf("第1秒：尝试4个请求，允许了%d个\n", allowedInFirstSecond)
|
||||
|
||||
// 等待1秒
|
||||
time.Sleep(1 * time.Second)
|
||||
|
||||
// 第二秒：尝试发送3个请求，应该允许2个
|
||||
allowedInSecondSecond := 0
|
||||
for i := 0; i < 3; i++ {
|
||||
if ipLimiter.Allow() {
|
||||
allowedInSecondSecond++
|
||||
}
|
||||
}
|
||||
fmt.Printf("第2秒：尝试3个请求，允许了%d个\n", allowedInSecondSecond)
|
||||
|
||||
// 等待1秒
|
||||
time.Sleep(1 * time.Second)
|
||||
|
||||
// 第三秒：尝试发送3个请求，应该允许2个
|
||||
allowedInThirdSecond := 0
|
||||
for i := 0; i < 3; i++ {
|
||||
if ipLimiter.Allow() {
|
||||
allowedInThirdSecond++
|
||||
}
|
||||
}
|
||||
fmt.Printf("第3秒：尝试3个请求，允许了%d个\n", allowedInThirdSecond)
|
||||
|
||||
// 验证总体结果
|
||||
totalAllowed := allowedInFirstSecond + allowedInSecondSecond + allowedInThirdSecond
|
||||
fmt.Printf("总计：尝试10个请求，允许了%d个\n", totalAllowed)
|
||||
|
||||
if totalAllowed > 10 {
|
||||
t.Errorf("5秒内应该最多允许10个请求，实际允许了%d个", totalAllowed)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// 测试实际的CC防护场景
|
||||
func TestCCProtection(t *testing.T) {
|
||||
// 模拟配置：5秒内最多允许10个请求
|
||||
timeWindow := 5
|
||||
maxRequests := 10
|
||||
|
||||
// 计算每秒速率
|
||||
ratePerSecond := rate.Limit(float64(maxRequests) / float64(timeWindow))
|
||||
|
||||
// 创建限流器
|
||||
limiter := NewIPRateLimiter(ratePerSecond, maxRequests)
|
||||
|
||||
// 测试IP
|
||||
ip := "192.168.1.100"
|
||||
|
||||
fmt.Printf("模拟CC防护：%d秒内最多允许%d个请求（每秒%.2f个）\n",
|
||||
timeWindow, maxRequests, float64(maxRequests)/float64(timeWindow))
|
||||
|
||||
// 模拟正常用户：每秒2个请求
|
||||
t.Run("正常用户测试", func(t *testing.T) {
|
||||
//ipLimiter := limiter.GetLimiter(ip)
|
||||
|
||||
// 重置限流器
|
||||
limiter.ips = make(map[string]*rate.Limiter)
|
||||
|
||||
success := 0
|
||||
fail := 0
|
||||
|
||||
// 模拟5秒，每秒发送2个请求
|
||||
for second := 1; second <= timeWindow; second++ {
|
||||
for i := 0; i < 2; i++ {
|
||||
if limiter.GetLimiter(ip).Allow() {
|
||||
success++
|
||||
} else {
|
||||
fail++
|
||||
}
|
||||
}
|
||||
time.Sleep(1 * time.Second)
|
||||
}
|
||||
|
||||
fmt.Printf("正常用户（每秒2个请求）：成功=%d, 失败=%d\n", success, fail)
|
||||
|
||||
// 正常用户应该不会被限流
|
||||
if fail > 0 {
|
||||
t.Errorf("正常用户不应该被限流，但有%d个请求被拒绝", fail)
|
||||
}
|
||||
})
|
||||
|
||||
// 模拟CC攻击：短时间内发送大量请求
|
||||
t.Run("CC攻击测试", func(t *testing.T) {
|
||||
// 重置限流器
|
||||
limiter.ips = make(map[string]*rate.Limiter)
|
||||
|
||||
success := 0
|
||||
fail := 0
|
||||
|
||||
// 短时间内发送20个请求
|
||||
for i := 0; i < 20; i++ {
|
||||
if limiter.GetLimiter(ip).Allow() {
|
||||
success++
|
||||
} else {
|
||||
fail++
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Printf("CC攻击（短时间内20个请求）：成功=%d, 失败=%d\n", success, fail)
|
||||
|
||||
// 应该有请求被限流
|
||||
if fail == 0 {
|
||||
t.Error("CC攻击应该被限流，但所有请求都成功了")
|
||||
}
|
||||
|
||||
// 成功的请求不应超过最大允许数
|
||||
if success > maxRequests {
|
||||
t.Errorf("成功请求数不应超过%d，但实际为%d", maxRequests, success)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// 测试滑动窗口模式的限流器
|
||||
func TestWindowIPRateLimiter(t *testing.T) {
|
||||
// 创建一个滑动窗口限流器，设置为5秒内最多10个请求
|
||||
timeWindow := 5
|
||||
maxRequests := 10
|
||||
limiter := NewWindowIPRateLimiter(timeWindow, maxRequests)
|
||||
|
||||
// 测试IP
|
||||
ip := "192.168.1.200"
|
||||
|
||||
// 测试基本限流功能
|
||||
t.Run("滑动窗口基本限流测试", func(t *testing.T) {
|
||||
// 重置限流器
|
||||
limiter.requests = make(map[string][]time.Time)
|
||||
limiter.ips = make(map[string]*rate.Limiter)
|
||||
|
||||
// 连续发送11个请求
|
||||
allowedCount := 0
|
||||
for i := 0; i < 11; i++ {
|
||||
if limiter.Allow(ip) {
|
||||
allowedCount++
|
||||
}
|
||||
}
|
||||
|
||||
// 应该只有10个请求被允许
|
||||
if allowedCount != 10 {
|
||||
t.Errorf("滑动窗口模式下预期允许10个请求，实际允许了%d个请求", allowedCount)
|
||||
}
|
||||
|
||||
// 等待1秒后，应该仍然不允许新请求（因为窗口是5秒）
|
||||
time.Sleep(1 * time.Second)
|
||||
if limiter.Allow(ip) {
|
||||
t.Error("等待1秒后不应该允许新的请求，因为窗口期是5秒")
|
||||
}
|
||||
|
||||
// 等待5秒后，应该允许新请求
|
||||
time.Sleep(5 * time.Second)
|
||||
if !limiter.Allow(ip) {
|
||||
t.Error("等待5秒后应该允许新的请求")
|
||||
}
|
||||
})
|
||||
|
||||
// 测试滑动窗口特性
|
||||
t.Run("滑动窗口特性测试", func(t *testing.T) {
|
||||
// 重置限流器
|
||||
limiter.requests = make(map[string][]time.Time)
|
||||
limiter.ips = make(map[string]*rate.Limiter)
|
||||
|
||||
fmt.Println("开始测试滑动窗口特性...")
|
||||
|
||||
// 第1秒：发送5个请求
|
||||
allowedInFirstSecond := 0
|
||||
for i := 0; i < 5; i++ {
|
||||
if limiter.Allow(ip) {
|
||||
allowedInFirstSecond++
|
||||
}
|
||||
}
|
||||
fmt.Printf("第1秒：尝试5个请求，允许了%d个\n", allowedInFirstSecond)
|
||||
|
||||
// 等待3秒
|
||||
time.Sleep(3 * time.Second)
|
||||
|
||||
// 第4秒：再发送5个请求
|
||||
allowedInFourthSecond := 0
|
||||
for i := 0; i < 5; i++ {
|
||||
if limiter.Allow(ip) {
|
||||
allowedInFourthSecond++
|
||||
}
|
||||
}
|
||||
fmt.Printf("第4秒：尝试5个请求，允许了%d个\n", allowedInFourthSecond)
|
||||
|
||||
// 第4秒：尝试再发送1个请求，应该被拒绝
|
||||
if limiter.Allow(ip) {
|
||||
t.Error("第4秒：已经发送了10个请求，应该被限流")
|
||||
}
|
||||
|
||||
// 等待2秒（此时第1秒的请求应该过期）
|
||||
time.Sleep(2 * time.Second)
|
||||
|
||||
// 第6秒：应该可以发送5个新请求（因为第1秒的5个请求已经滑出窗口）
|
||||
allowedInSixthSecond := 0
|
||||
for i := 0; i < 5; i++ {
|
||||
if limiter.Allow(ip) {
|
||||
allowedInSixthSecond++
|
||||
}
|
||||
}
|
||||
fmt.Printf("第6秒：尝试5个请求，允许了%d个\n", allowedInSixthSecond)
|
||||
|
||||
// 验证结果
|
||||
if allowedInSixthSecond != 5 {
|
||||
t.Errorf("第6秒应该允许5个请求（因为窗口滑动），但实际允许了%d个", allowedInSixthSecond)
|
||||
}
|
||||
})
|
||||
|
||||
// 测试清理过期记录功能
|
||||
t.Run("清理过期记录测试", func(t *testing.T) {
|
||||
// 重置限流器
|
||||
limiter.requests = make(map[string][]time.Time)
|
||||
limiter.ips = make(map[string]*rate.Limiter)
|
||||
|
||||
// 添加一些IP和请求记录
|
||||
ips := []string{"192.168.1.201", "192.168.1.202", "192.168.1.203"}
|
||||
for _, testIP := range ips {
|
||||
for i := 0; i < 5; i++ {
|
||||
limiter.Allow(testIP)
|
||||
}
|
||||
}
|
||||
|
||||
// 验证记录数量
|
||||
if len(limiter.requests) != 3 {
|
||||
t.Errorf("应该有3个IP的请求记录，但实际有%d个", len(limiter.requests))
|
||||
}
|
||||
|
||||
// 等待窗口期过后
|
||||
time.Sleep(time.Duration(timeWindow+1) * time.Second)
|
||||
|
||||
// 清理过期记录
|
||||
limiter.CleanupOldRecords()
|
||||
|
||||
// 验证记录是否被清理
|
||||
if len(limiter.requests) != 0 {
|
||||
t.Errorf("过期记录应该被清理，但仍有%d个记录", len(limiter.requests))
|
||||
}
|
||||
})
|
||||
|
||||
// 比较两种模式的行为差异
|
||||
t.Run("两种模式行为对比", func(t *testing.T) {
|
||||
// 创建平均速率模式限流器
|
||||
ratePerSecond := rate.Limit(float64(maxRequests) / float64(timeWindow))
|
||||
rateLimiter := NewIPRateLimiter(ratePerSecond, maxRequests)
|
||||
|
||||
// 创建滑动窗口模式限流器
|
||||
windowLimiter := NewWindowIPRateLimiter(timeWindow, maxRequests)
|
||||
|
||||
// 重置限流器
|
||||
rateLimiter.ips = make(map[string]*rate.Limiter)
|
||||
windowLimiter.requests = make(map[string][]time.Time)
|
||||
windowLimiter.ips = make(map[string]*rate.Limiter)
|
||||
|
||||
fmt.Println("\n开始对比两种模式的行为差异...")
|
||||
|
||||
// 测试突发请求场景
|
||||
rateIP := "192.168.1.210"
|
||||
windowIP := "192.168.1.211"
|
||||
|
||||
// 平均速率模式：突发8个请求
|
||||
rateAllowed := 0
|
||||
for i := 0; i < 8; i++ {
|
||||
if rateLimiter.GetLimiter(rateIP).Allow() {
|
||||
rateAllowed++
|
||||
}
|
||||
}
|
||||
|
||||
// 滑动窗口模式：突发8个请求
|
||||
windowAllowed := 0
|
||||
for i := 0; i < 8; i++ {
|
||||
if windowLimiter.Allow(windowIP) {
|
||||
windowAllowed++
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Printf("突发8个请求：平均速率模式允许%d个，滑动窗口模式允许%d个\n",
|
||||
rateAllowed, windowAllowed)
|
||||
|
||||
// 等待3秒
|
||||
time.Sleep(3 * time.Second)
|
||||
|
||||
// 平均速率模式：再发送5个请求
|
||||
rateAllowedLater := 0
|
||||
for i := 0; i < 5; i++ {
|
||||
if rateLimiter.GetLimiter(rateIP).Allow() {
|
||||
rateAllowedLater++
|
||||
}
|
||||
}
|
||||
|
||||
// 滑动窗口模式：再发送5个请求
|
||||
windowAllowedLater := 0
|
||||
for i := 0; i < 5; i++ {
|
||||
if windowLimiter.Allow(windowIP) {
|
||||
windowAllowedLater++
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Printf("3秒后再发送5个请求：平均速率模式允许%d个，滑动窗口模式允许%d个\n",
|
||||
rateAllowedLater, windowAllowedLater)
|
||||
|
||||
// 验证结果
|
||||
totalRateAllowed := rateAllowed + rateAllowedLater
|
||||
totalWindowAllowed := windowAllowed + windowAllowedLater
|
||||
|
||||
fmt.Printf("总计：平均速率模式允许%d个，滑动窗口模式允许%d个\n",
|
||||
totalRateAllowed, totalWindowAllowed)
|
||||
|
||||
// 修改：平均速率模式会随时间生成新令牌，所以允许的请求可能超过初始设定
|
||||
// 滑动窗口模式不应该超过最大请求数
|
||||
if totalWindowAllowed > maxRequests {
|
||||
t.Errorf("滑动窗口模式不应超过%d个请求，但实际允许了%d个",
|
||||
maxRequests, totalWindowAllowed)
|
||||
}
|
||||
|
||||
// 添加注释说明平均速率模式的行为特点，而不是将其视为错误
|
||||
fmt.Printf("注意：平均速率模式允许了%d个请求，超过了初始设定的%d个，这是因为随着时间推移会生成新的令牌\n",
|
||||
totalRateAllowed, maxRequests)
|
||||
})
|
||||
}
|
||||
|
||||
// 测试实际CC防护场景（滑动窗口模式）
|
||||
func TestWindowCCProtection(t *testing.T) {
|
||||
// 模拟配置：5秒内最多允许10个请求
|
||||
timeWindow := 5
|
||||
maxRequests := 10
|
||||
|
||||
// 创建滑动窗口限流器
|
||||
limiter := NewWindowIPRateLimiter(timeWindow, maxRequests)
|
||||
|
||||
// 测试IP
|
||||
ip := "192.168.1.220"
|
||||
|
||||
fmt.Printf("\n模拟CC防护（滑动窗口模式）：%d秒内最多允许%d个请求\n",
|
||||
timeWindow, maxRequests)
|
||||
|
||||
// 模拟正常用户：每秒2个请求
|
||||
t.Run("正常用户测试（滑动窗口）", func(t *testing.T) {
|
||||
// 重置限流器
|
||||
limiter.requests = make(map[string][]time.Time)
|
||||
limiter.ips = make(map[string]*rate.Limiter)
|
||||
|
||||
success := 0
|
||||
fail := 0
|
||||
|
||||
// 模拟5秒，每秒发送2个请求
|
||||
for second := 1; second <= timeWindow; second++ {
|
||||
for i := 0; i < 2; i++ {
|
||||
if limiter.Allow(ip) {
|
||||
success++
|
||||
} else {
|
||||
fail++
|
||||
}
|
||||
}
|
||||
time.Sleep(1 * time.Second)
|
||||
}
|
||||
|
||||
fmt.Printf("正常用户（每秒2个请求）：成功=%d, 失败=%d\n", success, fail)
|
||||
|
||||
// 正常用户应该不会被限流
|
||||
if fail > 0 {
|
||||
t.Errorf("正常用户不应该被限流，但有%d个请求被拒绝", fail)
|
||||
}
|
||||
})
|
||||
|
||||
// 模拟CC攻击：短时间内发送大量请求
|
||||
t.Run("CC攻击测试（滑动窗口）", func(t *testing.T) {
|
||||
// 重置限流器
|
||||
limiter.requests = make(map[string][]time.Time)
|
||||
limiter.ips = make(map[string]*rate.Limiter)
|
||||
|
||||
success := 0
|
||||
fail := 0
|
||||
|
||||
// 短时间内发送20个请求
|
||||
for i := 0; i < 20; i++ {
|
||||
if limiter.Allow(ip) {
|
||||
success++
|
||||
} else {
|
||||
fail++
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Printf("CC攻击（短时间内20个请求）：成功=%d, 失败=%d\n", success, fail)
|
||||
|
||||
// 应该有请求被限流
|
||||
if fail == 0 {
|
||||
t.Error("CC攻击应该被限流，但所有请求都成功了")
|
||||
}
|
||||
|
||||
// 成功的请求不应超过最大允许数
|
||||
if success > maxRequests {
|
||||
t.Errorf("成功请求数不应超过%d，但实际为%d", maxRequests, success)
|
||||
}
|
||||
})
|
||||
|
||||
// 测试滑动窗口的特性：请求逐渐滑出窗口
|
||||
t.Run("滑动窗口特性测试（CC防护）", func(t *testing.T) {
|
||||
// 重置限流器
|
||||
limiter.requests = make(map[string][]time.Time)
|
||||
limiter.ips = make(map[string]*rate.Limiter)
|
||||
|
||||
// 第1秒：发送10个请求（达到限制）
|
||||
firstBatchSuccess := 0
|
||||
for i := 0; i < 10; i++ {
|
||||
if limiter.Allow(ip) {
|
||||
firstBatchSuccess++
|
||||
}
|
||||
}
|
||||
|
||||
// 第1秒：再发送1个请求（应该被拒绝）
|
||||
if limiter.Allow(ip) {
|
||||
t.Error("达到限制后应该被拒绝")
|
||||
}
|
||||
|
||||
fmt.Printf("第1秒：尝试10个请求，允许了%d个，额外请求被拒绝\n", firstBatchSuccess)
|
||||
|
||||
// 等待3秒（窗口未完全滑动）
|
||||
time.Sleep(3 * time.Second)
|
||||
|
||||
// 第4秒：尝试发送5个请求（应该全部被拒绝，因为窗口内仍有10个请求）
|
||||
fourthSecondSuccess := 0
|
||||
for i := 0; i < 5; i++ {
|
||||
if limiter.Allow(ip) {
|
||||
fourthSecondSuccess++
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Printf("第4秒：尝试5个请求，允许了%d个\n", fourthSecondSuccess)
|
||||
|
||||
// 等待3秒（窗口完全滑动，第1秒的请求应该全部滑出）
|
||||
time.Sleep(3 * time.Second)
|
||||
|
||||
// 第7秒：应该可以发送10个新请求
|
||||
seventhSecondSuccess := 0
|
||||
for i := 0; i < 10; i++ {
|
||||
if limiter.Allow(ip) {
|
||||
seventhSecondSuccess++
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Printf("第7秒：尝试10个请求，允许了%d个\n", seventhSecondSuccess)
|
||||
|
||||
// 验证结果
|
||||
if seventhSecondSuccess != 10 {
|
||||
t.Errorf("窗口完全滑动后应该允许10个请求，但实际允许了%d个", seventhSecondSuccess)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// 测试清除指定IP的滑动窗口记录
|
||||
func TestClearWindowForIP(t *testing.T) {
|
||||
// 创建一个滑动窗口限流器
|
||||
timeWindow := 5
|
||||
maxRequests := 10
|
||||
limiter := NewWindowIPRateLimiter(timeWindow, maxRequests)
|
||||
|
||||
// 测试IP
|
||||
ip := "192.168.1.250"
|
||||
|
||||
t.Run("清除滑动窗口记录测试", func(t *testing.T) {
|
||||
// 重置限流器
|
||||
limiter.requests = make(map[string][]time.Time)
|
||||
limiter.ips = make(map[string]*rate.Limiter)
|
||||
|
||||
// 先发送8个请求
|
||||
for i := 0; i < 8; i++ {
|
||||
limiter.Allow(ip)
|
||||
}
|
||||
|
||||
// 验证请求记录数量
|
||||
if len(limiter.requests[ip]) != 8 {
|
||||
t.Errorf("应该记录8个请求，但实际记录了%d个", len(limiter.requests[ip]))
|
||||
}
|
||||
|
||||
// 清除该IP的记录
|
||||
limiter.ClearWindowForIP(ip)
|
||||
|
||||
// 验证记录是否被清除
|
||||
if len(limiter.requests[ip]) != 0 {
|
||||
t.Errorf("清除后应该有0个请求记录，但实际有%d个", len(limiter.requests[ip]))
|
||||
}
|
||||
|
||||
// 验证是否可以继续发送请求
|
||||
allowedAfterClear := 0
|
||||
for i := 0; i < maxRequests; i++ {
|
||||
if limiter.Allow(ip) {
|
||||
allowedAfterClear++
|
||||
}
|
||||
}
|
||||
|
||||
// 清除后应该可以发送最大请求数
|
||||
if allowedAfterClear != maxRequests {
|
||||
t.Errorf("清除后应该允许%d个请求，但实际允许了%d个", maxRequests, allowedAfterClear)
|
||||
}
|
||||
|
||||
// 超过限制后应该被拒绝
|
||||
if limiter.Allow(ip) {
|
||||
t.Error("超过限制后应该被拒绝")
|
||||
}
|
||||
})
|
||||
|
||||
// 测试平均速率模式下的清除功能
|
||||
t.Run("平均速率模式下清除记录测试", func(t *testing.T) {
|
||||
// 创建平均速率模式限流器
|
||||
ratePerSecond := rate.Limit(float64(maxRequests) / float64(timeWindow))
|
||||
rateLimiter := NewIPRateLimiter(ratePerSecond, maxRequests)
|
||||
|
||||
// 发送足够多的请求，消耗所有令牌
|
||||
for i := 0; i < maxRequests; i++ {
|
||||
rateLimiter.GetLimiter(ip).Allow()
|
||||
}
|
||||
|
||||
// 此时应该被限流
|
||||
if rateLimiter.GetLimiter(ip).Allow() {
|
||||
t.Error("消耗所有令牌后应该被限流")
|
||||
}
|
||||
|
||||
// 清除该IP的记录
|
||||
rateLimiter.ClearWindowForIP(ip)
|
||||
|
||||
// 清除后应该可以继续发送请求
|
||||
if !rateLimiter.GetLimiter(ip).Allow() {
|
||||
t.Error("清除后应该允许新的请求")
|
||||
}
|
||||
})
|
||||
|
||||
// 测试不存在的IP
|
||||
t.Run("清除不存在IP的记录", func(t *testing.T) {
|
||||
// 重置限流器
|
||||
limiter.requests = make(map[string][]time.Time)
|
||||
limiter.ips = make(map[string]*rate.Limiter)
|
||||
|
||||
// 清除一个不存在的IP
|
||||
nonExistIP := "10.0.0.1"
|
||||
|
||||
// 不应该抛出异常
|
||||
limiter.ClearWindowForIP(nonExistIP)
|
||||
|
||||
// 验证是否可以正常使用
|
||||
if !limiter.Allow(nonExistIP) {
|
||||
t.Error("清除不存在的IP后，应该允许该IP的请求")
|
||||
}
|
||||
})
|
||||
|
||||
// 测试实际场景：IP被限流后手动解除
|
||||
t.Run("IP被限流后手动解除", func(t *testing.T) {
|
||||
// 重置限流器
|
||||
limiter.requests = make(map[string][]time.Time)
|
||||
limiter.ips = make(map[string]*rate.Limiter)
|
||||
|
||||
fmt.Println("\n测试IP被限流后手动解除的场景...")
|
||||
|
||||
// 发送足够多的请求，触发限流
|
||||
allowedCount := 0
|
||||
for i := 0; i < maxRequests+1; i++ {
|
||||
if limiter.Allow(ip) {
|
||||
allowedCount++
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Printf("发送%d个请求，允许了%d个，最后一个被拒绝\n",
|
||||
maxRequests+1, allowedCount)
|
||||
|
||||
// 验证是否被限流
|
||||
if limiter.Allow(ip) {
|
||||
t.Error("应该被限流")
|
||||
}
|
||||
|
||||
// 手动解除限流
|
||||
fmt.Println("手动解除限流...")
|
||||
limiter.ClearWindowForIP(ip)
|
||||
|
||||
// 验证是否解除成功
|
||||
newAllowedCount := 0
|
||||
for i := 0; i < 5; i++ {
|
||||
if limiter.Allow(ip) {
|
||||
newAllowedCount++
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Printf("解除限流后，发送5个请求，允许了%d个\n", newAllowedCount)
|
||||
|
||||
if newAllowedCount != 5 {
|
||||
t.Errorf("解除限流后应该允许5个请求，但实际允许了%d个", newAllowedCount)
|
||||
}
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user