From 512a3bb84a602f36317cc49e026ca191e723f7f5 Mon Sep 17 00:00:00 2001 From: samwaf Date: Fri, 28 Mar 2025 14:16:03 +0800 Subject: [PATCH] fix:cc protect rate/window mode #IBVCH3 --- api/waf_anticc.go | 1 + enums/task_enum.go | 1 + global/global.go | 3 + main.go | 9 + model/anticc.go | 24 +- model/request/waf_anticc_req.go | 28 +- service/waf_service/waf_anticc.go | 2 + wafdb/patch_sql.go | 7 + wafenginecore/checkcc.go | 15 +- wafenginecore/wafengine.go | 45 +- wafenginecore/wafworker.go | 15 +- waftask/entrance.go | 11 +- waftask/task_cc.go | 9 + webplugin/ipratelimiter.go | 164 ++++++- webplugin/ipratelimiter_test.go | 695 ++++++++++++++++++++++++++++++ 15 files changed, 985 insertions(+), 44 deletions(-) create mode 100644 waftask/task_cc.go create mode 100644 webplugin/ipratelimiter_test.go diff --git a/api/waf_anticc.go b/api/waf_anticc.go index ff35d27..dba1144 100644 --- a/api/waf_anticc.go +++ b/api/waf_anticc.go @@ -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) diff --git a/enums/task_enum.go b/enums/task_enum.go index a01396b..66cbe78 100644 --- a/enums/task_enum.go +++ b/enums/task_enum.go @@ -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记录 ) diff --git a/global/global.go b/global/global.go index c778b79..35ec076 100644 --- a/global/global.go +++ b/global/global.go @@ -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 = "" //微信访问密钥 diff --git a/main.go b/main.go index b664b40..30b248a 100644 --- a/main.go +++ b/main.go @@ -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 } } diff --git a/model/anticc.go b/model/anticc.go index 9ccce00..ba7b342 100644 --- a/model/anticc.go +++ b/model/anticc.go @@ -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"` //备注 } diff --git a/model/request/waf_anticc_req.go b/model/request/waf_anticc_req.go index 7fb5977..1e1987f 100644 --- a/model/request/waf_anticc_req.go +++ b/model/request/waf_anticc_req.go @@ -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 diff --git a/service/waf_service/waf_anticc.go b/service/waf_service/waf_anticc.go index de67663..8715065 100644 --- a/service/waf_service/waf_anticc.go +++ b/service/waf_service/waf_anticc.go @@ -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 diff --git a/wafdb/patch_sql.go b/wafdb/patch_sql.go index dfe4b42..e58db40 100644 --- a/wafdb/patch_sql.go +++ b/wafdb/patch_sql.go @@ -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()) diff --git a/wafenginecore/checkcc.go b/wafenginecore/checkcc.go index f1c35f1..c0abb8a 100644 --- a/wafenginecore/checkcc.go +++ b/wafenginecore/checkcc.go @@ -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 diff --git a/wafenginecore/wafengine.go b/wafenginecore/wafengine.go index 454b854..d524faf 100644 --- a/wafenginecore/wafengine.go +++ b/wafenginecore/wafengine.go @@ -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)) +} diff --git a/wafenginecore/wafworker.go b/wafenginecore/wafworker.go index afc4474..ae40fa8 100644 --- a/wafenginecore/wafworker.go +++ b/wafenginecore/wafworker.go @@ -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白名单 diff --git a/waftask/entrance.go b/waftask/entrance.go index 277c233..4aab9ec 100644 --- a/waftask/entrance.go +++ b/waftask/entrance.go @@ -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 } diff --git a/waftask/task_cc.go b/waftask/task_cc.go new file mode 100644 index 0000000..dec88d9 --- /dev/null +++ b/waftask/task_cc.go @@ -0,0 +1,9 @@ +package waftask + +import ( + "SamWaf/global" +) + +func TaskCC() { + global.GWAF_CHAN_CLEAR_CC_WINDOWS <- 1 +} diff --git a/webplugin/ipratelimiter.go b/webplugin/ipratelimiter.go index 1303656..168a691 100644 --- a/webplugin/ipratelimiter.go +++ b/webplugin/ipratelimiter.go @@ -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 +} diff --git a/webplugin/ipratelimiter_test.go b/webplugin/ipratelimiter_test.go new file mode 100644 index 0000000..7d84d66 --- /dev/null +++ b/webplugin/ipratelimiter_test.go @@ -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) + } + }) +}