重构Docker代理，支持认证透传和可配置Docker Hub上游，并补充边界测试

2026-08-10 14:12:05 +08:00 · 2026-05-16 02:58:31 +08:00
parent ba83a44492
commit 53cc1761ce
2 changed files with 937 additions and 517 deletions
@@ -1,6 +1,56 @@
 package handlers

-import "testing"
+import (
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"path/filepath"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/gin-gonic/gin"
+	"hubproxy/config"
+	"hubproxy/utils"
+)
+
+type zeroReader struct{}
+
+func (zeroReader) Read(p []byte) (int, error) {
+	for i := range p {
+		p[i] = 0
+	}
+	return len(p), nil
+}
+
+type discardResponseWriter struct {
+	header http.Header
+	status int
+	bytes  int64
+}
+
+func newDiscardResponseWriter() *discardResponseWriter {
+	return &discardResponseWriter{header: make(http.Header)}
+}
+
+func (w *discardResponseWriter) Header() http.Header {
+	return w.header
+}
+
+func (w *discardResponseWriter) WriteHeader(status int) {
+	w.status = status
+}
+
+func (w *discardResponseWriter) Write(p []byte) (int, error) {
+	if w.status == 0 {
+		w.status = http.StatusOK
+	}
+	w.bytes += int64(len(p))
+	return len(p), nil
+}

 func TestParseRegistryPath(t *testing.T) {
 	tests := []struct {
@@ -28,3 +78,546 @@ func TestParseRegistryPathInvalid(t *testing.T) {
 		t.Fatalf("invalid path parsed as %q %q %q", image, apiType, reference)
 	}
 }
+
+type testEnv interface {
+	Helper()
+	TempDir() string
+	Setenv(string, string)
+	Fatal(...interface{})
+}
+
+func initDockerProxyTest(t testEnv, configBody string) {
+	t.Helper()
+
+	path := filepath.Join(t.TempDir(), "config.toml")
+	if err := os.WriteFile(path, []byte(configBody), 0644); err != nil {
+		t.Fatal(err)
+	}
+
+	t.Setenv("CONFIG_PATH", path)
+	if err := config.LoadConfig(); err != nil {
+		t.Fatal(err)
+	}
+	utils.InitHTTPClients()
+}
+
+func TestRewriteAuthChallengePreservesScopeAndUsesProxyRealm(t *testing.T) {
+	target := registryTarget{
+		Name:        "ghcr.io",
+		AuthService: "ghcr.io",
+	}
+
+	got := rewriteAuthChallenge(
+		`Bearer realm="https://ghcr.io/token",service="ghcr.io",scope="repository:owner/image:pull"`,
+		target,
+		"https://proxy.example.com",
+	)
+
+	want := `Bearer realm="https://proxy.example.com/token/ghcr.io",service="ghcr.io",scope="repository:owner/image:pull"`
+	if got != want {
+		t.Fatalf("challenge = %q, want %q", got, want)
+	}
+}
+
+func TestBuildAuthURLForDockerHubAddsLibraryScopeAndService(t *testing.T) {
+	got, err := buildAuthURL(
+		defaultRegistryTarget(),
+		"service=ignored&scope=repository%3Aalpine%3Apull&client_id=docker",
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !strings.HasPrefix(got, dockerHubAuthRealm+"?") {
+		t.Fatalf("auth URL = %q", got)
+	}
+	if !strings.Contains(got, "service=registry.docker.io") {
+		t.Fatalf("auth URL missing service: %q", got)
+	}
+	if !strings.Contains(got, "scope=repository%3Alibrary%2Falpine%3Apull") {
+		t.Fatalf("auth URL missing normalized scope: %q", got)
+	}
+}
+
+func TestDockerIODefaultTargetUsesBuiltInWhenUnconfigured(t *testing.T) {
+	initDockerProxyTest(t, "")
+
+	target := defaultRegistryTarget()
+	if target.Upstream != dockerHubUpstream {
+		t.Fatalf("Upstream = %q, want %q", target.Upstream, dockerHubUpstream)
+	}
+	if target.AuthRealm != dockerHubAuthRealm {
+		t.Fatalf("AuthRealm = %q, want %q", target.AuthRealm, dockerHubAuthRealm)
+	}
+	if !target.AutoLibraryPrefix {
+		t.Fatal("AutoLibraryPrefix = false, want true")
+	}
+}
+
+func TestDockerIODefaultTargetCanBeOverriddenByConfig(t *testing.T) {
+	initDockerProxyTest(t, `
+[registries."docker.io"]
+upstream = "mirror.local"
+authHost = "auth.mirror.local/token"
+authType = "docker"
+enabled = true
+`)
+
+	target := defaultRegistryTarget()
+	if target.Upstream != "https://mirror.local" {
+		t.Fatalf("Upstream = %q, want custom mirror", target.Upstream)
+	}
+	if target.AuthRealm != "https://auth.mirror.local/token" {
+		t.Fatalf("AuthRealm = %q, want custom auth realm", target.AuthRealm)
+	}
+	if target.AuthService != dockerHubAuthService {
+		t.Fatalf("AuthService = %q, want %q", target.AuthService, dockerHubAuthService)
+	}
+	if !target.AutoLibraryPrefix {
+		t.Fatal("AutoLibraryPrefix = false, want true")
+	}
+}
+
+func TestDockerIODefaultTargetIgnoresDisabledOverride(t *testing.T) {
+	initDockerProxyTest(t, `
+[registries."docker.io"]
+upstream = "mirror.local"
+authHost = "auth.mirror.local/token"
+authType = "docker"
+enabled = false
+`)
+
+	target := defaultRegistryTarget()
+	if target.Upstream != dockerHubUpstream {
+		t.Fatalf("Upstream = %q, want built-in %q", target.Upstream, dockerHubUpstream)
+	}
+}
+
+func TestProxyDockerRegistryTransparentlyForwardsAuthAndRewritesChallenge(t *testing.T) {
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != "/v2/team/app/manifests/latest" {
+			t.Fatalf("upstream path = %q", r.URL.Path)
+		}
+		if got := r.Header.Get("Authorization"); got != "Bearer client-token" {
+			t.Fatalf("Authorization = %q", got)
+		}
+		if got := r.Header.Get("Accept"); got != "application/vnd.docker.distribution.manifest.v2+json" {
+			t.Fatalf("Accept = %q", got)
+		}
+		if got := r.Header.Get("Range"); got != "bytes=0-99" {
+			t.Fatalf("Range = %q", got)
+		}
+
+		w.Header().Set("WWW-Authenticate", `Bearer realm="https://upstream.example/token",service="upstream.example",scope="repository:team/app:pull"`)
+		w.WriteHeader(http.StatusUnauthorized)
+	}))
+	defer upstream.Close()
+
+	initDockerProxyTest(t, `
+[registries."test.local"]
+upstream = "`+upstream.URL+`"
+authHost = "https://auth.test.local/token"
+authType = "anonymous"
+enabled = true
+`)
+
+	gin.SetMode(gin.TestMode)
+	router := gin.New()
+	router.Any("/v2/*path", ProxyDockerRegistryGin)
+
+	req := httptest.NewRequest(http.MethodGet, "/v2/test.local/team/app/manifests/latest", nil)
+	req.Host = "proxy.example.com"
+	req.Header.Set("X-Forwarded-Proto", "https")
+	req.Header.Set("Authorization", "Bearer client-token")
+	req.Header.Set("Accept", "application/vnd.docker.distribution.manifest.v2+json")
+	req.Header.Set("Range", "bytes=0-99")
+	w := httptest.NewRecorder()
+
+	router.ServeHTTP(w, req)
+
+	if w.Code != http.StatusUnauthorized {
+		t.Fatalf("status = %d, want 401; body=%s", w.Code, w.Body.String())
+	}
+
+	wantChallenge := `Bearer realm="https://proxy.example.com/token/test.local",service="` + strings.TrimPrefix(upstream.URL, "http://") + `",scope="repository:team/app:pull"`
+	if got := w.Header().Get("WWW-Authenticate"); got != wantChallenge {
+		t.Fatalf("WWW-Authenticate = %q, want %q", got, wantChallenge)
+	}
+}
+
+func TestProxyDockerAuthForwardsBasicCredentials(t *testing.T) {
+	authServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if got := r.Header.Get("Authorization"); got != "Basic dXNlcjpwYXNz" {
+			t.Fatalf("Authorization = %q", got)
+		}
+		if got := r.URL.Query().Get("service"); got != "127.0.0.1" {
+			t.Fatalf("service = %q", got)
+		}
+		if got := r.URL.Query().Get("scope"); got != "repository:team/app:pull" {
+			t.Fatalf("scope = %q", got)
+		}
+
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"token":"secret","expires_in":3600}`))
+	}))
+	defer authServer.Close()
+
+	initDockerProxyTest(t, `
+[registries."test.local"]
+upstream = "https://127.0.0.1"
+authHost = "`+authServer.URL+`"
+authType = "anonymous"
+enabled = true
+`)
+
+	gin.SetMode(gin.TestMode)
+	router := gin.New()
+	router.Any("/token/*path", ProxyDockerAuthGin)
+
+	req := httptest.NewRequest(http.MethodGet, "/token/test.local?scope=repository:team/app:pull", nil)
+	req.Header.Set("Authorization", "Basic dXNlcjpwYXNz")
+	w := httptest.NewRecorder()
+
+	router.ServeHTTP(w, req)
+
+	if w.Code != http.StatusOK {
+		t.Fatalf("status = %d, want 200; body=%s", w.Code, w.Body.String())
+	}
+	if got := w.Body.String(); !strings.Contains(got, `"token":"secret"`) {
+		t.Fatalf("body = %q", got)
+	}
+}
+
+func TestDockerHubShortNameIsProxiedWithLibraryPrefix(t *testing.T) {
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != "/v2/library/nginx/manifests/latest" {
+			t.Fatalf("upstream path = %q", r.URL.Path)
+		}
+		w.Header().Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json")
+		w.Header().Set("Docker-Content-Digest", "sha256:abc")
+		_, _ = w.Write([]byte(`{"schemaVersion":2}`))
+	}))
+	defer upstream.Close()
+
+	gin.SetMode(gin.TestMode)
+	router := gin.New()
+	router.Any("/v2/*path", ProxyDockerRegistryGin)
+
+	target := defaultRegistryTarget()
+	target.Upstream = upstream.URL
+
+	req := httptest.NewRequest(http.MethodGet, "/v2/nginx/manifests/latest", nil)
+	w := httptest.NewRecorder()
+
+	c, _ := gin.CreateTestContext(w)
+	c.Request = req
+	proxyRegistryHTTP(c, target, "/v2/library/nginx/manifests/latest")
+
+	if w.Code != http.StatusOK {
+		t.Fatalf("status = %d, want 200; body=%s", w.Code, w.Body.String())
+	}
+	if got := w.Header().Get("Docker-Content-Digest"); got != "sha256:abc" {
+		t.Fatalf("Docker-Content-Digest = %q", got)
+	}
+}
+
+func TestProxyDockerRegistryHeadReturnsHeadersWithoutBody(t *testing.T) {
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodHead {
+			t.Fatalf("method = %s, want HEAD", r.Method)
+		}
+		w.Header().Set("Content-Length", "123")
+		w.Header().Set("Docker-Content-Digest", "sha256:head")
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer upstream.Close()
+
+	initDockerProxyTest(t, `
+[registries."test.local"]
+upstream = "`+upstream.URL+`"
+authHost = "https://auth.test.local/token"
+authType = "anonymous"
+enabled = true
+`)
+
+	gin.SetMode(gin.TestMode)
+	router := gin.New()
+	router.Any("/v2/*path", ProxyDockerRegistryGin)
+
+	req := httptest.NewRequest(http.MethodHead, "/v2/test.local/team/app/blobs/sha256:abc", nil)
+	w := httptest.NewRecorder()
+	router.ServeHTTP(w, req)
+
+	if w.Code != http.StatusOK {
+		t.Fatalf("status = %d, want 200", w.Code)
+	}
+	if got := w.Header().Get("Docker-Content-Digest"); got != "sha256:head" {
+		t.Fatalf("Docker-Content-Digest = %q", got)
+	}
+	if body := w.Body.String(); body != "" {
+		t.Fatalf("HEAD body = %q, want empty", body)
+	}
+}
+
+func TestProxyDockerRegistryStreamsBlobAndSkipsHopByHopHeaders(t *testing.T) {
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Connection", "close")
+		w.Header().Set("Transfer-Encoding", "chunked")
+		w.Header().Set("Content-Type", "application/octet-stream")
+		_, _ = io.WriteString(w, "layer-data")
+	}))
+	defer upstream.Close()
+
+	initDockerProxyTest(t, `
+[registries."test.local"]
+upstream = "`+upstream.URL+`"
+authHost = "https://auth.test.local/token"
+authType = "anonymous"
+enabled = true
+`)
+
+	gin.SetMode(gin.TestMode)
+	router := gin.New()
+	router.Any("/v2/*path", ProxyDockerRegistryGin)
+
+	req := httptest.NewRequest(http.MethodGet, "/v2/test.local/team/app/blobs/sha256:abc", nil)
+	w := httptest.NewRecorder()
+	router.ServeHTTP(w, req)
+
+	if w.Code != http.StatusOK {
+		t.Fatalf("status = %d, want 200; body=%s", w.Code, w.Body.String())
+	}
+	if got := w.Body.String(); got != "layer-data" {
+		t.Fatalf("body = %q", got)
+	}
+	if got := w.Header().Get("Connection"); got != "" {
+		t.Fatalf("Connection header leaked: %q", got)
+	}
+}
+
+func TestProxyDockerRegistryUsesNsQueryForContainerd(t *testing.T) {
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != "/v2/team/app/manifests/latest" {
+			t.Fatalf("upstream path = %q", r.URL.Path)
+		}
+		if got := r.URL.Query().Get("ns"); got != "test.local" {
+			t.Fatalf("ns query = %q", got)
+		}
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer upstream.Close()
+
+	initDockerProxyTest(t, `
+[registries."test.local"]
+upstream = "`+upstream.URL+`"
+authHost = "https://auth.test.local/token"
+authType = "anonymous"
+enabled = true
+`)
+
+	gin.SetMode(gin.TestMode)
+	router := gin.New()
+	router.Any("/v2/*path", ProxyDockerRegistryGin)
+
+	req := httptest.NewRequest(http.MethodGet, "/v2/team/app/manifests/latest?ns=test.local", nil)
+	w := httptest.NewRecorder()
+	router.ServeHTTP(w, req)
+
+	if w.Code != http.StatusOK {
+		t.Fatalf("status = %d, want 200; body=%s", w.Code, w.Body.String())
+	}
+}
+
+func TestProxyDockerAuthCachesOnlyAnonymousTokenRequests(t *testing.T) {
+	var hits int32
+	authServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		count := atomic.AddInt32(&hits, 1)
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = fmt.Fprintf(w, `{"token":"token-%d","expires_in":3600}`, count)
+	}))
+	defer authServer.Close()
+
+	initDockerProxyTest(t, `
+[registries."test.local"]
+upstream = "https://test.local"
+authHost = "`+authServer.URL+`"
+authType = "anonymous"
+enabled = true
+`)
+	utils.GlobalCache = &utils.UniversalCache{}
+
+	gin.SetMode(gin.TestMode)
+	router := gin.New()
+	router.Any("/token/*path", ProxyDockerAuthGin)
+
+	for i := 0; i < 2; i++ {
+		req := httptest.NewRequest(http.MethodGet, "/token/test.local?scope=repository:team/app:pull", nil)
+		w := httptest.NewRecorder()
+		router.ServeHTTP(w, req)
+		if w.Code != http.StatusOK {
+			t.Fatalf("anonymous request %d status = %d; body=%s", i, w.Code, w.Body.String())
+		}
+		if got := w.Body.String(); !strings.Contains(got, `"token":"token-1"`) {
+			t.Fatalf("anonymous request %d body = %q", i, got)
+		}
+	}
+
+	if got := atomic.LoadInt32(&hits); got != 1 {
+		t.Fatalf("anonymous token hits = %d, want 1", got)
+	}
+
+	for i := 0; i < 2; i++ {
+		req := httptest.NewRequest(http.MethodGet, "/token/test.local?scope=repository:team/app:pull", nil)
+		req.Header.Set("Authorization", "Basic dXNlcjpwYXNz")
+		w := httptest.NewRecorder()
+		router.ServeHTTP(w, req)
+		if w.Code != http.StatusOK {
+			t.Fatalf("authenticated request %d status = %d; body=%s", i, w.Code, w.Body.String())
+		}
+	}
+
+	if got := atomic.LoadInt32(&hits); got != 3 {
+		t.Fatalf("authenticated token hits total = %d, want 3", got)
+	}
+}
+
+func TestProxyDockerAuthRejectsUnknownRegistry(t *testing.T) {
+	initDockerProxyTest(t, "")
+
+	gin.SetMode(gin.TestMode)
+	router := gin.New()
+	router.Any("/token/*path", ProxyDockerAuthGin)
+
+	req := httptest.NewRequest(http.MethodGet, "/token/missing.local?scope=repository:team/app:pull", nil)
+	w := httptest.NewRecorder()
+	router.ServeHTTP(w, req)
+
+	if w.Code != http.StatusBadRequest {
+		t.Fatalf("status = %d, want 400; body=%s", w.Code, w.Body.String())
+	}
+}
+
+func TestProxyDockerRegistryConcurrentRequests(t *testing.T) {
+	const requests = 64
+	var hits int32
+
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		atomic.AddInt32(&hits, 1)
+		if got := r.Header.Get("Authorization"); got == "" {
+			t.Fatal("missing Authorization")
+		}
+		w.Header().Set("Content-Type", "application/octet-stream")
+		_, _ = w.Write([]byte("ok"))
+	}))
+	defer upstream.Close()
+
+	initDockerProxyTest(t, `
+[registries."test.local"]
+upstream = "`+upstream.URL+`"
+authHost = "https://auth.test.local/token"
+authType = "anonymous"
+enabled = true
+`)
+
+	gin.SetMode(gin.TestMode)
+	router := gin.New()
+	router.Any("/v2/*path", ProxyDockerRegistryGin)
+
+	var wg sync.WaitGroup
+	errs := make(chan string, requests)
+	for i := 0; i < requests; i++ {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			req := httptest.NewRequest(http.MethodGet, "/v2/test.local/team/app/blobs/sha256:abc", nil)
+			req.Header.Set("Authorization", fmt.Sprintf("Bearer token-%d", i))
+			w := httptest.NewRecorder()
+			router.ServeHTTP(w, req)
+			if w.Code != http.StatusOK || w.Body.String() != "ok" {
+				errs <- fmt.Sprintf("request %d status=%d body=%q", i, w.Code, w.Body.String())
+			}
+		}(i)
+	}
+	wg.Wait()
+	close(errs)
+
+	for err := range errs {
+		t.Fatal(err)
+	}
+	if got := atomic.LoadInt32(&hits); got != requests {
+		t.Fatalf("hits = %d, want %d", got, requests)
+	}
+}
+
+func TestProxyDockerRegistryLargeBlobStreamsWithoutRecorderBuffer(t *testing.T) {
+	const blobSize = 8 << 20
+
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/octet-stream")
+		w.Header().Set("Content-Length", fmt.Sprintf("%d", blobSize))
+		_, _ = io.CopyN(w, zeroReader{}, blobSize)
+	}))
+	defer upstream.Close()
+
+	initDockerProxyTest(t, `
+[registries."test.local"]
+upstream = "`+upstream.URL+`"
+authHost = "https://auth.test.local/token"
+authType = "anonymous"
+enabled = true
+`)
+
+	gin.SetMode(gin.TestMode)
+	router := gin.New()
+	router.Any("/v2/*path", ProxyDockerRegistryGin)
+
+	req := httptest.NewRequest(http.MethodGet, "/v2/test.local/team/app/blobs/sha256:large", nil)
+	w := newDiscardResponseWriter()
+	router.ServeHTTP(w, req)
+
+	if w.status != http.StatusOK {
+		t.Fatalf("status = %d, want 200", w.status)
+	}
+	if w.bytes != blobSize {
+		t.Fatalf("streamed bytes = %d, want %d", w.bytes, blobSize)
+	}
+	if got := w.Header().Get("Content-Length"); got != fmt.Sprintf("%d", blobSize) {
+		t.Fatalf("Content-Length = %q", got)
+	}
+}
+
+func BenchmarkProxyDockerRegistryBlobStreaming(b *testing.B) {
+	const blobSize = 1 << 20
+
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/octet-stream")
+		w.Header().Set("Content-Length", fmt.Sprintf("%d", blobSize))
+		_, _ = io.CopyN(w, zeroReader{}, blobSize)
+	}))
+	defer upstream.Close()
+
+	initDockerProxyTest(b, `
+[registries."bench.local"]
+upstream = "`+upstream.URL+`"
+authHost = "https://auth.bench.local/token"
+authType = "anonymous"
+enabled = true
+`)
+
+	gin.SetMode(gin.TestMode)
+	router := gin.New()
+	router.Any("/v2/*path", ProxyDockerRegistryGin)
+
+	b.ReportAllocs()
+	b.SetBytes(blobSize)
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		req := httptest.NewRequest(http.MethodGet, "/v2/bench.local/team/app/blobs/sha256:bench", nil)
+		w := newDiscardResponseWriter()
+		router.ServeHTTP(w, req)
+		if w.status != http.StatusOK || w.bytes != blobSize {
+			b.Fatalf("status=%d bytes=%d", w.status, w.bytes)
+		}
+	}
+}