From 587c1f2144872997d69cabb14c782faea8989510 Mon Sep 17 00:00:00 2001 From: sky22333 Date: Sat, 11 Jul 2026 20:42:23 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=20Registry=20token=20?= =?UTF-8?q?=E8=B7=AF=E7=94=B1=E4=B8=8E=E7=A6=BB=E7=BA=BF=E4=B8=8B=E8=BD=BD?= =?UTF-8?q?=E9=94=99=E8=AF=AF=E5=93=8D=E5=BA=94=EF=BC=9B=E5=8E=BB=E6=8E=89?= =?UTF-8?q?=E5=86=97=E4=BD=99=E6=AD=A3=E5=88=99?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/config/config.go | 12 ---- src/handlers/docker.go | 67 ++++++++++++------ src/handlers/docker_test.go | 127 +++++++++++++++++++++++++++++++++- src/handlers/github.go | 1 - src/handlers/github_test.go | 3 + src/handlers/imagetar.go | 55 ++++++++++----- src/handlers/imagetar_test.go | 35 ++++++++++ src/public/index.html | 3 +- 8 files changed, 248 insertions(+), 55 deletions(-) diff --git a/src/config/config.go b/src/config/config.go index 2c1e75e..2227fee 100644 --- a/src/config/config.go +++ b/src/config/config.go @@ -276,15 +276,3 @@ func overrideFromEnv(cfg *AppConfig) { } } } - -// CreateDefaultConfigFile 创建默认配置文件 -func CreateDefaultConfigFile() error { - cfg := DefaultConfig() - - data, err := toml.Marshal(cfg) - if err != nil { - return fmt.Errorf("序列化默认配置失败: %v", err) - } - - return os.WriteFile("config.toml", data, 0644) -} diff --git a/src/handlers/docker.go b/src/handlers/docker.go index 39dc554..6dea065 100644 --- a/src/handlers/docker.go +++ b/src/handlers/docker.go @@ -108,9 +108,6 @@ func handleRegistryRequest(c *gin.Context, path string) { if registryDomain, remainingPath := registryDetector.detectRegistryDomain(c, pathWithoutV2); registryDomain != "" { if registryDetector.isRegistryEnabled(registryDomain) { - c.Set("target_registry_domain", registryDomain) - c.Set("target_path", remainingPath) - handleMultiRegistryRequest(c, registryDomain, remainingPath) return } @@ -355,20 +352,7 @@ func (r *ResponseRecorder) Write(data []byte) (int, error) { } func proxyDockerAuthOriginal(c *gin.Context) { - var authURL string - if targetDomain, exists := c.Get("target_registry_domain"); exists { - if mapping, found := registryDetector.getRegistryMapping(targetDomain.(string)); found { - authURL = "https://" + mapping.AuthHost + c.Request.URL.Path - } else { - authURL = "https://auth.docker.io" + c.Request.URL.Path - } - } else { - authURL = "https://auth.docker.io" + c.Request.URL.Path - } - - if c.Request.URL.RawQuery != "" { - authURL += "?" + c.Request.URL.RawQuery - } + authURL := buildDockerAuthURL(c) client := &http.Client{ Timeout: 30 * time.Second, @@ -423,12 +407,53 @@ func proxyDockerAuthOriginal(c *gin.Context) { } } -// rewriteAuthHeader 重写认证头 +// buildDockerAuthURL 根据 token 请求的 service 参数选择上游认证地址。 +// AuthHost 已包含路径（如 ghcr.io/token、quay.io/v2/auth），不再拼接本机 Path，避免 /token/token。 +func buildDockerAuthURL(c *gin.Context) string { + authHost := resolveAuthHost(c.Query("service")) + var authURL string + if authHost != "" { + authURL = "https://" + authHost + } else { + authURL = "https://auth.docker.io" + c.Request.URL.Path + } + if c.Request.URL.RawQuery != "" { + authURL += "?" + c.Request.URL.RawQuery + } + return authURL +} + +// resolveAuthHost 用 service 匹配已启用 Registry 的 AuthHost；Docker Hub 返回空串走默认路径。 +func resolveAuthHost(service string) string { + if service == "" || service == "registry.docker.io" || service == "docker.io" { + return "" + } + + cfg := config.GetConfig() + for domain, mapping := range cfg.Registries { + if !mapping.Enabled || mapping.AuthHost == "" { + continue + } + if service == domain || service == mapping.Upstream { + return mapping.AuthHost + } + } + return "" +} + +// rewriteAuthHeader 将上游认证 realm 统一改写到本机 /token，避免 quay 等变成 /v2/auth 误入 Registry 路由。 func rewriteAuthHeader(authHeader, proxyHost string) string { + proxyToken := "http://" + proxyHost + "/token" + + cfg := config.GetConfig() + for _, mapping := range cfg.Registries { + if mapping.AuthHost == "" { + continue + } + authHeader = strings.ReplaceAll(authHeader, "https://"+mapping.AuthHost, proxyToken) + } + authHeader = strings.ReplaceAll(authHeader, "https://auth.docker.io/token", proxyToken) authHeader = strings.ReplaceAll(authHeader, "https://auth.docker.io", "http://"+proxyHost) - authHeader = strings.ReplaceAll(authHeader, "https://ghcr.io", "http://"+proxyHost) - authHeader = strings.ReplaceAll(authHeader, "https://gcr.io", "http://"+proxyHost) - authHeader = strings.ReplaceAll(authHeader, "https://quay.io", "http://"+proxyHost) return authHeader } diff --git a/src/handlers/docker_test.go b/src/handlers/docker_test.go index e4c2f9d..524068b 100644 --- a/src/handlers/docker_test.go +++ b/src/handlers/docker_test.go @@ -1,6 +1,15 @@ package handlers -import "testing" +import ( + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "testing" + + "github.com/gin-gonic/gin" + "hubproxy/config" +) func TestParseRegistryPath(t *testing.T) { tests := []struct { @@ -28,3 +37,119 @@ func TestParseRegistryPathInvalid(t *testing.T) { t.Fatalf("invalid path parsed as %q %q %q", image, apiType, reference) } } + +func TestResolveAuthHost(t *testing.T) { + path := filepath.Join(t.TempDir(), "config.toml") + data := []byte(` +[registries."ghcr.io"] +upstream = "ghcr.io" +authHost = "ghcr.io/token" +authType = "github" +enabled = true + +[registries."quay.io"] +upstream = "quay.io" +authHost = "quay.io/v2/auth" +authType = "quay" +enabled = true +`) + if err := os.WriteFile(path, data, 0644); err != nil { + t.Fatal(err) + } + t.Setenv("CONFIG_PATH", path) + if err := config.LoadConfig(); err != nil { + t.Fatal(err) + } + + if got := resolveAuthHost(""); got != "" { + t.Fatalf("empty service = %q", got) + } + if got := resolveAuthHost("registry.docker.io"); got != "" { + t.Fatalf("docker hub service = %q", got) + } + if got := resolveAuthHost("ghcr.io"); got != "ghcr.io/token" { + t.Fatalf("ghcr.io = %q", got) + } + if got := resolveAuthHost("quay.io"); got != "quay.io/v2/auth" { + t.Fatalf("quay.io = %q", got) + } +} + +func TestBuildDockerAuthURL(t *testing.T) { + path := filepath.Join(t.TempDir(), "config.toml") + data := []byte(` +[registries."ghcr.io"] +upstream = "ghcr.io" +authHost = "ghcr.io/token" +enabled = true +`) + if err := os.WriteFile(path, data, 0644); err != nil { + t.Fatal(err) + } + t.Setenv("CONFIG_PATH", path) + if err := config.LoadConfig(); err != nil { + t.Fatal(err) + } + + gin.SetMode(gin.TestMode) + + t.Run("docker hub keeps path", func(t *testing.T) { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodGet, "/token?service=registry.docker.io&scope=repository:library/nginx:pull", nil) + got := buildDockerAuthURL(c) + want := "https://auth.docker.io/token?service=registry.docker.io&scope=repository:library/nginx:pull" + if got != want { + t.Fatalf("got %q want %q", got, want) + } + }) + + t.Run("ghcr uses AuthHost without duplicating path", func(t *testing.T) { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodGet, "/token?service=ghcr.io&scope=repository:foo/bar:pull", nil) + got := buildDockerAuthURL(c) + want := "https://ghcr.io/token?service=ghcr.io&scope=repository:foo/bar:pull" + if got != want { + t.Fatalf("got %q want %q", got, want) + } + }) +} + +func TestRewriteAuthHeader(t *testing.T) { + path := filepath.Join(t.TempDir(), "config.toml") + data := []byte(` +[registries."quay.io"] +upstream = "quay.io" +authHost = "quay.io/v2/auth" +enabled = true + +[registries."ghcr.io"] +upstream = "ghcr.io" +authHost = "ghcr.io/token" +enabled = true +`) + if err := os.WriteFile(path, data, 0644); err != nil { + t.Fatal(err) + } + t.Setenv("CONFIG_PATH", path) + if err := config.LoadConfig(); err != nil { + t.Fatal(err) + } + + got := rewriteAuthHeader(`Bearer realm="https://quay.io/v2/auth",service="quay.io"`, "proxy.example.com") + want := `Bearer realm="http://proxy.example.com/token",service="quay.io"` + if got != want { + t.Fatalf("quay rewrite: got %q want %q", got, want) + } + + got = rewriteAuthHeader(`Bearer realm="https://ghcr.io/token",service="ghcr.io"`, "proxy.example.com") + want = `Bearer realm="http://proxy.example.com/token",service="ghcr.io"` + if got != want { + t.Fatalf("ghcr rewrite: got %q want %q", got, want) + } + + got = rewriteAuthHeader(`Bearer realm="https://auth.docker.io/token",service="registry.docker.io"`, "proxy.example.com") + want = `Bearer realm="http://proxy.example.com/token",service="registry.docker.io"` + if got != want { + t.Fatalf("docker hub rewrite: got %q want %q", got, want) + } +} diff --git a/src/handlers/github.go b/src/handlers/github.go index 2c45868..5a8ad3c 100644 --- a/src/handlers/github.go +++ b/src/handlers/github.go @@ -24,7 +24,6 @@ var ( regexp.MustCompile(`^(?:https?://)?api\.github\.com/repos/([^/]+)/([^/]+)/.*`), regexp.MustCompile(`^(?:https?://)?huggingface\.co(?:/spaces)?/([^/]+)/(.+)`), regexp.MustCompile(`^(?:https?://)?cdn-lfs\.hf\.co(?:/spaces)?/([^/]+)/([^/]+)(?:/(.*))?`), - regexp.MustCompile(`^(?:https?://)?download\.docker\.com/([^/]+)/.*\.(tgz|zip)`), regexp.MustCompile(`^(?:https?://)?(github|opengraph)\.githubassets\.com/([^/]+)/.+?`), } ) diff --git a/src/handlers/github_test.go b/src/handlers/github_test.go index d0acf35..3d1725a 100644 --- a/src/handlers/github_test.go +++ b/src/handlers/github_test.go @@ -29,4 +29,7 @@ func TestCheckGitHubURLRejectsOtherHosts(t *testing.T) { if got := CheckGitHubURL("https://example.com/user/repo/file"); got != nil { t.Fatalf("unexpected match: %#v", got) } + if got := CheckGitHubURL("https://download.docker.com/linux/static/stable/x86_64/docker.tgz"); got != nil { + t.Fatalf("download.docker.com should be rejected: %#v", got) + } } diff --git a/src/handlers/imagetar.go b/src/handlers/imagetar.go index bcb841e..a618a05 100644 --- a/src/handlers/imagetar.go +++ b/src/handlers/imagetar.go @@ -289,27 +289,22 @@ func (is *ImageStreamer) StreamImageToWriter(ctx context.Context, imageRef strin contextOptions := append(is.remoteOptions, remote.WithContext(ctx)) - desc, err := is.getImageDescriptorWithPlatform(ref, contextOptions, options.Platform) + desc, err := is.getImageDescriptor(ref, contextOptions) if err != nil { return fmt.Errorf("获取镜像描述失败: %w", err) } switch desc.MediaType { case types.OCIImageIndex, types.DockerManifestList: - return is.streamMultiArchImage(ctx, desc, writer, options, contextOptions, imageRef) + return is.streamMultiArchImage(ctx, desc, writer, options, imageRef) case types.OCIManifestSchema1, types.DockerManifestSchema2: - return is.streamSingleImage(ctx, desc, writer, options, contextOptions, imageRef) + return is.streamSingleImage(ctx, desc, writer, options, imageRef) default: - return is.streamSingleImage(ctx, desc, writer, options, contextOptions, imageRef) + return is.streamSingleImage(ctx, desc, writer, options, imageRef) } } // getImageDescriptor 获取镜像描述符 func (is *ImageStreamer) getImageDescriptor(ref name.Reference, options []remote.Option) (*remote.Descriptor, error) { - return is.getImageDescriptorWithPlatform(ref, options, "") -} - -// getImageDescriptorWithPlatform 获取指定平台的镜像描述符 -func (is *ImageStreamer) getImageDescriptorWithPlatform(ref name.Reference, options []remote.Option, platform string) (*remote.Descriptor, error) { return remote.Get(ref, options...) } @@ -324,20 +319,47 @@ func setDownloadHeaders(c *gin.Context, filename string, compressed bool) { } } +// writeDownloadError 仅在尚未写出响应体时返回 JSON；流已开始则只记日志，避免损坏 tar。 +func writeDownloadError(c *gin.Context, err error, message string) { + if c.Writer.Written() { + log.Printf("%s: %v", message, err) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": message + ": " + err.Error()}) +} + // StreamImageToGin 流式响应到Gin func (is *ImageStreamer) StreamImageToGin(ctx context.Context, imageRef string, c *gin.Context, options *StreamOptions) error { if options == nil { options = &StreamOptions{UseCompressedLayers: true} } + ref, err := name.ParseReference(imageRef) + if err != nil { + return fmt.Errorf("解析镜像引用失败: %w", err) + } + + contextOptions := append(is.remoteOptions, remote.WithContext(ctx)) + desc, err := is.getImageDescriptor(ref, contextOptions) + if err != nil { + return fmt.Errorf("获取镜像描述失败: %w", err) + } + filename := strings.ReplaceAll(imageRef, "/", "_") + ".tar" setDownloadHeaders(c, filename, options.Compression) - return is.StreamImageToWriter(ctx, imageRef, c.Writer, options) + switch desc.MediaType { + case types.OCIImageIndex, types.DockerManifestList: + return is.streamMultiArchImage(ctx, desc, c.Writer, options, imageRef) + case types.OCIManifestSchema1, types.DockerManifestSchema2: + return is.streamSingleImage(ctx, desc, c.Writer, options, imageRef) + default: + return is.streamSingleImage(ctx, desc, c.Writer, options, imageRef) + } } // streamMultiArchImage 处理多架构镜像 -func (is *ImageStreamer) streamMultiArchImage(ctx context.Context, desc *remote.Descriptor, writer io.Writer, options *StreamOptions, remoteOptions []remote.Option, imageRef string) error { +func (is *ImageStreamer) streamMultiArchImage(ctx context.Context, desc *remote.Descriptor, writer io.Writer, options *StreamOptions, imageRef string) error { img, err := is.selectPlatformImage(desc, options) if err != nil { return err @@ -347,7 +369,7 @@ func (is *ImageStreamer) streamMultiArchImage(ctx context.Context, desc *remote. } // streamSingleImage 处理单架构镜像 -func (is *ImageStreamer) streamSingleImage(ctx context.Context, desc *remote.Descriptor, writer io.Writer, options *StreamOptions, remoteOptions []remote.Option, imageRef string) error { +func (is *ImageStreamer) streamSingleImage(ctx context.Context, desc *remote.Descriptor, writer io.Writer, options *StreamOptions, imageRef string) error { img, err := desc.Image() if err != nil { return fmt.Errorf("获取镜像失败: %w", err) @@ -583,7 +605,7 @@ func (is *ImageStreamer) streamSingleImageForBatch(ctx context.Context, tarWrite contextOptions := append(is.remoteOptions, remote.WithContext(ctx)) - desc, err := is.getImageDescriptorWithPlatform(ref, contextOptions, options.Platform) + desc, err := is.getImageDescriptor(ref, contextOptions) if err != nil { return nil, nil, fmt.Errorf("获取镜像描述失败: %w", err) } @@ -784,8 +806,7 @@ func handleDirectImageDownload(c *gin.Context) { log.Printf("下载镜像: %s (平台: %s)", req.Image, formatPlatformText(req.Platform)) if err := globalImageStreamer.StreamImageToGin(ctx, req.Image, c, options); err != nil { - log.Printf("镜像下载失败: %v", err) - c.JSON(http.StatusInternalServerError, gin.H{"error": "镜像下载失败: " + err.Error()}) + writeDownloadError(c, err, "镜像下载失败") return } } @@ -821,12 +842,10 @@ func handleSimpleBatchDownload(c *gin.Context) { log.Printf("批量下载 %d 个镜像 (平台: %s)", len(req.Images), formatPlatformText(req.Platform)) filename := fmt.Sprintf("batch_%d_images.tar", len(req.Images)) - setDownloadHeaders(c, filename, options.Compression) if err := globalImageStreamer.StreamMultipleImages(ctx, req.Images, c.Writer, options); err != nil { - log.Printf("批量镜像下载失败: %v", err) - c.JSON(http.StatusInternalServerError, gin.H{"error": "批量镜像下载失败: " + err.Error()}) + writeDownloadError(c, err, "批量镜像下载失败") return } return diff --git a/src/handlers/imagetar_test.go b/src/handlers/imagetar_test.go index 1d52199..8aa0e5e 100644 --- a/src/handlers/imagetar_test.go +++ b/src/handlers/imagetar_test.go @@ -1,8 +1,14 @@ package handlers import ( + "errors" + "net/http" + "net/http/httptest" + "strings" "testing" "time" + + "github.com/gin-gonic/gin" ) func TestDownloadDebouncer(t *testing.T) { @@ -58,3 +64,32 @@ func TestGenerateContentFingerprintStable(t *testing.T) { t.Fatalf("unexpected fingerprints: %q %q %q", a, b, c) } } + +func TestWriteDownloadErrorSkipsJSONAfterBodyStarted(t *testing.T) { + gin.SetMode(gin.TestMode) + + t.Run("before write returns json", func(t *testing.T) { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + writeDownloadError(c, errors.New("boom"), "镜像下载失败") + if w.Code != http.StatusInternalServerError { + t.Fatalf("status = %d", w.Code) + } + body := w.Body.String() + if !strings.Contains(body, "镜像下载失败") || !strings.Contains(body, "boom") { + t.Fatalf("body = %q", body) + } + }) + + t.Run("after write skips json", func(t *testing.T) { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + if _, err := c.Writer.Write([]byte("tar-bytes")); err != nil { + t.Fatal(err) + } + writeDownloadError(c, errors.New("boom"), "镜像下载失败") + if got := w.Body.String(); got != "tar-bytes" { + t.Fatalf("body corrupted: %q", got) + } + }) +} diff --git a/src/public/index.html b/src/public/index.html index 36113d8..90aac41 100644 --- a/src/public/index.html +++ b/src/public/index.html @@ -752,8 +752,7 @@ link.startsWith("raw.githubusercontent.com/") || link.startsWith("gist.githubusercontent.com/") || link.startsWith("huggingface.co/") || - link.startsWith("cdn-lfs.hf.co/") || - link.startsWith("download.docker.com/") + link.startsWith("cdn-lfs.hf.co/") ) { formattedLink = "https://" + currentHost + "/https://" + link; } else {