From fc77ddb1ef4323adae8c74f9603cfb5af414e957 Mon Sep 17 00:00:00 2001 From: sky22333 Date: Sat, 16 May 2026 05:39:54 +0800 Subject: [PATCH] =?UTF-8?q?=E9=99=8D=E4=BD=8E=E6=97=A5=E5=BF=97IO=20+=20?= =?UTF-8?q?=E5=8C=BF=E5=90=8Dmanifest=E7=BC=93=E5=AD=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/handlers/docker.go | 68 +++++++++++++++++++++++++-- src/handlers/docker_test.go | 93 +++++++++++++++++++++++++++++++++++++ src/utils/ratelimiter.go | 28 ++++++----- 3 files changed, 174 insertions(+), 15 deletions(-) diff --git a/src/handlers/docker.go b/src/handlers/docker.go index 30d3f0c..bb7c0fe 100644 --- a/src/handlers/docker.go +++ b/src/handlers/docker.go @@ -21,10 +21,11 @@ type registryTarget struct { } const ( - dockerHubName = "docker.io" - dockerHubUpstream = "https://registry-1.docker.io" - dockerHubAuthRealm = "https://auth.docker.io/token" - dockerHubAuthService = "registry.docker.io" + dockerHubName = "docker.io" + dockerHubUpstream = "https://registry-1.docker.io" + dockerHubAuthRealm = "https://auth.docker.io/token" + dockerHubAuthService = "registry.docker.io" + maxCachedManifestSize = 4 << 20 ) var hopByHopHeaders = map[string]struct{}{ @@ -349,6 +350,13 @@ func addLibraryPrefixToScope(scope string) string { } func proxyRegistryHTTP(c *gin.Context, target registryTarget, upstreamPath string) { + if cacheKey, ok := manifestCacheKey(c, target, upstreamPath); ok { + if cachedItem := utils.GlobalCache.Get(cacheKey); cachedItem != nil { + utils.WriteCachedResponse(c, cachedItem) + return + } + } + targetURL := target.Upstream + upstreamPath if c.Request.URL.RawQuery != "" { targetURL += "?" + c.Request.URL.RawQuery @@ -374,11 +382,63 @@ func proxyRegistryHTTP(c *gin.Context, target registryTarget, upstreamPath strin if c.Request.Method == http.MethodHead { return } + + if cacheKey, ok := manifestCacheKey(c, target, upstreamPath); ok && canCacheManifestResponse(resp) { + body, err := io.ReadAll(resp.Body) + if err != nil { + fmt.Printf("Failed to read manifest response: %v\n", err) + return + } + contentType := resp.Header.Get("Content-Type") + utils.GlobalCache.Set(cacheKey, body, contentType, cacheHeaders(resp.Header), utils.GetManifestTTL(manifestReference(upstreamPath))) + c.Data(resp.StatusCode, contentType, body) + return + } + if _, err := io.Copy(c.Writer, resp.Body); err != nil { fmt.Printf("Failed to stream registry response: %v\n", err) } } +func manifestCacheKey(c *gin.Context, target registryTarget, upstreamPath string) (string, bool) { + if c.Request.Method != http.MethodGet || c.GetHeader("Authorization") != "" || !utils.IsCacheEnabled() { + return "", false + } + if !strings.Contains(upstreamPath, "/manifests/") { + return "", false + } + + key := strings.Join([]string{ + target.Name, + upstreamPath, + c.Request.URL.RawQuery, + strings.Join(c.Request.Header.Values("Accept"), ","), + }, "|") + return utils.BuildCacheKey("manifest", key), true +} + +func canCacheManifestResponse(resp *http.Response) bool { + return resp.StatusCode == http.StatusOK && + resp.ContentLength > 0 && + resp.ContentLength <= maxCachedManifestSize +} + +func manifestReference(upstreamPath string) string { + _, _, reference := parseRegistryPath(strings.TrimPrefix(upstreamPath, "/v2/")) + return reference +} + +func cacheHeaders(headers http.Header) map[string]string { + cached := make(map[string]string) + for name, values := range headers { + if shouldSkipResponseHeader(name) || strings.EqualFold(name, "WWW-Authenticate") || len(values) == 0 { + continue + } + cached[name] = values[0] + } + return cached +} + func forwardSelectedRequestHeaders(dst http.Header, src http.Header) { for _, name := range forwardedRequestHeaders { for _, value := range src.Values(name) { diff --git a/src/handlers/docker_test.go b/src/handlers/docker_test.go index 4f69479..daa554b 100644 --- a/src/handlers/docker_test.go +++ b/src/handlers/docker_test.go @@ -528,6 +528,99 @@ enabled = true } } +func TestProxyDockerRegistryCachesAnonymousManifestByAccept(t *testing.T) { + 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", r.Header.Get("Accept")) + 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("Accept", "application/vnd.docker.distribution.manifest.v2+json") + 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) + } + + req := httptest.NewRequest(http.MethodGet, "/v2/test.local/team/app/manifests/latest", nil) + req.Header.Set("Accept", "application/vnd.oci.image.index.v1+json") + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("second accept status = %d; body=%s", w.Code, w.Body.String()) + } + if !strings.Contains(w.Body.String(), `"hit":2`) { + t.Fatalf("second accept body = %q", w.Body.String()) + } +} + +func TestProxyDockerRegistryDoesNotCacheAuthenticatedManifest(t *testing.T) { + 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 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" { diff --git a/src/utils/ratelimiter.go b/src/utils/ratelimiter.go index 3512cc9..5f01dbd 100644 --- a/src/utils/ratelimiter.go +++ b/src/utils/ratelimiter.go @@ -3,6 +3,7 @@ package utils import ( "fmt" "net" + "os" "strings" "sync" "time" @@ -17,6 +18,8 @@ const ( MaxIPCacheSize = 10000 ) +var debugRateLimitLog = strings.EqualFold(os.Getenv("DEBUG_RATE_LIMIT_LOG"), "true") + // IPRateLimiter IP限流器结构体 type IPRateLimiter struct { ips map[string]*rateLimiterEntry @@ -234,17 +237,20 @@ func RateLimitMiddleware(limiter *IPRateLimiter) gin.HandlerFunc { cleanIP := extractIPFromAddress(ip) - normalizedIP := normalizeIPForRateLimit(cleanIP) - if cleanIP != normalizedIP { - fmt.Printf("请求IP: %s (提纯后: %s, 限流段: %s), X-Forwarded-For: %s, X-Real-IP: %s\n", - ip, cleanIP, normalizedIP, - c.GetHeader("X-Forwarded-For"), - c.GetHeader("X-Real-IP")) - } else { - fmt.Printf("请求IP: %s (提纯后: %s), X-Forwarded-For: %s, X-Real-IP: %s\n", - ip, cleanIP, - c.GetHeader("X-Forwarded-For"), - c.GetHeader("X-Real-IP")) + if debugRateLimitLog { + normalizedIP := normalizeIPForRateLimit(cleanIP) + if cleanIP != normalizedIP { + fmt.Printf("请求IP: %s (提纯后: %s, 限流段: %s), X-Forwarded-For: %s, X-Real-IP: %s\n", + ip, cleanIP, normalizedIP, + c.GetHeader("X-Forwarded-For"), + c.GetHeader("X-Real-IP")) + } else { + fmt.Printf("请求IP: %s (提纯后: %s), X-Forwarded-For: %s, X-Real-IP: %s\n", + ip, cleanIP, + c.GetHeader("X-Forwarded-For"), + c.GetHeader("X-Real-IP")) + } + } ipLimiter, allowed := limiter.GetLimiter(cleanIP)