diff --git a/src/ratelimiter.go b/src/ratelimiter.go index 620c41f..91db2a2 100644 --- a/src/ratelimiter.go +++ b/src/ratelimiter.go @@ -260,8 +260,11 @@ func RateLimitMiddleware(limiter *IPRateLimiter) gin.HandlerFunc { return } - // 检查是否允许本次请求 - if !ipLimiter.Allow() { + // 智能限流判断:检查是否应该跳过限流计数 + shouldSkip := smartLimiter.ShouldSkipRateLimit(cleanIP, c.Request.URL.Path) + + // 只有在不跳过的情况下才检查限流 + if !shouldSkip && !ipLimiter.Allow() { c.JSON(429, gin.H{ "error": "请求频率过快,暂时限制访问", }) diff --git a/src/smart_ratelimit.go b/src/smart_ratelimit.go new file mode 100644 index 0000000..436d5f4 --- /dev/null +++ b/src/smart_ratelimit.go @@ -0,0 +1,125 @@ +package main + +import ( + "strings" + "sync" + "time" +) + +// SmartRateLimit 智能限流会话管理 +type SmartRateLimit struct { + sessions sync.Map // IP -> *PullSession +} + +// PullSession Docker拉取会话 +type PullSession struct { + LastManifestTime time.Time + RequestCount int +} + +// 全局智能限流实例 +var smartLimiter = &SmartRateLimit{} + +// 硬编码的智能限流参数 - 无需配置管理 +const ( + // manifest请求后的活跃窗口时间 + activeWindowDuration = 5 * time.Minute + // 活跃窗口内最大免费blob请求数(防止滥用) + maxFreeBlobRequests = 100 + // 会话清理间隔 + sessionCleanupInterval = 10 * time.Minute + // 会话过期时间 + sessionExpireTime = 30 * time.Minute +) + +func init() { + // 启动会话清理协程 + go smartLimiter.cleanupSessions() +} + +// ShouldSkipRateLimit 判断是否应该跳过限流计数 +// 返回true表示跳过限流,false表示正常计入限流 +func (s *SmartRateLimit) ShouldSkipRateLimit(ip, path string) bool { + // 提取请求类型 + requestType, _ := parseRequestInfo(path) + + // 只对manifest和blob请求做智能处理 + if requestType != "manifests" && requestType != "blobs" { + return false // 其他请求正常计入限流 + } + + // 获取或创建会话 + sessionKey := ip + sessionInterface, _ := s.sessions.LoadOrStore(sessionKey, &PullSession{}) + session := sessionInterface.(*PullSession) + + now := time.Now() + + if requestType == "manifests" { + // manifest请求:始终计入限流,但更新会话状态 + session.LastManifestTime = now + session.RequestCount = 0 // 重置计数 + return false // manifest请求正常计入限流 + } + + // blob请求:检查是否在活跃窗口内 + if requestType == "blobs" { + // 检查是否在活跃拉取窗口内 + if !session.LastManifestTime.IsZero() && + now.Sub(session.LastManifestTime) <= activeWindowDuration { + + // 在活跃窗口内,检查是否超过最大免费请求数 + session.RequestCount++ + if session.RequestCount <= maxFreeBlobRequests { + return true // 跳过限流计数 + } + } + } + + return false // 正常计入限流 +} + +// parseRequestInfo 解析请求路径,提取请求类型和镜像引用 +func parseRequestInfo(path string) (requestType, imageRef string) { + // 清理路径前缀 + path = strings.TrimPrefix(path, "/v2/") + + // 查找manifest或blob路径 + if idx := strings.Index(path, "/manifests/"); idx != -1 { + return "manifests", path[:idx] + } + if idx := strings.Index(path, "/blobs/"); idx != -1 { + return "blobs", path[:idx] + } + if idx := strings.Index(path, "/tags/"); idx != -1 { + return "tags", path[:idx] + } + + return "unknown", "" +} + +// cleanupSessions 定期清理过期会话,防止内存泄露 +func (s *SmartRateLimit) cleanupSessions() { + ticker := time.NewTicker(sessionCleanupInterval) + defer ticker.Stop() + + for range ticker.C { + now := time.Now() + expiredKeys := make([]string, 0) + + // 找出过期的会话 + s.sessions.Range(func(key, value interface{}) bool { + session := value.(*PullSession) + if !session.LastManifestTime.IsZero() && + now.Sub(session.LastManifestTime) > sessionExpireTime { + expiredKeys = append(expiredKeys, key.(string)) + } + return true + }) + + // 删除过期会话 + for _, key := range expiredKeys { + s.sessions.Delete(key) + } + } +} \ No newline at end of file