Fix support for public.ecr.aws

Signed-off-by: Matheus Pimenta <[email protected]>

Author: matheuscscp

.github/workflows/integration-aws.yaml

@@ -20,6 +20,7 @@ jobs:
         auth-mode:
         - node-identity
         - workload-identity
+      fail-fast: false
     defaults:
       run:
         working-directory: ./oci/tests/integration
@@ -38,7 +39,7 @@ jobs:
         with:
           role-to-assume: arn:aws:iam::${{ secrets.AWS_ACCOUNT_ID }}:role/${{ secrets.OCI_E2E_AWS_ASSUME_ROLE_NAME }}
           role-session-name: OCI_GH_Actions
-          aws-region: ${{ vars.AWS_REGION }}
+          aws-region: us-east-1
       - name: Setup QEMU
         uses: docker/setup-qemu-action@29109295f81e9208d7d86ff1c6c12d2833863392 # v3.6.0
       - name: Setup Docker Buildx
@@ -56,13 +57,13 @@ jobs:
       - name: Run tests
         run: . .env && make test-aws
         env:
-          AWS_REGION: ${{ vars.AWS_REGION }}
-          TF_VAR_cross_region: ${{ vars.OCI_E2E_TF_VAR_cross_region }}
+          AWS_REGION: us-east-1
+          TF_VAR_cross_region: us-east-2
           TF_VAR_enable_wi: ${{ (matrix.auth-mode == 'workload-identity' && 'true') || 'false' }}
       - name: Ensure resource cleanup
         if: ${{ always() }}
         run: . .env && make destroy-aws
         env:
-          AWS_REGION: ${{ vars.AWS_REGION }}
-          TF_VAR_cross_region: ${{ vars.OCI_E2E_TF_VAR_cross_region }}
+          AWS_REGION: us-east-1
+          TF_VAR_cross_region: us-east-2
           TF_VAR_enable_wi: ${{ (matrix.auth-mode == 'workload-identity' && 'true') || 'false' }}

auth/aws/implementation.go

@@ -22,14 +22,16 @@ import (
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/service/ecr"
+	"github.com/aws/aws-sdk-go-v2/service/ecrpublic"
 	"github.com/aws/aws-sdk-go-v2/service/sts"
 )
 
 // Implementation provides the required methods of the AWS libraries.
 type Implementation interface {
 	LoadDefaultConfig(ctx context.Context, optFns ...func(*config.LoadOptions) error) (aws.Config, error)
 	AssumeRoleWithWebIdentity(ctx context.Context, params *sts.AssumeRoleWithWebIdentityInput, options sts.Options) (*sts.AssumeRoleWithWebIdentityOutput, error)
-	GetAuthorizationToken(ctx context.Context, cfg aws.Config) (*ecr.GetAuthorizationTokenOutput, error)
+	GetAuthorizationToken(ctx context.Context, cfg aws.Config) (any, error)
+	GetPublicAuthorizationToken(ctx context.Context, cfg aws.Config) (any, error)
 }
 
 type implementation struct{}
@@ -42,6 +44,10 @@ func (implementation) AssumeRoleWithWebIdentity(ctx context.Context, params *sts
 	return sts.New(options).AssumeRoleWithWebIdentity(ctx, params)
 }
 
