108 lines
2.4 KiB
Go
108 lines
2.4 KiB
Go
package main
|
|
|
|
import (
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// SmartRateLimit 智能限流会话管理
|
|
type SmartRateLimit struct {
|
|
sessions sync.Map
|
|
}
|
|
|
|
// PullSession Docker拉取会话
|
|
type PullSession struct {
|
|
LastManifestTime time.Time
|
|
RequestCount int
|
|
}
|
|
|
|
// 全局智能限流实例
|
|
var smartLimiter = &SmartRateLimit{}
|
|
|
|
const (
|
|
// manifest请求后的活跃窗口时间
|
|
activeWindowDuration = 3 * time.Minute
|
|
// 活跃窗口内最大免费blob请求数(防止滥用)
|
|
maxFreeBlobRequests = 100
|
|
sessionCleanupInterval = 10 * time.Minute
|
|
sessionExpireTime = 30 * time.Minute
|
|
)
|
|
|
|
func init() {
|
|
go smartLimiter.cleanupSessions()
|
|
}
|
|
|
|
// ShouldSkipRateLimit 判断是否应该跳过限流计数
|
|
func (s *SmartRateLimit) ShouldSkipRateLimit(ip, path string) bool {
|
|
requestType, _ := parseRequestInfo(path)
|
|
|
|
if requestType != "manifests" && requestType != "blobs" {
|
|
return false
|
|
}
|
|
|
|
sessionKey := ip
|
|
sessionInterface, _ := s.sessions.LoadOrStore(sessionKey, &PullSession{})
|
|
session := sessionInterface.(*PullSession)
|
|
|
|
now := time.Now()
|
|
|
|
if requestType == "manifests" {
|
|
session.LastManifestTime = now
|
|
session.RequestCount = 0
|
|
return false
|
|
}
|
|
|
|
if requestType == "blobs" {
|
|
if !session.LastManifestTime.IsZero() &&
|
|
now.Sub(session.LastManifestTime) <= activeWindowDuration {
|
|
|
|
session.RequestCount++
|
|
if session.RequestCount <= maxFreeBlobRequests {
|
|
return true
|
|
}
|
|
}
|
|
}
|
|
|
|
return false
|
|
}
|
|
|
|
func parseRequestInfo(path string) (requestType, imageRef string) {
|
|
path = strings.TrimPrefix(path, "/v2/")
|
|
|
|
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)
|
|
}
|
|
}
|
|
} |