diff --git a/proxy/docker.go b/proxy/docker.go index 4712ae3..1471171 100644 --- a/proxy/docker.go +++ b/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)