From a2550bc80e867d3e3f432a1b262065cf8c31eed2 Mon Sep 17 00:00:00 2001
From: user123456 <newemail@example.com>
Date: Wed, 11 Jun 2025 14:18:49 +0800
Subject: [PATCH] =?UTF-8?q?docker=E5=8C=BF=E5=90=8DToken=E7=BC=93=E5=AD=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/config.go      | 13 ++++++++
 src/config.toml    |  7 ++++
 src/docker.go      | 64 ++++++++++++++++++++++++++++++++++-
 src/token_cache.go | 83 ++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 166 insertions(+), 1 deletion(-)
 create mode 100644 src/token_cache.go

diff --git a/src/config.go b/src/config.go
index 4eb70c7..c9ccfc5 100644
--- a/src/config.go
+++ b/src/config.go
@@ -47,6 +47,12 @@ type AppConfig struct {
 
 	// 新增：Registry映射配置
 	Registries map[string]RegistryMapping `toml:"registries"`
+
+	// Token缓存配置
+	TokenCache struct {
+		Enabled    bool   `toml:"enabled"`    // 是否启用token缓存
+		DefaultTTL string `toml:"defaultTTL"` // 默认缓存时间
+	} `toml:"tokenCache"`
 }
 
 var (
@@ -118,6 +124,13 @@ func DefaultConfig() *AppConfig {
 				Enabled:  true,
 			},
 		},
+		TokenCache: struct {
+			Enabled    bool   `toml:"enabled"`
+			DefaultTTL string `toml:"defaultTTL"`
+		}{
+			Enabled:    true, // docker认证的匿名Token缓存配置，用于提升性能
+			DefaultTTL: "20m",
+		},
 	}
 }
 
diff --git a/src/config.toml b/src/config.toml
index e4696d9..dae8a22 100644
--- a/src/config.toml
+++ b/src/config.toml
@@ -81,3 +81,10 @@ enabled = true
 # authHost = "harbor.company.com/service/token"
 # authType = "basic"
 # enabled = false
+
+# docker的匿名Token缓存配置，显著提升性能
+[tokenCache]
+# 是否启用token缓存
+enabled = true
+# 默认缓存时间（自动从响应中提取实际TTL）
+defaultTTL = "20m"
diff --git a/src/docker.go b/src/docker.go
index a4468d0..003784f 100644
--- a/src/docker.go
+++ b/src/docker.go
@@ -305,8 +305,70 @@ func handleTagsRequest(c *gin.Context, imageRef string) {
 	c.JSON(http.StatusOK, response)
 }
 
