标记匿名token区分匿名manifest

src/handlers/docker.go

@@ -1,11 +1,16 @@
 package handlers
 
 import (
+	"crypto/sha256"
+	"encoding/hex"
+	"encoding/json"
 	"fmt"
 	"io"
 	"net/http"
 	"net/url"
 	"strings"
+	"sync"
+	"time"
 
 	"github.com/gin-gonic/gin"
 	"hubproxy/config"
@@ -26,6 +31,7 @@ const (
 	dockerHubAuthRealm    = "https://auth.docker.io/token"
 	dockerHubAuthService  = "registry.docker.io"
 	maxCachedManifestSize = 4 << 20
+	maxAnonymousTokens    = 4096
 )
 
 var hopByHopHeaders = map[string]struct{}{
@@ -50,6 +56,13 @@ var forwardedRequestHeaders = []string{
 	"If-Unmodified-Since",
 }
 
+type anonymousTokenStore struct {
+	mu      sync.Mutex
+	entries map[string]time.Time
+}
+
+var anonymousTokens = &anonymousTokenStore{entries: make(map[string]time.Time)}
+
 // 保留初始化入口，在线代理无状态。
 func InitDockerProxy() {}
 
@@ -276,12 +289,83 @@ func ProxyDockerAuthGin(c *gin.Context) {
 
 	copyResponseHeaders(c, resp.Header, target)
 	if cacheable && resp.StatusCode == http.StatusOK && len(body) > 0 {
-		utils.GlobalCache.SetToken(cacheKey, string(body), utils.ExtractTTLFromResponse(body))
+		ttl := utils.ExtractTTLFromResponse(body)
+		utils.GlobalCache.SetToken(cacheKey, string(body), ttl)
+		anonymousTokens.RememberFromResponse(body, ttl)
 	}
 
 	c.Data(resp.StatusCode, resp.Header.Get("Content-Type"), body)
 }
 
+func (s *anonymousTokenStore) RememberFromResponse(body []byte, ttl time.Duration) {
+	token := tokenFromAuthResponse(body)
+	if token == "" || ttl <= 0 {
+		return
+	}
+
+	now := time.Now()
+	expiresAt := now.Add(ttl)
+	key := tokenHash(token)
+
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.cleanupLocked(now)
+	if len(s.entries) >= maxAnonymousTokens {
+		return
+	}
+	s.entries[key] = expiresAt
+}
+
+func (s *anonymousTokenStore) IsKnown(token string) bool {
+	if token == "" {
+		return false
+	}
+
+	key := tokenHash(token)
+	now := time.Now()
+
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	expiresAt, ok := s.entries[key]
+	if !ok {
+		return false
+	}
+	if now.After(expiresAt) {
+		delete(s.entries, key)
+		return false
+	}
+	return true
+}
+
+func (s *anonymousTokenStore) cleanupLocked(now time.Time) {
+	for key, expiresAt := range s.entries {
+		if now.After(expiresAt) {
+			delete(s.entries, key)
+		}
+	}
+}
+
+func tokenFromAuthResponse(body []byte) string {
+	var tokenResp struct {
+		Token       string `json:"token"`
+		AccessToken string `json:"access_token"`
+	}
+	if err := json.Unmarshal(body, &tokenResp); err != nil {
+		return ""
+	}
+	if tokenResp.Token != "" {
+		return tokenResp.Token
+	}
+	return tokenResp.AccessToken
+}
+
+func tokenHash(token string) string {
+	sum := sha256.Sum256([]byte(token))
+	return hex.EncodeToString(sum[:])
+}
+
 func buildAuthURL(target registryTarget, rawQuery string) (string, error) {
 	authURL, err := url.Parse(target.AuthRealm)
 	if err != nil {
@@ -401,7 +485,7 @@ func proxyRegistryHTTP(c *gin.Context, target registryTarget, upstreamPath strin
 }
 
 func manifestCacheKey(c *gin.Context, target registryTarget, upstreamPath string) (string, bool) {
-	if c.Request.Method != http.MethodGet || c.GetHeader("Authorization") != "" || !utils.IsCacheEnabled() {
+	if c.Request.Method != http.MethodGet || !utils.IsCacheEnabled() || !isAnonymousManifestRequest(c) {
 		return "", false
 	}
 	if !strings.Contains(upstreamPath, "/manifests/") {
@@ -417,6 +501,21 @@ func manifestCacheKey(c *gin.Context, target registryTarget, upstreamPath string
 	return utils.BuildCacheKey("manifest", key), true
 }
 
+func isAnonymousManifestRequest(c *gin.Context) bool {
+	authHeader := strings.TrimSpace(c.GetHeader("Authorization"))
+	if authHeader == "" {
+		return true
+	}
+
+	const bearerPrefix = "bearer "
+	if len(authHeader) <= len(bearerPrefix) || !strings.EqualFold(authHeader[:len(bearerPrefix)], bearerPrefix) {
+		return false
+	}
+
+	token := strings.TrimSpace(authHeader[len(bearerPrefix):])
+	return anonymousTokens.IsKnown(token)
+}
+
 func canCacheManifestResponse(resp *http.Response) bool {
 	return resp.StatusCode == http.StatusOK &&
 		resp.ContentLength > 0 &&

src/handlers/docker_test.go

@@ -11,6 +11,7 @@ import (
 	"sync"
 	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/gin-gonic/gin"
 	"hubproxy/config"
@@ -621,6 +622,93 @@ enabled = true
 	}
 }
 
+func TestProxyDockerRegistryCachesKnownAnonymousBearerManifest(t *testing.T) {
+	anonymousTokens = &anonymousTokenStore{entries: make(map[string]time.Time)}
+	anonymousTokens.RememberFromResponse([]byte(`{"token":"anonymous-token","expires_in":3600}`), time.Hour)
+
+	var hits int32
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		count := atomic.AddInt32(&hits, 1)
+		body := fmt.Sprintf(`{"schemaVersion":2,"hit":%d}`, count)
+		w.Header().Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json")
+		w.Header().Set("Content-Length", fmt.Sprintf("%d", len(body)))
+		_, _ = w.Write([]byte(body))
+	}))
+	defer upstream.Close()
+
+	initDockerProxyTest(t, `
+[registries."test.local"]
+upstream = "`+upstream.URL+`"
+authHost = "https://auth.test.local/token"
+authType = "anonymous"
+enabled = true
+`)
+	utils.GlobalCache = &utils.UniversalCache{}
+
+	gin.SetMode(gin.TestMode)
+	router := gin.New()
+	router.Any("/v2/*path", ProxyDockerRegistryGin)
+
+	for i := 0; i < 2; i++ {
+		req := httptest.NewRequest(http.MethodGet, "/v2/test.local/team/app/manifests/latest", nil)
+		req.Header.Set("Authorization", "Bearer anonymous-token")
+		w := httptest.NewRecorder()
+		router.ServeHTTP(w, req)
+		if w.Code != http.StatusOK {
+			t.Fatalf("request %d status = %d; body=%s", i, w.Code, w.Body.String())
+		}
+		if !strings.Contains(w.Body.String(), `"hit":1`) {
+			t.Fatalf("request %d body = %q", i, w.Body.String())
+		}
+	}
+	if got := atomic.LoadInt32(&hits); got != 1 {
+		t.Fatalf("hits = %d, want 1", got)
+	}
+}
+
+func TestProxyDockerRegistryDoesNotCacheUnknownBearerManifest(t *testing.T) {
+	anonymousTokens = &anonymousTokenStore{entries: make(map[string]time.Time)}
+
+	var hits int32
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		count := atomic.AddInt32(&hits, 1)
+		body := fmt.Sprintf(`{"schemaVersion":2,"hit":%d}`, count)
+		w.Header().Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json")
+		w.Header().Set("Content-Length", fmt.Sprintf("%d", len(body)))
+		_, _ = w.Write([]byte(body))
+	}))
+	defer upstream.Close()
+
+	initDockerProxyTest(t, `
+[registries."test.local"]
+upstream = "`+upstream.URL+`"
+authHost = "https://auth.test.local/token"
+authType = "anonymous"
+enabled = true
+`)
+	utils.GlobalCache = &utils.UniversalCache{}
+
+	gin.SetMode(gin.TestMode)
+	router := gin.New()
+	router.Any("/v2/*path", ProxyDockerRegistryGin)
+
+	for i := 1; i <= 2; i++ {
+		req := httptest.NewRequest(http.MethodGet, "/v2/test.local/team/app/manifests/latest", nil)
+		req.Header.Set("Authorization", "Bearer user-token")
+		w := httptest.NewRecorder()
+		router.ServeHTTP(w, req)
+		if w.Code != http.StatusOK {
+			t.Fatalf("request %d status = %d; body=%s", i, w.Code, w.Body.String())
+		}
+		if !strings.Contains(w.Body.String(), fmt.Sprintf(`"hit":%d`, i)) {
+			t.Fatalf("request %d body = %q", i, w.Body.String())
+		}
+	}
+	if got := atomic.LoadInt32(&hits); got != 2 {
+		t.Fatalf("hits = %d, want 2", got)
+	}
+}
+
 func TestProxyDockerRegistryUsesNsQueryForContainerd(t *testing.T) {
 	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if r.URL.Path != "/v2/team/app/manifests/latest" {