From d565ec678a46b4141eb23e3bd990c8f02eafa418 Mon Sep 17 00:00:00 2001
From: XOF <xofine@noreply.localhost>
Date: Sun, 28 Dec 2025 14:41:03 +0800
Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0=20main.go?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 main.go | 38 ++++++++++++++++++++------------------
 1 file changed, 20 insertions(+), 18 deletions(-)

diff --git a/main.go b/main.go
index 8e93492..a3c3ffd 100644
--- a/main.go
+++ b/main.go
@@ -115,10 +115,8 @@ func preprocessArgs() {
 		newArgs = append(newArgs, arg)
 	}
 
-	if len(newArgs) > 0 {
+	if len(newArgs) > 1 {
 		os.Args = newArgs
-	} else {
-		logrus.Warn("命令行参数为空，使用原始参数")
 	}
 }
 
@@ -231,7 +229,9 @@ func handleRequest(w http.ResponseWriter, r *http.Request) {
 	// 健康检查
 	if path == "/health" || path == "/healthz" {
 		w.WriteHeader(http.StatusOK)
-		w.Write([]byte("OK"))
+		if _, err := w.Write([]byte("OK")); err != nil {
+			logrus.Errorf("健康检查响应失败: %v", err)
+		}
 		return
 	}
 
@@ -262,20 +262,26 @@ func handleRequest(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+func getCleanHost(r *http.Request) string {
+	host := r.Host
+	if idx := strings.Index(host, ":"); idx != -1 {
+		return host[:idx]
+	}
+	return host
+}
+
 func handleRegistryRequest(w http.ResponseWriter, r *http.Request) {
 	const targetHost = "registry-1.docker.io"
 
 	ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
 	defer cancel()
 
-	pathParts := strings.Split(r.URL.Path, "/")
-	v2PathParts := pathParts[2:]
-	pathString := strings.Join(v2PathParts, "/")
+	path := strings.TrimPrefix(r.URL.Path, "/v2/")
 
 	url := &url.URL{
 		Scheme:   "https",
 		Host:     targetHost,
-		Path:     "/v2/" + pathString,
+		Path:     "/v2/" + path,
 		RawQuery: r.URL.RawQuery,
 	}
 
@@ -304,7 +310,7 @@ func handleRegistryRequest(w http.ResponseWriter, r *http.Request) {
 	respHeaders := copyHeaders(resp.Header)
 
 	if respHeaders.Get("WWW-Authenticate") != "" {
-		currentDomain := r.Host
+		currentDomain := getCleanHost(r)
 		respHeaders.Set("WWW-Authenticate",
 			fmt.Sprintf(`Bearer realm="https://%s/auth/token", service="registry.docker.io"`, currentDomain))
 	}
@@ -334,14 +340,12 @@ func handleAuthRequest(w http.ResponseWriter, r *http.Request) {
 	ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
 	defer cancel()
 
-	pathParts := strings.Split(r.URL.Path, "/")
-	authPathParts := pathParts[2:]
-	pathString := strings.Join(authPathParts, "/")
+	path := strings.TrimPrefix(r.URL.Path, "/auth/")
 
 	url := &url.URL{
 		Scheme:   "https",
 		Host:     targetHost,
-		Path:     "/" + pathString,
+		Path:     "/" + path,
 		RawQuery: r.URL.RawQuery,
 	}
 
@@ -387,14 +391,12 @@ func handleCloudflareRequest(w http.ResponseWriter, r *http.Request) {
 	ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
 	defer cancel()
 
-	pathParts := strings.Split(r.URL.Path, "/")
-	cfPathParts := pathParts[2:]
-	pathString := strings.Join(cfPathParts, "/")
+	path := strings.TrimPrefix(r.URL.Path, "/production-cloudflare/")
 
 	url := &url.URL{
 		Scheme:   "https",
 		Host:     targetHost,
-		Path:     "/" + pathString,
+		Path:     "/" + path,
 		RawQuery: r.URL.RawQuery,
 	}
 
@@ -442,7 +444,7 @@ func handleAuthChallenge(w http.ResponseWriter, r *http.Request, resp *http.Resp
 	}
 
 	if authHeader := w.Header().Get("WWW-Authenticate"); authHeader != "" {
-		currentDomain := r.Host
+		currentDomain := getCleanHost(r)
 		w.Header().Set("WWW-Authenticate",
 			fmt.Sprintf(`Bearer realm="https://%s/auth/token", service="registry.docker.io"`, currentDomain))
 	}