-// ProxyDockerAuthGin Docker认证代理
+// ProxyDockerAuthGin Docker认证代理（带缓存优化）
 func ProxyDockerAuthGin(c *gin.Context) {
+	// 检查是否启用token缓存
+	if isTokenCacheEnabled() {
+		proxyDockerAuthWithCache(c)
+	} else {
+		proxyDockerAuthOriginal(c)
+	}
+}
+
+// proxyDockerAuthWithCache 带缓存的认证代理
+func proxyDockerAuthWithCache(c *gin.Context) {
+	// 1. 构建缓存key（基于完整的查询参数）
+	cacheKey := buildCacheKey(c.Request.URL.RawQuery)
+	
+	// 2. 尝试从缓存获取token
+	if cachedToken := globalTokenCache.Get(cacheKey); cachedToken != "" {
+		writeTokenResponse(c, cachedToken)
+		return
+	}
+	
+	// 3. 缓存未命中，创建响应记录器
+	recorder := &ResponseRecorder{
+		ResponseWriter: c.Writer,
+		statusCode:     200,
+	}
+	c.Writer = recorder
+	
+	// 4. 调用原有认证逻辑
+	proxyDockerAuthOriginal(c)
+	
+	// 5. 如果认证成功，缓存响应
+	if recorder.statusCode == 200 && len(recorder.body) > 0 {
+		ttl := extractTTLFromResponse(recorder.body)
+		globalTokenCache.Set(cacheKey, string(recorder.body), ttl)
+	}
+	
+	// 6. 写入实际响应（如果还没写入）
+	if !recorder.written {
+		c.Writer = recorder.ResponseWriter
+		c.Data(recorder.statusCode, "application/json", recorder.body)
+	}
+}
+
+// ResponseRecorder HTTP响应记录器
+type ResponseRecorder struct {
+	gin.ResponseWriter
+	statusCode int
+	body       []byte
+	written    bool
+}
+
+func (r *ResponseRecorder) WriteHeader(code int) {
+	r.statusCode = code
+}
+
+func (r *ResponseRecorder) Write(data []byte) (int, error) {
+	r.body = append(r.body, data...)
+	r.written = true
+	return r.ResponseWriter.Write(data)
+}
+
+// proxyDockerAuthOriginal Docker认证代理（原始逻辑，保持不变）
+func proxyDockerAuthOriginal(c *gin.Context) {
 	// 检查是否有目标Registry域名（来自Context）
 	var authURL string
 	if targetDomain, exists := c.Get("target_registry_domain"); exists {
diff --git a/src/token_cache.go b/src/token_cache.go
new file mode 100644
index 0000000..6041163
--- /dev/null
+++ b/src/token_cache.go
@@ -0,0 +1,83 @@
+package main
+
+import (
+	"crypto/md5"
+	"encoding/json"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/gin-gonic/gin"
+)
+
+// CachedToken 缓存的Token信息
+type CachedToken struct {
+	Token     string
+	ExpiresAt time.Time
+}
+
+// SimpleTokenCache 极简Token缓存
+type SimpleTokenCache struct {
+	cache sync.Map // 线程安全的并发映射
+}
+
+var globalTokenCache = &SimpleTokenCache{}
+
+// Get 获取缓存的token，如果不存在或过期返回空字符串
+func (c *SimpleTokenCache) Get(key string) string {
+	if v, ok := c.cache.Load(key); ok {
+		if cached := v.(*CachedToken); time.Now().Before(cached.ExpiresAt) {
+			return cached.Token
+		}
+		// 自动清理过期token，保持内存整洁
+		c.cache.Delete(key)
+	}
+	return ""
+}
+
+// Set 设置token缓存，自动计算过期时间
+func (c *SimpleTokenCache) Set(key, token string, ttl time.Duration) {
+	c.cache.Store(key, &CachedToken{
+		Token:     token,
+		ExpiresAt: time.Now().Add(ttl),
+	})
+}
+
+// buildCacheKey 构建稳定的缓存key
+func buildCacheKey(query string) string {
+	// 使用MD5确保key的一致性和简洁性
+	return fmt.Sprintf("token:%x", md5.Sum([]byte(query)))
+}
+
+// extractTTLFromResponse 从响应中智能提取TTL
+func extractTTLFromResponse(responseBody []byte) time.Duration {
+	var tokenResp struct {
+		ExpiresIn int `json:"expires_in"`
+	}
+	
+	// 默认30分钟TTL，确保稳定性
+	defaultTTL := 30 * time.Minute
+	
+	if json.Unmarshal(responseBody, &tokenResp) == nil && tokenResp.ExpiresIn > 0 {
+		// 使用响应中的过期时间，但提前5分钟过期确保安全边际
+		safeTTL := time.Duration(tokenResp.ExpiresIn-300) * time.Second
+		if safeTTL > 5*time.Minute { // 确保至少有5分钟的缓存时间
+			return safeTTL
+		}
+	}
+	
+	return defaultTTL
+}
+
+// writeTokenResponse 写入token响应
+func writeTokenResponse(c *gin.Context, cachedBody string) {
+	// 直接返回缓存的完整响应体，保持格式一致性
+	c.Header("Content-Type", "application/json")
+	c.String(200, cachedBody)
+}
+
+// isTokenCacheEnabled 检查token缓存是否启用
+func isTokenCacheEnabled() bool {
+	cfg := GetConfig()
+	return cfg.TokenCache.Enabled
+} 
\ No newline at end of file