proxy/docker.go

@@ -3,6 +3,9 @@ package proxy
 import (
 	"bytes"
 	"context"
+
+	"github.com/go-json-experiment/json"
+
 	"fmt"
 	"net/http"
 	"net/url"
@@ -14,7 +17,6 @@ import (
 
 	"github.com/WJQSERVER-STUDIO/go-utils/iox"
 	"github.com/WJQSERVER-STUDIO/go-utils/limitreader"
-	"github.com/go-json-experiment/json"
 	"github.com/infinite-iroha/touka"
 )
 
@@ -127,6 +129,7 @@ func GhcrToTarget(c *touka.Context, cfg *config.Config, target string, path stri
 
 // GhcrRequest 执行对Docker注册表的HTTP请求, 处理认证和重定向
 func GhcrRequest(ctx context.Context, c *touka.Context, u string, image *imageInfo, cfg *config.Config, target string) {
+
 	var (
 		method string
 		req    *http.Request
@@ -134,6 +137,17 @@ func GhcrRequest(ctx context.Context, c *touka.Context, u string, image *imageIn
 		err    error
 	)
 
+	// 当请求上下文被取消时, 确保关闭响应和请求体
+	go func() {
+		<-ctx.Done()
+		if resp != nil && resp.Body != nil {
+			_ = resp.Body.Close()
+		}
+		if req != nil && req.Body != nil {
+			_ = req.Body.Close()
+		}
+	}()
+
 	method = c.Request.Method
 	ghcrclient := c.GetHTTPC()
 	bodyByte, err := c.GetReqBodyFull()
@@ -186,10 +200,10 @@ func GhcrRequest(ctx context.Context, c *touka.Context, u string, image *imageIn
 		shouldRetry := string(c.GetRequestURIPath()) != "/v2/"
 		originalStatusCode := resp.StatusCode
 		c.Debugf("Initial request failed with status %d. Retry eligibility: %t", originalStatusCode, shouldRetry)
+		_ = resp.Body.Close() // 关闭当前响应体
 
 		if shouldRetry {
 			if image == nil {
-				_ = resp.Body.Close() // 终止流程, 关闭当前响应体
 				ErrorPage(c, NewErrorWithStatusLookup(originalStatusCode, "Unauthorized"))
 				return
 			}
@@ -198,12 +212,6 @@ func GhcrRequest(ctx context.Context, c *touka.Context, u string, image *imageIn
 
 			if token != "" {
 				c.Debugf("Successfully obtained auth token. Retrying request.")
-				_ = resp.Body.Close() // 在发起重试请求前, 关闭旧的响应体
-
-				// 更新kv
-				c.Debugf("Update Cache Token: %s", token)
-				cache.Put(image.Image, token)
-
 				// 重新构建并发送请求
 				rb_retry := ghcrclient.NewRequestBuilder(method, u)
 				rb_retry.NoDefaultHeaders()
@@ -235,7 +243,6 @@ func GhcrRequest(ctx context.Context, c *touka.Context, u string, image *imageIn
 				resp = resp_retry // 更新响应为重试后的响应
 			} else {
 				c.Warnf("Failed to obtain auth token. Cannot retry.")
-				// 获取令牌失败, 将继续处理原始的401/404响应, 其响应体仍然打开
 			}
 		}
 	}
@@ -244,14 +251,12 @@ func GhcrRequest(ctx context.Context, c *touka.Context, u string, image *imageIn
 	if resp.StatusCode == http.StatusFound || resp.StatusCode == http.StatusTemporaryRedirect {
 		location := resp.Header.Get("Location")
 		if location == "" {
-			_ = resp.Body.Close() // 终止流程, 关闭当前响应体
 			HandleError(c, "Redirect response missing Location header")
 			return
 		}
 
 		redirectURL, err := url.Parse(location)
 		if err != nil {
-			_ = resp.Body.Close() // 终止流程, 关闭当前响应体
 			HandleError(c, fmt.Sprintf("Failed to parse redirect location: %v", err))
 			return
 		}
@@ -264,7 +269,7 @@ func GhcrRequest(ctx context.Context, c *touka.Context, u string, image *imageIn
 		}
 
 		c.Debugf("Handling redirect. Status: %d, Final Location: %s", resp.StatusCode, redirectURL.String())
-		_ = resp.Body.Close() // 明确关闭重定向响应的响应体, 因为我们将发起新请求
+		_ = resp.Body.Close() // 关闭当前响应体
 
 		// 创建并发送重定向请求, 通常使用 GET 方法
 		redirectReq, err := http.NewRequestWithContext(ctx, "GET", redirectURL.String(), nil)
@@ -286,13 +291,13 @@ func GhcrRequest(ctx context.Context, c *touka.Context, u string, image *imageIn
 
 	// 如果最终响应是 404, 则读取响应体并返回自定义错误页面
 	if resp.StatusCode == 404 {
-		defer resp.Body.Close() // 使用defer确保在函数返回前关闭响应体
 		bodyBytes, err := iox.ReadAll(resp.Body)
 		if err != nil {
 			c.Warnf("Failed to read upstream 404 response body: %v", err)
 		} else {
 			c.Warnf("Upstream 404 response body: %s", string(bodyBytes))
 		}
+		_ = resp.Body.Close()
 		ErrorPage(c, NewErrorWithStatusLookup(404, "Page Not Found (From Upstream)"))
 		return
 	}
@@ -316,7 +321,7 @@ func GhcrRequest(ctx context.Context, c *touka.Context, u string, image *imageIn
 		// 如果内容大小超出限制, 返回 301 重定向到原始上游URL
 		if err == nil && bodySize > sizelimit {
 			finalURL := resp.Request.URL.String()
-			_ = resp.Body.Close() // 明确关闭响应体, 因为我们将重定向而不是流式传输
+			_ = resp.Body.Close() // 关闭响应体
 			c.Redirect(301, finalURL)
 			c.Warnf("%s %s %s %s %s Final-URL: %s Size-Limit-Exceeded: %d", c.ClientIP(), c.Request.Method, c.Request.URL.Path, c.UserAgent(), c.Request.Proto, finalURL, bodySize)
 			return
@@ -327,7 +332,6 @@ func GhcrRequest(ctx context.Context, c *touka.Context, u string, image *imageIn
 	c.SetHeaders(resp.Header)
 	// 设置客户端响应状态码
 	c.Status(resp.StatusCode)
-	// bodyReader 的所有权将转移给 SetBodyStream, 不再由此函数管理关闭
 	bodyReader := resp.Body
 
 	// 如果启用了带宽限制, 则使用限速读取器
@@ -340,7 +344,7 @@ func GhcrRequest(ctx context.Context, c *touka.Context, u string, image *imageIn
 		c.SetBodyStream(bodyReader, bodySize)
 		return
 	}
-	c.SetBodyStream(bodyReader, -1)
+	c.SetBodyStream(bodyReader, -1) // Content-Length 未知
 }
 
 // AuthToken 用于解析认证响应中的令牌
@@ -359,6 +363,7 @@ func ChallengeReq(target string, image *imageInfo, ctx context.Context, c *touka
 	rb401 := ghcrclient.NewRequestBuilder("GET", "https://"+target+"/v2/")
 	rb401.NoDefaultHeaders()
 	rb401.WithContext(ctx)
+	//rb401.AddHeader("User-Agent", "docker/28.1.1 go/go1.23.8 git-commit/01f442b kernel/6.12.25-amd64 os/linux arch/amd64 UpstreamClient(Docker-Client/28.1.1 ")
 	req401, err = rb401.Build()
 	if err != nil {
 		HandleError(c, fmt.Sprintf("Failed to create request: %v", err))
@@ -371,7 +376,9 @@ func ChallengeReq(target string, image *imageInfo, ctx context.Context, c *touka
 		HandleError(c, fmt.Sprintf("Failed to send request: %v", err))
 		return
 	}
-	defer resp401.Body.Close() // 确保响应体关闭
+	defer func() {
+		_ = resp401.Body.Close() // 确保响应体关闭
+	}()
 
 	// 解析 Www-Authenticate 头部, 获取认证领域和参数
 	bearer, err := parseBearerWWWAuthenticateHeader(resp401.Header.Get("Www-Authenticate"))
@@ -387,6 +394,7 @@ func ChallengeReq(target string, image *imageInfo, ctx context.Context, c *touka
 	getAuthRB := ghcrclient.NewRequestBuilder("GET", bearer.Realm).
 		NoDefaultHeaders().
 		WithContext(ctx).
+		//AddHeader("User-Agent", "docker/28.1.1 go/go1.23.8 git-commit/01f442b kernel/6.12.25-amd64 os/linux arch/amd64 UpstreamClient(Docker-Client/28.1.1 ").
 		SetHeader("Host", bearer.Service).
 		AddQueryParam("service", bearer.Service).
 		AddQueryParam("scope", scope)
@@ -402,7 +410,9 @@ func ChallengeReq(target string, image *imageInfo, ctx context.Context, c *touka
 		c.Errorf("Failed to send request: %v", err)
 		return
 	}
-	defer authResp.Body.Close() // 确保响应体关闭
+	defer func() {
+		_ = authResp.Body.Close() // 确保响应体关闭
+	}()
 
 	// 读取认证响应体
 	bodyBytes, err := iox.ReadAll(authResp.Body)