From bbc720490112968c09aa04a597f8192afca2ea35 Mon Sep 17 00:00:00 2001
From: Rogee <rogee@ipao.vip>
Date: Mon, 23 Mar 2026 15:42:45 +0800
Subject: [PATCH] fix: limit docker library rewrite to Docker Hub

---
 internal/hubmodule/docker/hooks.go      | 7 +++++++
 internal/hubmodule/docker/hooks_test.go | 8 ++++++++
 tests/integration/cache_flow_test.go    | 6 +++---
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/internal/hubmodule/docker/hooks.go b/internal/hubmodule/docker/hooks.go
index 4111c3c..5a2bd34 100644
--- a/internal/hubmodule/docker/hooks.go
+++ b/internal/hubmodule/docker/hooks.go
@@ -1,6 +1,7 @@
 package docker
 
 import (
+	"net"
 	"strings"
 
 	"github.com/any-hub/any-hub/internal/proxy/hooks"
@@ -15,6 +16,9 @@ func init() {
 }
 
 func normalizePath(ctx *hooks.RequestContext, clean string, rawQuery []byte) (string, []byte) {
+	if ctx == nil || !isDockerHubHost(ctx.UpstreamHost) {
+		return clean, rawQuery
+	}
 	repo, rest, ok := splitDockerRepoPath(clean)
 	if !ok || repo == "" || strings.Contains(repo, "/") || repo == "library" {
 		return clean, rawQuery
@@ -54,6 +58,9 @@ func contentType(_ *hooks.RequestContext, locatorPath string) string {
 }
 
 func isDockerHubHost(host string) bool {
+	if parsedHost, _, err := net.SplitHostPort(host); err == nil {
+		host = parsedHost
+	}
 	switch strings.ToLower(host) {
 	case "registry-1.docker.io", "docker.io", "index.docker.io":
 		return true
diff --git a/internal/hubmodule/docker/hooks_test.go b/internal/hubmodule/docker/hooks_test.go
index 53a8b56..5e2a563 100644
--- a/internal/hubmodule/docker/hooks_test.go
+++ b/internal/hubmodule/docker/hooks_test.go
@@ -19,6 +19,14 @@ func TestNormalizePathAddsLibraryForDockerHub(t *testing.T) {
 	}
 }
 
+func TestNormalizePathSkipsLibraryForNonDockerHub(t *testing.T) {
+	ctx := &hooks.RequestContext{UpstreamHost: "registry.k8s.io"}
+	path, _ := normalizePath(ctx, "/v2/kube-apiserver/manifests/v1.35.3", nil)
+	if path != "/v2/kube-apiserver/manifests/v1.35.3" {
+		t.Fatalf("expected non-docker hub path to remain unchanged, got %s", path)
+	}
+}
+
 func TestSplitDockerRepoPath(t *testing.T) {
 	repo, rest, ok := splitDockerRepoPath("/v2/library/nginx/manifests/latest")
 	if !ok || repo != "library/nginx" || rest != "/manifests/latest" {
diff --git a/tests/integration/cache_flow_test.go b/tests/integration/cache_flow_test.go
index 0cf9811..a89a8f9 100644
--- a/tests/integration/cache_flow_test.go
+++ b/tests/integration/cache_flow_test.go
@@ -222,8 +222,8 @@ func TestDockerManifestHeadDoesNotOverwriteCache(t *testing.T) {
 	}
 }
 
-func TestDockerNamespaceFallbackAddsLibrary(t *testing.T) {
-	stub := newCacheFlowStub(t, dockerManifestPath)
+func TestDockerNonDockerHubUpstreamKeepsOriginalPath(t *testing.T) {
+	stub := newCacheFlowStub(t, dockerManifestNoNamespacePath)
 	defer stub.Close()
 
 	storageDir := t.TempDir()
@@ -277,7 +277,7 @@ func TestDockerNamespaceFallbackAddsLibrary(t *testing.T) {
 	}
 	if resp.StatusCode != fiber.StatusOK {
 		body, _ := io.ReadAll(resp.Body)
-		t.Fatalf("expected 200 when fallback applies, got %d (body=%s)", resp.StatusCode, string(body))
+		t.Fatalf("expected 200 when upstream keeps original path, got %d (body=%s)", resp.StatusCode, string(body))
 	}
 	resp.Body.Close()