更新 main.go

main.go

@@ -3,6 +3,7 @@ package main
 
 import (
 	"context"
+	"encoding/json"
 	"flag"
 	"fmt"
 	"io"
@@ -12,16 +13,15 @@ import (
 	"os/signal"
 	"strconv"
 	"strings"
+	"sync"
 	"syscall"
 	"time"
 
 	"github.com/sirupsen/logrus"
 )
 
-// Version 用于嵌入构建版本号
 var Version = "dev"
 
-// Config 定义配置结构体
 type Config struct {
 	ListenAddress string
 	Port          int
@@ -33,26 +33,103 @@ var config Config
 
 var client = &http.Client{
 	CheckRedirect: func(req *http.Request, via []*http.Request) error {
-		// 不跟随重定向，让 Docker 客户端自己处理
 		return http.ErrUseLastResponse
 	},
 	Timeout: 30 * time.Second,
 	Transport: &http.Transport{
 		DisableKeepAlives:     false,
 		MaxIdleConns:          100,
+		MaxIdleConnsPerHost:   10,
 		IdleConnTimeout:       90 * time.Second,
 		TLSHandshakeTimeout:   10 * time.Second,
 		ExpectContinueTimeout: 1 * time.Second,
 	},
 }
 
+// Token 缓存
+type TokenCache struct {
+	mu    sync.RWMutex
+	cache map[string]*CachedToken
+}
+
+type CachedToken struct {
+	Token     string
+	ExpiresAt time.Time
+}
+
+var tokenCache = &TokenCache{
+	cache: make(map[string]*CachedToken),
+}
+
+func (tc *TokenCache) Get(key string) (string, bool) {
+	tc.mu.RLock()
+	defer tc.mu.RUnlock()
+	
+	if cached, ok := tc.cache[key]; ok {
+		if time.Now().Before(cached.ExpiresAt) {
+			return cached.Token, true
+		}
+		delete(tc.cache, key)
+	}
+	return "", false
+}
+
+func (tc *TokenCache) Set(key, token string, ttl time.Duration) {
+	tc.mu.Lock()
+	defer tc.mu.Unlock()
+	
+	tc.cache[key] = &CachedToken{
+		Token:     token,
+		ExpiresAt: time.Now().Add(ttl),
+	}
+}
+
+// Manifest 缓存
+type ManifestCache struct {
+	mu    sync.RWMutex
+	cache map[string]*CachedManifest
+}
+
+type CachedManifest struct {
+	Data      []byte
+	Headers   http.Header
+	ExpiresAt time.Time
+}
+
+var manifestCache = &ManifestCache{
+	cache: make(map[string]*CachedManifest),
+}
+
+func (mc *ManifestCache) Get(key string) (*CachedManifest, bool) {
+	mc.mu.RLock()
+	defer mc.mu.RUnlock()
+	
+	if cached, ok := mc.cache[key]; ok {
+		if time.Now().Before(cached.ExpiresAt) {
+			return cached, true
+		}
+		delete(mc.cache, key)
+	}
+	return nil, false
+}
+
+func (mc *ManifestCache) Set(key string, data []byte, headers http.Header, ttl time.Duration) {
+	mc.mu.Lock()
+	defer mc.mu.Unlock()
+	
+	mc.cache[key] = &CachedManifest{
+		Data:      data,
+		Headers:   headers,
+		ExpiresAt: time.Now().Add(ttl),
+	}
+}
+
 type CustomFormatter struct {
 	logrus.TextFormatter
 }
 
 func (f *CustomFormatter) Format(entry *logrus.Entry) ([]byte, error) {
 	timestamp := entry.Time.Format("2006-01-02 15:04:05.000")
-
 	var levelColor string
 	switch entry.Level {
 	case logrus.DebugLevel:
@@ -66,16 +143,9 @@ func (f *CustomFormatter) Format(entry *logrus.Entry) ([]byte, error) {
 	case logrus.FatalLevel, logrus.PanicLevel:
 		levelColor = "\033[35m"
 	}
-
 	resetColor := "\033[0m"
-	logMessage := fmt.Sprintf("%s %s[%s]%s %s\n",
+	return []byte(fmt.Sprintf("%s %s[%s]%s %s\n",
-		timestamp,
+		timestamp, levelColor, strings.ToUpper(entry.Level.String()), resetColor, entry.Message)), nil
-		levelColor,
-		strings.ToUpper(entry.Level.String()),
-		resetColor,
-		entry.Message)
-
-	return []byte(logMessage), nil
 }
 
 func init() {
@@ -90,15 +160,10 @@ func init() {
 
 func preprocessArgs() {
 	alias := map[string]string{
-		"--listen":    "-l",
+		"--listen": "-l", "--port": "-p", "--log-level": "-ll", "--disguise": "-w",
-		"--port":      "-p",
-		"--log-level": "-ll",
-		"--disguise":  "-w",
 	}
-
 	newArgs := make([]string, 0, len(os.Args))
 	newArgs = append(newArgs, os.Args[0])
-
 	for _, arg := range os.Args[1:] {
 		if strings.HasPrefix(arg, "--") && strings.Contains(arg, "=") {
 			parts := strings.SplitN(arg, "=", 2)
@@ -110,14 +175,13 @@ func preprocessArgs() {
 		}
 		newArgs = append(newArgs, arg)
 	}
-
 	if len(newArgs) > 1 {
 		os.Args = newArgs
 	}
 }
 
 func usage() {
-	const helpText = `HubP - Docker Hub 代理服务器
+	fmt.Fprintf(os.Stderr, `HubP - Docker Hub 代理服务器
 
 参数说明:
     -l, --listen       监听地址 (默认: 0.0.0.0)
@@ -127,9 +191,7 @@ func usage() {
 
 示例:
     ./HubP -l 0.0.0.0 -p 18184 -ll debug -w www.bing.com
-    ./HubP --listen=0.0.0.0 --port=18184 --log-level=debug --disguise=www.bing.com`
+`)
-
-	fmt.Fprintf(os.Stderr, "%s\n", helpText)
 }
 
 func validateConfig() error {
@@ -146,15 +208,10 @@ func main() {
 	preprocessArgs()
 	flag.Usage = usage
 
-	defaultListenAddress := getEnv("HUBP_LISTEN", "0.0.0.0")
+	flag.StringVar(&config.ListenAddress, "l", getEnv("HUBP_LISTEN", "0.0.0.0"), "监听地址")
-	defaultPort := getEnvAsInt("HUBP_PORT", 18184)
+	flag.IntVar(&config.Port, "p", getEnvAsInt("HUBP_PORT", 18184), "监听端口")
-	defaultLogLevel := getEnv("HUBP_LOG_LEVEL", "debug")
+	flag.StringVar(&config.LogLevel, "ll", getEnv("HUBP_LOG_LEVEL", "debug"), "日志级别")
-	defaultDisguiseURL := getEnv("HUBP_DISGUISE", "onlinealarmkur.com")
+	flag.StringVar(&config.DisguiseURL, "w", getEnv("HUBP_DISGUISE", "onlinealarmkur.com"), "伪装网站 URL")
-
-	flag.StringVar(&config.ListenAddress, "l", defaultListenAddress, "监听地址")
-	flag.IntVar(&config.Port, "p", defaultPort, "监听端口")
-	flag.StringVar(&config.LogLevel, "ll", defaultLogLevel, "日志级别")
-	flag.StringVar(&config.DisguiseURL, "w", defaultDisguiseURL, "伪装网站 URL")
 
 	if err := flag.CommandLine.Parse(os.Args[1:]); err != nil {
 		logrus.Fatal("解析命令行参数失败：", err)
@@ -176,10 +233,7 @@ func main() {
 	addr := fmt.Sprintf("%s:%d", config.ListenAddress, config.Port)
 	http.HandleFunc("/", handleRequest)
 
-	server := &http.Server{
+	server := &http.Server{Addr: addr, Handler: http.DefaultServeMux}
-		Addr:    addr,
-		Handler: http.DefaultServeMux,
-	}
 
 	go func() {
 		logrus.Info("服务启动成功")
@@ -203,10 +257,7 @@ func main() {
 }
 
 func printStartupInfo() {
-	const blue = "\033[34m"
+	const blue, green, reset = "\033[34m", "\033[32m", "\033[0m"
-	const green = "\033[32m"
-	const reset = "\033[0m"
-
 	fmt.Println(blue + "\n╔════════════════════════════════════════════════════════════╗" + reset)
 	fmt.Println(blue + "║" + green + "               HubP Docker Hub 代理服务器               " + blue + "║" + reset)
 	fmt.Printf(blue+"║"+green+"               版本: %-33s"+blue+"║\n"+reset, Version)
@@ -222,12 +273,9 @@ func printStartupInfo() {
 func handleRequest(w http.ResponseWriter, r *http.Request) {
 	path := r.URL.Path
 
-	// 健康检查
 	if path == "/health" || path == "/healthz" {
 		w.WriteHeader(http.StatusOK)
-		if _, err := w.Write([]byte("OK")); err != nil {
+		w.Write([]byte("OK"))
-			logrus.Errorf("健康检查响应失败: %v", err)
-		}
 		return
 	}
 
@@ -242,9 +290,7 @@ func handleRequest(w http.ResponseWriter, r *http.Request) {
 		} else {
 			routeTag = "[伪装]"
 		}
-
+		logrus.Debugf("%s 请求: [%s %s] 来自 %s", routeTag, r.Method, r.URL.String(), r.RemoteAddr)
-		logrus.Debugf("%s 请求: [%s %s] 来自 %s",
-			routeTag, r.Method, r.URL.String(), r.RemoteAddr)
 	}
 
 	if strings.HasPrefix(path, "/v2/") {
@@ -268,12 +314,27 @@ func getCleanHost(r *http.Request) string {
 
 func handleRegistryRequest(w http.ResponseWriter, r *http.Request) {
 	const targetHost = "registry-1.docker.io"
+	path := strings.TrimPrefix(r.URL.Path, "/v2/")
+
+	// Manifest 缓存检查
+	if r.Method == http.MethodGet && strings.Contains(path, "/manifests/") {
+		cacheKey := fmt.Sprintf("manifest:%s", path)
+		if cached, ok := manifestCache.Get(cacheKey); ok {
+			logrus.Debugf("Docker镜像: 使用缓存的 manifest")
+			for k, v := range cached.Headers {
+				for _, val := range v {
+					w.Header().Add(k, val)
+				}
+			}
+			w.WriteHeader(http.StatusOK)
+			w.Write(cached.Data)
+			return
+		}
+	}
 
 	ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
 	defer cancel()
 
-	path := strings.TrimPrefix(r.URL.Path, "/v2/")
-
 	url := &url.URL{
 		Scheme:   "https",
 		Host:     targetHost,
@@ -289,11 +350,7 @@ func handleRegistryRequest(w http.ResponseWriter, r *http.Request) {
 	resp, err := sendRequestWithContext(ctx, r.Method, url.String(), headers, r.Body)
 	if err != nil {
 		logrus.Errorf("Docker镜像: 请求失败 - %v", err)
-		if logrus.IsLevelEnabled(logrus.DebugLevel) {
+		http.Error(w, "服务暂时不可用", http.StatusBadGateway)
-			http.Error(w, fmt.Sprintf("代理错误: %v", err), http.StatusBadGateway)
-		} else {
-			http.Error(w, "服务暂时不可用", http.StatusBadGateway)
-		}
 		return
 	}
 	defer resp.Body.Close()
@@ -304,7 +361,6 @@ func handleRegistryRequest(w http.ResponseWriter, r *http.Request) {
 	}
 
 	respHeaders := copyHeaders(resp.Header)
-
 	if respHeaders.Get("WWW-Authenticate") != "" {
 		currentDomain := getCleanHost(r)
 		respHeaders.Set("WWW-Authenticate",
@@ -318,6 +374,22 @@ func handleRegistryRequest(w http.ResponseWriter, r *http.Request) {
 	}
 	w.WriteHeader(resp.StatusCode)
 
+	// 缓存 manifest
+	if resp.StatusCode == http.StatusOK && r.Method == http.MethodGet && strings.Contains(path, "/manifests/") {
+		data, err := io.ReadAll(resp.Body)
+		if err == nil {
+			cacheKey := fmt.Sprintf("manifest:%s", path)
+			ttl := 10 * time.Minute
+			if strings.Contains(path, "sha256:") {
+				ttl = 1 * time.Hour
+			}
+			manifestCache.Set(cacheKey, data, respHeaders, ttl)
+			w.Write(data)
+			logrus.Debugf("Docker镜像: manifest 已缓存 [大小: %.2f KB]", float64(len(data))/1024)
+			return
+		}
+	}
+
 	written, err := io.Copy(w, resp.Body)
 	if err != nil {
 		logrus.Errorf("Docker镜像: 传输响应失败 - %v", err)
@@ -333,11 +405,20 @@ func handleRegistryRequest(w http.ResponseWriter, r *http.Request) {
 func handleAuthRequest(w http.ResponseWriter, r *http.Request) {
 	const targetHost = "auth.docker.io"
 
+	// Token 缓存检查
+	cacheKey := r.URL.RawQuery
+	if token, ok := tokenCache.Get(cacheKey); ok {
+		logrus.Debugf("认证服务: 使用缓存的 token")
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte(token))
+		return
+	}
+
 	ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
 	defer cancel()
 
 	path := strings.TrimPrefix(r.URL.Path, "/auth/")
-
 	url := &url.URL{
 		Scheme:   "https",
 		Host:     targetHost,
@@ -353,11 +434,7 @@ func handleAuthRequest(w http.ResponseWriter, r *http.Request) {
 	resp, err := sendRequestWithContext(ctx, r.Method, url.String(), headers, r.Body)
 	if err != nil {
 		logrus.Errorf("认证服务: 请求失败 - %v", err)
-		if logrus.IsLevelEnabled(logrus.DebugLevel) {
+		http.Error(w, "服务暂时不可用", http.StatusBadGateway)
-			http.Error(w, fmt.Sprintf("代理错误: %v", err), http.StatusBadGateway)
-		} else {
-			http.Error(w, "服务暂时不可用", http.StatusBadGateway)
-		}
 		return
 	}
 	defer resp.Body.Close()
@@ -369,6 +446,24 @@ func handleAuthRequest(w http.ResponseWriter, r *http.Request) {
 	}
 	w.WriteHeader(resp.StatusCode)
 
+	// 缓存 token
+	if resp.StatusCode == http.StatusOK {
+		data, err := io.ReadAll(resp.Body)
+		if err == nil {
+			var tokenResp map[string]interface{}
+			if json.Unmarshal(data, &tokenResp) == nil {
+				ttl := 5 * time.Minute
+				if expiresIn, ok := tokenResp["expires_in"].(float64); ok {
+					ttl = time.Duration(expiresIn*0.9) * time.Second
+				}
+				tokenCache.Set(cacheKey, string(data), ttl)
+				logrus.Debugf("认证服务: token 已缓存 [TTL: %v]", ttl)
+			}
+			w.Write(data)
+			return
+		}
+	}
+
 	written, err := io.Copy(w, resp.Body)
 	if err != nil {
 		logrus.Errorf("认证服务: 传输响应失败 - %v", err)
@@ -383,12 +478,10 @@ func handleAuthRequest(w http.ResponseWriter, r *http.Request) {
 
 func handleCloudflareRequest(w http.ResponseWriter, r *http.Request) {
 	const targetHost = "production.cloudflare.docker.com"
-
 	ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
 	defer cancel()
 
 	path := strings.TrimPrefix(r.URL.Path, "/production-cloudflare/")
-
 	url := &url.URL{
 		Scheme:   "https",
 		Host:     targetHost,
@@ -404,11 +497,7 @@ func handleCloudflareRequest(w http.ResponseWriter, r *http.Request) {
 	resp, err := sendRequestWithContext(ctx, r.Method, url.String(), headers, r.Body)
 	if err != nil {
 		logrus.Errorf("Cloudflare: 请求失败 - %v", err)
-		if logrus.IsLevelEnabled(logrus.DebugLevel) {
+		http.Error(w, "服务暂时不可用", http.StatusBadGateway)
-			http.Error(w, fmt.Sprintf("代理错误: %v", err), http.StatusBadGateway)
-		} else {
-			http.Error(w, "服务暂时不可用", http.StatusBadGateway)
-		}
 		return
 	}
 	defer resp.Body.Close()
@@ -446,11 +535,7 @@ func handleAuthChallenge(w http.ResponseWriter, r *http.Request, resp *http.Resp
 	}
 
 	w.WriteHeader(resp.StatusCode)
-
+	io.Copy(w, resp.Body)
-	_, err := io.Copy(w, resp.Body)
-	if err != nil {
-		logrus.Errorf("认证响应传输失败: %v", err)
-	}
 }
 
 func handleDisguise(w http.ResponseWriter, r *http.Request) {
@@ -505,7 +590,7 @@ func sendRequestWithContext(ctx context.Context, method, url string, headers htt
 	}
 
 	req.Header = headers
-	
+
 	startTime := time.Now()
 	resp, err := client.Do(req)