Unify AWS ECR and GCP Artifact Registry behavior in Docker reverse proxy

packages/docker-reverse-proxy/internal/handlers/proxy_aws.go

@@ -0,0 +1,107 @@
+package handlers
+
+import (
+	"fmt"
+	"log"
+	"net/http"
+	"os"
+	"strings"
+
+	"github.com/e2b-dev/infra/packages/docker-reverse-proxy/internal/utils"
+)
+
+type AWSProxyHandler struct {
+	apiStore       *APIStore
+	repositoryName string
+	region         string
+	accountID      string
+}
+
+func NewAWSProxyHandler(apiStore *APIStore) *AWSProxyHandler {
+	return &AWSProxyHandler{
+		apiStore:       apiStore,
+		repositoryName: os.Getenv("AWS_DOCKER_REPOSITORY_NAME"),
+		region:         os.Getenv("AWS_REGION"),
+		accountID:      os.Getenv("AWS_ACCOUNT_ID"),
+	}
+}
+
+func (a *AWSProxyHandler) ValidatePath(path string) bool {
+	repoPrefix := "/v2/e2b/custom-envs/"
+	realRepoPrefix := fmt.Sprintf("/v2/%s.dkr.ecr.%s.amazonaws.com/%s/", a.accountID, a.region, a.repositoryName)
+
+	return strings.HasPrefix(path, repoPrefix) || strings.HasPrefix(path, realRepoPrefix)
+}
+
+func (a *AWSProxyHandler) TransformPath(originalPath string, templateID string) string {
+	repoPrefix := "/v2/e2b/custom-envs/"
+	realRepoPrefix := fmt.Sprintf("/v2/%s.dkr.ecr.%s.amazonaws.com/%s/", a.accountID, a.region, a.repositoryName)
+
+	return strings.Replace(originalPath, repoPrefix, realRepoPrefix, 1)
+}
+
+func (a *AWSProxyHandler) GetExpectedTagFormat(templateID string, buildID string) string {
+	return fmt.Sprintf("%s_%s", templateID, buildID)
+}
+
+func (a *AWSProxyHandler) HandleProxy(w http.ResponseWriter, req *http.Request, token *AuthToken) error {
+	path := req.URL.String()
+
+	repoPrefix := "/v2/e2b/custom-envs/"
+	realRepoPrefix := fmt.Sprintf("/v2/%s.dkr.ecr.%s.amazonaws.com/%s/", a.accountID, a.region, a.repositoryName)
+
+	if !a.ValidatePath(path) {
+		// The request shouldn't need any other endpoints, we deny access
+		log.Printf("No matching route found for path: %s\n", path)
+		w.WriteHeader(http.StatusForbidden)
+		return fmt.Errorf("no matching route found for path: %s", path)
+	}
+
+	templateID := token.TemplateID
+
+	// Uploading blobs doesn't have the template ID in the path for AWS ECR
+	if strings.HasPrefix(path, fmt.Sprintf("%sblobs/uploads/", realRepoPrefix)) {
+		a.apiStore.ServeHTTP(w, req)
+		return nil
+	}
+
+	pathInRepo := strings.TrimPrefix(path, repoPrefix)
+
+	// For AWS, we expect the format: templateID:templateID_buildID
+	// Extract the templateID from the path
+	pathParts := strings.Split(pathInRepo, "/")
+	if len(pathParts) > 0 {
+		templateWithTag := strings.Split(pathParts[0], ":")
+		if len(templateWithTag) > 0 {
+			pathTemplateID := templateWithTag[0]
+
+			// If the template ID in the path is different from the token template ID, deny access
+			if pathTemplateID != templateID {
+				w.WriteHeader(http.StatusForbidden)
+				log.Printf("Access denied for template: %s (expected: %s)\n", pathTemplateID, templateID)
+				return fmt.Errorf("access denied for template: %s", pathTemplateID)
+			}
+
+			// For AWS, we need to transform the tag format
+			// From: templateID:templateID_buildID to repository:templateID_buildID
+			if len(templateWithTag) > 1 {
+				// Validate that the tag follows the expected format (templateID_buildID)
+				expectedPrefix := templateID + "_"
+				if !strings.HasPrefix(templateWithTag[1], expectedPrefix) {
+					w.WriteHeader(http.StatusBadRequest)
+					log.Printf("Invalid tag format for AWS: %s (expected format: %s_buildID)\n", templateWithTag[1], templateID)
+					return fmt.Errorf("invalid tag format: %s", templateWithTag[1])
+				}
+			}
+		}
+	}
+
+	// Transform the path for AWS ECR
+	req.URL.Path = a.TransformPath(req.URL.Path, templateID)
+
+	// Set the Authorization header for the request to the real docker registry
+	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token.DockerToken))
+
+	a.apiStore.ServeHTTP(w, req)
+	return nil
+}

