diff --git a/src/handlers/docker.go b/src/handlers/docker.go index bb7c0fe..2ef0dd5 100644 --- a/src/handlers/docker.go +++ b/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 && diff --git a/src/handlers/docker_test.go b/src/handlers/docker_test.go index daa554b..846dc0b 100644 --- a/src/handlers/docker_test.go +++ b/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" {