-func (implementation) GetAuthorizationToken(ctx context.Context, cfg aws.Config) (*ecr.GetAuthorizationTokenOutput, error) {
+func (implementation) GetAuthorizationToken(ctx context.Context, cfg aws.Config) (any, error) {
 	return ecr.NewFromConfig(cfg).GetAuthorizationToken(ctx, &ecr.GetAuthorizationTokenInput{})
 }
+
+func (implementation) GetPublicAuthorizationToken(ctx context.Context, cfg aws.Config) (any, error) {
+	return ecrpublic.NewFromConfig(cfg).GetAuthorizationToken(ctx, &ecrpublic.GetAuthorizationTokenInput{})
+}

auth/aws/implementation_test.go

@@ -27,6 +27,8 @@ import (
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/service/ecr"
 	ecrtypes "github.com/aws/aws-sdk-go-v2/service/ecr/types"
+	"github.com/aws/aws-sdk-go-v2/service/ecrpublic"
+	ecrpublictypes "github.com/aws/aws-sdk-go-v2/service/ecrpublic/types"
 	"github.com/aws/aws-sdk-go-v2/service/sts"
 	ststypes "github.com/aws/aws-sdk-go-v2/service/sts/types"
 	. "github.com/onsi/gomega"
@@ -35,6 +37,8 @@ import (
 type mockImplementation struct {
 	t *testing.T
 
+	publicECR bool
+
 	argRoleARN         string
 	argRoleSessionName string
 	argOIDCToken       string
@@ -101,7 +105,33 @@ func (m *mockImplementation) AssumeRoleWithWebIdentity(ctx context.Context, para
 	}, nil
 }
 
-func (m *mockImplementation) GetAuthorizationToken(ctx context.Context, cfg aws.Config) (*ecr.GetAuthorizationTokenOutput, error) {
+func (m *mockImplementation) GetAuthorizationToken(ctx context.Context, cfg aws.Config) (any, error) {
+	m.t.Helper()
+	g := NewWithT(m.t)
+	g.Expect(m.publicECR).To(BeFalse())
+	m.checkGetAuthorizationToken(ctx, cfg)
+	return &ecr.GetAuthorizationTokenOutput{
+		AuthorizationData: []ecrtypes.AuthorizationData{{
+			AuthorizationToken: aws.String(base64.StdEncoding.EncodeToString([]byte(m.returnUsername + ":" + m.returnPassword))),
+			ExpiresAt:          aws.Time(m.returnCreds.Expires),
+		}},
+	}, nil
+}
+
+func (m *mockImplementation) GetPublicAuthorizationToken(ctx context.Context, cfg aws.Config) (any, error) {
+	m.t.Helper()
+	g := NewWithT(m.t)
+	g.Expect(m.publicECR).To(BeTrue())
+	m.checkGetAuthorizationToken(ctx, cfg)
+	return &ecrpublic.GetAuthorizationTokenOutput{
+		AuthorizationData: &ecrpublictypes.AuthorizationData{
+			AuthorizationToken: aws.String(base64.StdEncoding.EncodeToString([]byte(m.returnUsername + ":" + m.returnPassword))),
+			ExpiresAt:          aws.Time(m.returnCreds.Expires),
+		},
+	}, nil
+}
+
+func (m *mockImplementation) checkGetAuthorizationToken(ctx context.Context, cfg aws.Config) {
 	m.t.Helper()
 	g := NewWithT(m.t)
 	g.Expect(cfg.Region).To(Equal(m.argRegion))
@@ -114,11 +144,6 @@ func (m *mockImplementation) GetAuthorizationToken(ctx context.Context, cfg aws.
 	proxyURL, err := cfg.HTTPClient.(*http.Client).Transport.(*http.Transport).Proxy(nil)
 	g.Expect(err).NotTo(HaveOccurred())
 	g.Expect(proxyURL).To(Equal(m.argProxyURL))
-	return &ecr.GetAuthorizationTokenOutput{
-		AuthorizationData: []ecrtypes.AuthorizationData{{
-			AuthorizationToken: aws.String(base64.StdEncoding.EncodeToString([]byte(m.returnUsername + ":" + m.returnPassword))),
-		}},
-	}, nil
 }
 
 func (m *mockCredentialsProvider) Retrieve(ctx context.Context) (aws.Credentials, error) {

auth/aws/provider.go

@@ -29,6 +29,8 @@ import (
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/service/ecr"
+	"github.com/aws/aws-sdk-go-v2/service/ecrpublic"
 	"github.com/aws/aws-sdk-go-v2/service/sts"
 	"github.com/google/go-containerregistry/pkg/authn"
 	corev1 "k8s.io/api/core/v1"
@@ -70,8 +72,8 @@ func (p Provider) NewControllerToken(ctx context.Context, opts ...auth.Option) (
 		case o.ArtifactRepository != "":
 			// We can safely ignore the error here, auth.GetToken() has already called
 			// ParseArtifactRepository() and validated the repository at this point.
-			ecrRegion, _ := p.ParseArtifactRepository(o.ArtifactRepository)
-			stsRegion = ecrRegion
+			registryInput, _ := p.ParseArtifactRepository(o.ArtifactRepository)
+			stsRegion = getECRRegionFromRegistryInput(registryInput)
 		// EKS sets this environment variable automatically if the controller pod is
 		// properly configured with IRSA or EKS Pod Identity, so we can rely on this
 		// and communicate this to users since this is controller-level configuration.
@@ -140,8 +142,8 @@ func (p Provider) NewTokenForServiceAccount(ctx context.Context, oidcToken strin
 		case o.ArtifactRepository != "":
 			// We can safely ignore the error here, auth.GetToken() has already called
 			// ParseArtifactRepository() and validated the repository at this point.
-			ecrRegion, _ := p.ParseArtifactRepository(o.ArtifactRepository)
-			stsRegion = ecrRegion
+			registryInput, _ := p.ParseArtifactRepository(o.ArtifactRepository)
+			stsRegion = getECRRegionFromRegistryInput(registryInput)
 		// In this case we can't rely on IRSA or EKS Pod Identity for the controller
 		// pod because this is object-level configuration, so we show a different
 		// error message.
@@ -204,20 +206,21 @@ func (p Provider) NewTokenForServiceAccount(ctx context.Context, oidcToken strin
 // It covers both public AWS partitions like amazonaws.com, China partitions like amazonaws.com.cn, and non-public partitions.
 const registryPattern = `([0-9+]*).dkr.ecr(?:-fips)?\.([^/.]*)\.(amazonaws\.com[.cn]*|sc2s\.sgov\.gov|c2s\.ic\.gov|cloud\.adc-e\.uk|csp\.hci\.ic\.gov)`
 
+const publicECR = "public.ecr.aws"
+
 var registryRegex = regexp.MustCompile(registryPattern)
 
 // ParseArtifactRepository implements auth.Provider.
-// ParseArtifactRepository returns the ECR region.
+// ParseArtifactRepository returns the ECR region, unless the registry
+// is public.ecr.aws, in which case it returns public.ecr.aws.
 func (Provider) ParseArtifactRepository(artifactRepository string) (string, error) {
 	registry, err := auth.GetRegistryFromArtifactRepository(artifactRepository)
 	if err != nil {
 		return "", err
 	}
 
-	// Region is required to be us-east-1 for public.ecr.aws:
-	// https://docs.aws.amazon.com/AmazonECR/latest/public/public-registry-auth.html#public-registry-auth-token
-	if registry == "public.ecr.aws" {
-		return "us-east-1", nil
+	if registry == publicECR {
+		return publicECR, nil
 	}
 
 	parts := registryRegex.FindAllStringSubmatch(registry, -1)
@@ -231,36 +234,70 @@ func (Provider) ParseArtifactRepository(artifactRepository string) (string, erro
 	return ecrRegion, nil
 }
 
+func getECRRegionFromRegistryInput(registryInput string) string {
+	if registryInput == publicECR {
+		// Region is required to be us-east-1 for public ECR:
+		// https://docs.aws.amazon.com/AmazonECR/latest/public/public-registry-auth.html#public-registry-auth-token
+		return "us-east-1"
+	}
+	return registryInput
+}
+
 // NewArtifactRegistryCredentials implements auth.Provider.
-func (p Provider) NewArtifactRegistryCredentials(ctx context.Context, ecrRegion string,
+func (p Provider) NewArtifactRegistryCredentials(ctx context.Context, registryInput string,
 	accessToken auth.Token, opts ...auth.Option) (*auth.ArtifactRegistryCredentials, error) {
 
 	var o auth.Options
 	o.Apply(opts...)
 
+	authTokenFunc := p.impl().GetAuthorizationToken
+	if registryInput == publicECR {
+		authTokenFunc = p.impl().GetPublicAuthorizationToken
+	}
+
 	conf := aws.Config{
-		Region:      ecrRegion,
+		Region:      getECRRegionFromRegistryInput(registryInput),
 		Credentials: accessToken.(*Token).CredentialsProvider(),
 	}
 
 	if hc := o.GetHTTPClient(); hc != nil {
 		conf.HTTPClient = hc
 	}
 
-	resp, err := p.impl().GetAuthorizationToken(ctx, conf)
+	respAny, err := authTokenFunc(ctx, conf)
 	if err != nil {
 		return nil, err
 	}
 
 	// Parse the authorization token.
-	if len(resp.AuthorizationData) == 0 {
-		return nil, fmt.Errorf("no authorization data returned")
-	}
-	tokenResp := resp.AuthorizationData[0]
-	if tokenResp.AuthorizationToken == nil {
-		return nil, fmt.Errorf("authorization token is nil")
+	var token string
+	var expiresAt time.Time
+	switch resp := respAny.(type) {
+	case *ecr.GetAuthorizationTokenOutput:
+		if len(resp.AuthorizationData) == 0 {
+			return nil, fmt.Errorf("no authorization data returned")
+		}
+		if resp.AuthorizationData[0].AuthorizationToken == nil {
+			return nil, fmt.Errorf("authorization token is nil")
+		}
+		if resp.AuthorizationData[0].ExpiresAt == nil {
+			return nil, fmt.Errorf("authorization token expiration is nil")
+		}
+		token = *resp.AuthorizationData[0].AuthorizationToken
+		expiresAt = *resp.AuthorizationData[0].ExpiresAt
+	case *ecrpublic.GetAuthorizationTokenOutput:
+		if resp.AuthorizationData == nil {
+			return nil, fmt.Errorf("no authorization data returned")
+		}
+		if resp.AuthorizationData.AuthorizationToken == nil {
+			return nil, fmt.Errorf("authorization token is nil")
+		}
+		if resp.AuthorizationData.ExpiresAt == nil {
+			return nil, fmt.Errorf("authorization token expiration is nil")
+		}
+		token = *resp.AuthorizationData.AuthorizationToken
+		expiresAt = *resp.AuthorizationData.ExpiresAt
 	}
-	token := *tokenResp.AuthorizationToken
 	b, err := base64.StdEncoding.DecodeString(token)
 	if err != nil {
 		return nil, fmt.Errorf("failed to parse authorization token: %w", err)
@@ -269,10 +306,6 @@ func (p Provider) NewArtifactRegistryCredentials(ctx context.Context, ecrRegion
 	if len(s) != 2 {
 		return nil, fmt.Errorf("invalid authorization token format")
 	}
-	var expiresAt time.Time
-	if exp := tokenResp.ExpiresAt; exp != nil {
-		expiresAt = *exp
-	}
 	return &auth.ArtifactRegistryCredentials{
 		Authenticator: authn.FromConfig(authn.AuthConfig{
 			Username: s[0],

auth/aws/provider_test.go

@@ -225,39 +225,67 @@ func TestProvider_GetIdentity(t *testing.T) {
 }
 
 func TestProvider_NewArtifactRegistryCredentials(t *testing.T) {
-	g := NewWithT(t)
+	for _, tt := range []struct {
+		name              string
+		registryInput     string
+		expectedPublicECR bool
+		expectedRegion    string
+	}{
+		{
+			name:              "non public ECR",
+			registryInput:     "us-east-1",
+			expectedRegion:    "us-east-1",
+			expectedPublicECR: false,
+		},
+		{
+			name:              "non public ECR",
+			registryInput:     "us-west-2",
+			expectedRegion:    "us-west-2",
+			expectedPublicECR: false,
+		},
+		{
+			name:              "public ECR",
+			registryInput:     "public.ecr.aws",
+			expectedRegion:    "us-east-1", // Public ECR is always us-east-1
+			expectedPublicECR: true,
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			g := NewWithT(t)
 
-	impl := &mockImplementation{
-		t:                t,
-		argRegion:        "us-east-1",
-		argProxyURL:      &url.URL{Scheme: "http", Host: "proxy.example.com"},
-		argCredsProvider: credentials.NewStaticCredentialsProvider("access-key-id", "secret-access-key", "session-token"),
-		returnUsername:   "username",
-		returnPassword:   "password",
-	}
+			impl := &mockImplementation{
+				t:                t,
+				publicECR:        tt.expectedPublicECR,
+				argRegion:        tt.expectedRegion,
+				argProxyURL:      &url.URL{Scheme: "http", Host: "proxy.example.com"},
+				argCredsProvider: credentials.NewStaticCredentialsProvider("access-key-id", "secret-access-key", "session-token"),
+				returnUsername:   "username",
+				returnPassword:   "password",
+			}
 
-	ecrRegion := "us-east-1"
-	accessToken := &aws.Token{
-		Credentials: types.Credentials{
-			AccessKeyId:     awssdk.String("access-key-id"),
-			SecretAccessKey: awssdk.String("secret-access-key"),
-			SessionToken:    awssdk.String("session-token"),
-		},
-	}
-	opts := []auth.Option{
-		auth.WithProxyURL(url.URL{Scheme: "http", Host: "proxy.example.com"}),
-	}
+			accessToken := &aws.Token{
+				Credentials: types.Credentials{
+					AccessKeyId:     awssdk.String("access-key-id"),
+					SecretAccessKey: awssdk.String("secret-access-key"),
+					SessionToken:    awssdk.String("session-token"),
+				},
+			}
+			opts := []auth.Option{
+				auth.WithProxyURL(url.URL{Scheme: "http", Host: "proxy.example.com"}),
+			}
 
-	provider := aws.Provider{Implementation: impl}
-	creds, err := provider.NewArtifactRegistryCredentials(
-		context.Background(), ecrRegion, accessToken, opts...)
-	g.Expect(err).NotTo(HaveOccurred())
-	g.Expect(creds).To(Equal(&auth.ArtifactRegistryCredentials{
-		Authenticator: authn.FromConfig(authn.AuthConfig{
-			Username: "username",
-			Password: "password",
-		}),
-	}))
+			provider := aws.Provider{Implementation: impl}
+			creds, err := provider.NewArtifactRegistryCredentials(
+				context.Background(), tt.registryInput, accessToken, opts...)
+			g.Expect(err).NotTo(HaveOccurred())
+			g.Expect(creds).To(Equal(&auth.ArtifactRegistryCredentials{
+				Authenticator: authn.FromConfig(authn.AuthConfig{
+					Username: "username",
+					Password: "password",
+				}),
+			}))
+		})
+	}
 }
 
 func TestProvider_ParseArtifactRepository(t *testing.T) {
@@ -322,7 +350,7 @@ func TestProvider_ParseArtifactRepository(t *testing.T) {
 		},
 		{
 			artifactRepository: "public.ecr.aws/foo/bar",
-			expectedRegion:     "us-east-1",
+			expectedRegion:     "public.ecr.aws",
 			expectValid:        true,
 		},
 	}

auth/go.mod

@@ -12,6 +12,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/config v1.29.14
 	github.com/aws/aws-sdk-go-v2/credentials v1.17.67
 	github.com/aws/aws-sdk-go-v2/service/ecr v1.43.3
+	github.com/aws/aws-sdk-go-v2/service/ecrpublic v1.33.0
 	github.com/aws/aws-sdk-go-v2/service/sts v1.33.19
 	github.com/coreos/go-oidc/v3 v3.14.1
 	github.com/fluxcd/pkg/cache v0.9.0