From 7fd03005b2ed76bf10bd031d03d2931dbbdaeb28 Mon Sep 17 00:00:00 2001
From: David Sauer <davedamoon@gmail.com>
Date: Wed, 23 Oct 2024 12:31:00 +0200
Subject: [PATCH] fixup! fixup! use / from docker hub

---
 lib/http.go | 46 +++++++++++++++++++++++++++++++---------------
 1 file changed, 31 insertions(+), 15 deletions(-)

diff --git a/lib/http.go b/lib/http.go
index 6c6aade..783bd1c 100644
--- a/lib/http.go
+++ b/lib/http.go
@@ -12,19 +12,22 @@ import (
 	"github.com/aws/aws-sdk-go-v2/service/ecr"
 )
 
-func fetchDockerHubPrefix(ctx context.Context, svc *ecr.Client) (string, error) {
+func fetchPullThroughCachePrefixes(ctx context.Context, svc *ecr.Client) (map[string]string, error) {
 	result, err := svc.DescribePullThroughCacheRules(ctx, &ecr.DescribePullThroughCacheRulesInput{})
 	if err != nil {
-		return "", fmt.Errorf("failed to describe pull-through cache rules: %v", err)
+		return nil, fmt.Errorf("failed to describe pull-through cache rules: %v", err)
 	}
 
+	prefixes := make(map[string]string)
 	for _, rule := range result.PullThroughCacheRules {
-		if *rule.UpstreamRegistryUrl == "registry-1.docker.io" {
-			return *rule.EcrRepositoryPrefix, nil
-		}
+		prefixes[*rule.UpstreamRegistryUrl] = *rule.EcrRepositoryPrefix
+	}
+
+	if len(prefixes) == 0 {
+		return nil, fmt.Errorf("no pull-through cache rules found")
 	}
 
-	return "", fmt.Errorf("no Docker Hub pull-through cache rule found")
+	return prefixes, nil
 }
 
 func RunHttpServer(ctx context.Context, port int) error {
@@ -34,14 +37,14 @@ func RunHttpServer(ctx context.Context, port int) error {
 		return fmt.Errorf("unable to load SDK config, %v", err)
 	}
 
-	// Fetch Docker Hub prefix
-	dockerHubPrefix, err := fetchDockerHubPrefix(ctx, svc)
+	// Fetch pull-through cache prefixes
+	prefixes, err := fetchPullThroughCachePrefixes(ctx, svc)
 	if err != nil {
-		return fmt.Errorf("failed to fetch Docker Hub prefix: %v", err)
+		return fmt.Errorf("failed to fetch pull-through cache prefixes: %v", err)
 	}
 
 	// Setup HTTP server
-	http.HandleFunc("/", handler(ctx, svc, dockerHubPrefix))
+	http.HandleFunc("/", handler(ctx, svc, prefixes))
 
 	addr := fmt.Sprintf("127.0.0.1:%d", port)
 
@@ -73,15 +76,28 @@ func getECRAuthToken(ctx context.Context, svc *ecr.Client) (string, string, erro
 }
 
 // newProxy creates a new ReverseProxy that forwards requests to the ECR domain.
-func newProxy(ecrURL *url.URL, authToken string, dockerHubPrefix string) *httputil.ReverseProxy {
+func newProxy(ecrURL *url.URL, authToken string, prefixes map[string]string) *httputil.ReverseProxy {
 	proxy := httputil.NewSingleHostReverseProxy(ecrURL)
 	originalDirector := proxy.Director
 
 	proxy.Director = func(req *http.Request) {
 		log.Printf("Original request path: %s", req.URL.Path)
-		if strings.HasPrefix(req.URL.Path, "/v2/") {
-			req.URL.Path = "/" + dockerHubPrefix + req.URL.Path
+
+		// Check if the request path starts with any of the pull-through cache prefixes
+		shouldModify := true
+		for _, prefix := range prefixes {
+			if strings.HasPrefix(req.URL.Path, "/"+prefix) {
+				shouldModify = false
+				break
+			}
 		}
+
+		// If the path doesn't start with any prefix, prepend the prefix of dockerhub
+		if shouldModify {
+			dockerHubPrefix := prefixes["registry-1.docker.io"]
+			req.URL.Path = strings.Replace(req.URL.Path, "/v2/", "/v2/"+dockerHubPrefix, 1)
+		}
+
 		log.Printf("Modified request path: %s", req.URL.Path)
 
 		originalDirector(req)
@@ -99,7 +115,7 @@ func newProxy(ecrURL *url.URL, authToken string, dockerHubPrefix string) *httput
 }
 
 // handler forwards incoming requests to the ECR endpoint with proper authorization headers.
-func handler(ctx context.Context, svc *ecr.Client, dockerHubPrefix string) http.HandlerFunc {
+func handler(ctx context.Context, svc *ecr.Client, prefixes map[string]string) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		// Retrieve a fresh token and the ECR domain.
 		authToken, proxyEndpoint, err := getECRAuthToken(ctx, svc)
@@ -116,7 +132,7 @@ func handler(ctx context.Context, svc *ecr.Client, dockerHubPrefix string) http.
 		}
 
 		// Create a new ReverseProxy and forward the request.
-		proxy := newProxy(ecrURL, authToken, dockerHubPrefix)
+		proxy := newProxy(ecrURL, authToken, prefixes)
 		proxy.ServeHTTP(w, r)
 	}
 }