packages/docker-reverse-proxy/internal/handlers/proxy_factory.go

@@ -0,0 +1,44 @@
+package handlers
+
+import (
+	"fmt"
+	"os"
+	"strings"
+)
+
+// NewDockerProxyHandler creates the appropriate proxy handler based on environment configuration
+func NewDockerProxyHandler(apiStore *APIStore) (DockerProxyHandler, error) {
+	// Check for cloud provider environment variable
+	provider := strings.ToLower(os.Getenv("CLOUD_PROVIDER"))
+
+	// If not set, try to infer from other environment variables
+	if provider == "" {
+		if os.Getenv("GCP_PROJECT_ID") != "" {
+			provider = "gcp"
+		} else if os.Getenv("AWS_ACCOUNT_ID") != "" {
+			provider = "aws"
+		} else {
+			// Default to GCP for backward compatibility
+			provider = "gcp"
+		}
+	}
+
+	switch provider {
+	case "gcp":
+		return NewGCPProxyHandler(apiStore), nil
+	case "aws":
+		// Validate required AWS environment variables
+		if os.Getenv("AWS_DOCKER_REPOSITORY_NAME") == "" {
+			return nil, fmt.Errorf("AWS_DOCKER_REPOSITORY_NAME environment variable is required for AWS provider")
+		}
+		if os.Getenv("AWS_REGION") == "" {
+			return nil, fmt.Errorf("AWS_REGION environment variable is required for AWS provider")
+		}
+		if os.Getenv("AWS_ACCOUNT_ID") == "" {
+			return nil, fmt.Errorf("AWS_ACCOUNT_ID environment variable is required for AWS provider")
+		}
+		return NewAWSProxyHandler(apiStore), nil
+	default:
+		return nil, fmt.Errorf("unsupported cloud provider: %s (supported: gcp, aws)", provider)
+	}
+}

packages/docker-reverse-proxy/internal/handlers/proxy_gcp.go

@@ -0,0 +1,90 @@
+package handlers
+
+import (
+	"fmt"
+	"log"
+	"net/http"
+	"strings"
+
+	"github.com/e2b-dev/infra/packages/docker-reverse-proxy/internal/constants"
+	"github.com/e2b-dev/infra/packages/docker-reverse-proxy/internal/utils"
+	"github.com/e2b-dev/infra/packages/shared/pkg/consts"
+)
+
+type GCPProxyHandler struct {
+	apiStore *APIStore
+}
+
+func NewGCPProxyHandler(apiStore *APIStore) *GCPProxyHandler {
+	return &GCPProxyHandler{
+		apiStore: apiStore,
+	}
+}
+
+func (g *GCPProxyHandler) ValidatePath(path string) bool {
+	repoPrefix := "/v2/e2b/custom-envs/"
+	realRepoPrefix := fmt.Sprintf("/v2/%s/%s/", consts.GCPProject, consts.DockerRegistry)
+
+	return strings.HasPrefix(path, repoPrefix) || strings.HasPrefix(path, realRepoPrefix)
+}
+
+func (g *GCPProxyHandler) TransformPath(originalPath string, templateID string) string {
+	repoPrefix := "/v2/e2b/custom-envs/"
+	realRepoPrefix := fmt.Sprintf("/v2/%s/%s/", consts.GCPProject, consts.DockerRegistry)
+
+	return strings.Replace(originalPath, repoPrefix, realRepoPrefix, 1)
+}
+
+func (g *GCPProxyHandler) GetExpectedTagFormat(templateID string, buildID string) string {
+	return fmt.Sprintf("%s:%s", templateID, buildID)
+}
+
+func (g *GCPProxyHandler) HandleProxy(w http.ResponseWriter, req *http.Request, token *AuthToken) error {
+	path := req.URL.String()
+
+	// Allow access to the GCP artifact uploads.
+	// The url is generated by repository and sent as a Location header from the /blobs/upload request
+	// https://distribution.github.io/distribution/spec/api/#starting-an-upload
+	// Other methods than PATCH require the Authorization header
+	if strings.HasPrefix(path, constants.GCPArtifactUploadPrefix) {
+		g.apiStore.ServeHTTP(w, req)
+		return nil
+	}
+
+	repoPrefix := "/v2/e2b/custom-envs/"
+	realRepoPrefix := fmt.Sprintf("/v2/%s/%s/", consts.GCPProject, consts.DockerRegistry)
+
+	if !g.ValidatePath(path) {
+		// The request shouldn't need any other endpoints, we deny access
+		log.Printf("No matching route found for path: %s\n", path)
+		w.WriteHeader(http.StatusForbidden)
+		return fmt.Errorf("no matching route found for path: %s", path)
+	}
+
+	templateID := token.TemplateID
+
+	// Uploading blobs doesn't have the template ID in the path
+	if strings.HasPrefix(path, fmt.Sprintf("%spkg/blobs/uploads/", realRepoPrefix)) {
+		g.apiStore.ServeHTTP(w, req)
+		return nil
+	}
+
+	pathInRepo := strings.TrimPrefix(path, repoPrefix)
+	templateWithBuildID := strings.Split(strings.Split(pathInRepo, "/")[0], ":")
+
+	// If the template ID in the path is different from the token template ID, deny access
+	if templateWithBuildID[0] != templateID {
+		w.WriteHeader(http.StatusForbidden)
+		log.Printf("Access denied for template: %s\n", templateID)
+		return fmt.Errorf("access denied for template: %s", templateID)
+	}
+
+	// Set the host and access token for the real docker registry
+	req.URL.Path = g.TransformPath(req.URL.Path, templateID)
+
+	// Set the Authorization header for the request to the real docker registry
+	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token.DockerToken))
+
+	g.apiStore.ServeHTTP(w, req)
+	return nil
+}

