From d431bc841957980007abf542ec231ae5879b048c Mon Sep 17 00:00:00 2001 From: XOF Date: Sun, 28 Dec 2025 15:25:41 +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 | 235 ++++++++++++++++++++++++++++++++++++++------------------ 1 file changed, 160 insertions(+), 75 deletions(-) diff --git a/main.go b/main.go index 36f3ecc..cbf712a 100644 --- a/main.go +++ b/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", - timestamp, - levelColor, - strings.ToUpper(entry.Level.String()), - resetColor, - entry.Message) - - return []byte(logMessage), nil + return []byte(fmt.Sprintf("%s %s[%s]%s %s\n", + timestamp, levelColor, strings.ToUpper(entry.Level.String()), resetColor, entry.Message)), nil } func init() { @@ -90,15 +160,10 @@ func init() { func preprocessArgs() { alias := map[string]string{ - "--listen": "-l", - "--port": "-p", - "--log-level": "-ll", - "--disguise": "-w", + "--listen": "-l", "--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") - defaultPort := getEnvAsInt("HUBP_PORT", 18184) - defaultLogLevel := getEnv("HUBP_LOG_LEVEL", "debug") - defaultDisguiseURL := getEnv("HUBP_DISGUISE", "onlinealarmkur.com") - - 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") + flag.StringVar(&config.ListenAddress, "l", getEnv("HUBP_LISTEN", "0.0.0.0"), "监听地址") + flag.IntVar(&config.Port, "p", getEnvAsInt("HUBP_PORT", 18184), "监听端口") + flag.StringVar(&config.LogLevel, "ll", getEnv("HUBP_LOG_LEVEL", "debug"), "日志级别") + flag.StringVar(&config.DisguiseURL, "w", getEnv("HUBP_DISGUISE", "onlinealarmkur.com"), "伪装网站 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{ - Addr: addr, - Handler: http.DefaultServeMux, - } + server := &http.Server{Addr: addr, Handler: http.DefaultServeMux} go func() { logrus.Info("服务启动成功") @@ -203,10 +257,7 @@ func main() { } func printStartupInfo() { - const blue = "\033[34m" - const green = "\033[32m" - const reset = "\033[0m" - + const blue, green, reset = "\033[34m", "\033[32m", "\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 { - logrus.Errorf("健康检查响应失败: %v", err) - } + w.Write([]byte("OK")) 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, fmt.Sprintf("代理错误: %v", err), http.StatusBadGateway) - } else { - http.Error(w, "服务暂时不可用", http.StatusBadGateway) - } + 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, fmt.Sprintf("代理错误: %v", err), http.StatusBadGateway) - } else { - http.Error(w, "服务暂时不可用", http.StatusBadGateway) - } + 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, fmt.Sprintf("代理错误: %v", err), http.StatusBadGateway) - } else { - http.Error(w, "服务暂时不可用", http.StatusBadGateway) - } + 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) - - _, err := io.Copy(w, resp.Body) - if err != nil { - logrus.Errorf("认证响应传输失败: %v", err) - } + io.Copy(w, resp.Body) } 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)