packages/docker-reverse-proxy/internal/handlers/proxy_interface.go

@@ -0,0 +1,27 @@
+package handlers
+
+import (
+	"net/http"
+)
+
+// AuthToken represents the authentication token information
+type AuthToken struct {
+	TemplateID   string
+	DockerToken  string
+	ExpiresIn    int
+}
+
+// DockerProxyHandler defines the interface for handling docker registry proxy requests
+type DockerProxyHandler interface {
+	// HandleProxy processes the docker registry request
+	HandleProxy(w http.ResponseWriter, req *http.Request, token *AuthToken) error
+
+	// ValidatePath checks if the request path is valid for this provider
+	ValidatePath(path string) bool
+
+	// TransformPath converts the incoming path to the target registry path
+	TransformPath(originalPath string, templateID string) string
+
+	// GetExpectedTagFormat returns the expected tag format for this provider
+	GetExpectedTagFormat(templateID string, buildID string) string
+}

packages/shared/pkg/artifacts-registry/registry_aws.go

@@ -17,9 +17,16 @@ import (
 	"github.com/google/go-containerregistry/pkg/v1/remote"
 )
 
+// ECRClient interface for testability
+type ECRClient interface {
+	DescribeRepositories(ctx context.Context, input *ecr.DescribeRepositoriesInput, opts ...func(*ecr.Options)) (*ecr.DescribeRepositoriesOutput, error)
+	BatchDeleteImage(ctx context.Context, input *ecr.BatchDeleteImageInput, opts ...func(*ecr.Options)) (*ecr.BatchDeleteImageOutput, error)
+	GetAuthorizationToken(ctx context.Context, input *ecr.GetAuthorizationTokenInput, opts ...func(*ecr.Options)) (*ecr.GetAuthorizationTokenOutput, error)
+}
+
 type AWSArtifactsRegistry struct {
 	repositoryName string
-	client         *ecr.Client
+	client         ECRClient
 }
 
 var (
@@ -68,6 +75,12 @@ func (g *AWSArtifactsRegistry) Delete(ctx context.Context, templateId string, bu
 }
 
 func (g *AWSArtifactsRegistry) GetTag(ctx context.Context, templateId string, buildId string) (string, error) {
+	// Generate composite tag using templateId and buildId first (includes validation)
+	compositeTag, err := GenerateCompositeTag(templateId, buildId)
+	if err != nil {
+		return "", fmt.Errorf("failed to generate composite tag: %w", err)
+	}
+
 	res, err := g.client.DescribeRepositories(ctx, &ecr.DescribeRepositoriesInput{RepositoryNames: []string{g.repositoryName}})
 	if err != nil {
 		return "", fmt.Errorf("failed to describe aws ecr repository: %w", err)
@@ -77,7 +90,7 @@ func (g *AWSArtifactsRegistry) GetTag(ctx context.Context, templateId string, bu
 		return "", fmt.Errorf("repository %s not found", g.repositoryName)
 	}
 
-	return fmt.Sprintf("%s:%s", *res.Repositories[0].RepositoryUri, buildId), nil
+	return fmt.Sprintf("%s:%s", *res.Repositories[0].RepositoryUri, compositeTag), nil
 }
 
 func (g *AWSArtifactsRegistry) GetImage(ctx context.Context, templateId string, buildId string, platform containerregistry.Platform) (containerregistry.Image, error) {