full cleanup + tests

Author: gargipanatula

cmd/ecr-credential-provider/main.go

@@ -30,36 +30,24 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/ecr"
-	"github.com/aws/aws-sdk-go/service/ecrpublic"
-	"github.com/spf13/cobra"
 
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	"k8s.io/component-base/logs"
 	"k8s.io/klog/v2"
 	v1 "k8s.io/kubelet/pkg/apis/credentialprovider/v1"
 )
 
-const ecrPublicRegion string = "us-east-1"
-const ecrPublicHost string = "public.ecr.aws"
-
-var ecrPrivateHostPattern = regexp.MustCompile(`^(\d{12})\.dkr[\.\-]ecr(\-fips)?\.([a-zA-Z0-9][a-zA-Z0-9-_]*)\.(amazonaws\.com(?:\.cn)?|on\.(?:aws|amazonwebservices\.com\.cn)|sc2s\.sgov\.gov|c2s\.ic\.gov|cloud\.adc-e\.uk|csp\.hci\.ic\.gov)$`)
+var ecrPattern = regexp.MustCompile(`^(\d{12})\.dkr\.ecr(\-fips)?\.([a-zA-Z0-9][a-zA-Z0-9-_]*)\.(amazonaws\.com(\.cn)?|sc2s\.sgov\.gov|c2s\.ic\.gov|cloud\.adc-e\.uk|csp\.hci\.ic\.gov)$`)
 
 // ECR abstracts the calls we make to aws-sdk for testing purposes
 type ECR interface {
 	GetAuthorizationToken(input *ecr.GetAuthorizationTokenInput) (*ecr.GetAuthorizationTokenOutput, error)
 }
 
-// ECRPublic abstracts the calls we make to aws-sdk for testing purposes
-type ECRPublic interface {
-	GetAuthorizationToken(input *ecrpublic.GetAuthorizationTokenInput) (*ecrpublic.GetAuthorizationTokenOutput, error)
-}
-
 type ecrPlugin struct {
-	ecr       ECR
-	ecrPublic ECRPublic
+	ecr ECR
 }
 
-func defaultECRProvider(region string) (*ecr.ECR, error) {
+func defaultECRProvider(region string, registryID string) (*ecr.ECR, error) {
 	sess, err := session.NewSessionWithOptions(session.Options{
 		Config:            aws.Config{Region: aws.String(region)},
 		SharedConfigState: session.SharedConfigEnable,
@@ -71,107 +59,40 @@ func defaultECRProvider(region string) (*ecr.ECR, error) {
 	return ecr.New(sess), nil
 }
 
-func publicECRProvider() (*ecrpublic.ECRPublic, error) {
-	// ECR public registries are only in one region and only accessible from regions
-	// in the "aws" partition.
-	sess, err := session.NewSessionWithOptions(session.Options{
-		Config:            aws.Config{Region: aws.String(ecrPublicRegion)},
-		SharedConfigState: session.SharedConfigEnable,
-	})
-	if err != nil {
-		return nil, err
-	}
-
-	return ecrpublic.New(sess), nil
-}
-
-type credsData struct {
-	authToken *string
-	expiresAt *time.Time
-}
-
-func (e *ecrPlugin) getPublicCredsData() (*credsData, error) {
-	klog.Infof("Getting creds for public registry")
-	var err error
-
-	if e.ecrPublic == nil {
-		e.ecrPublic, err = publicECRProvider()
-	}
-	if err != nil {
-		return nil, err
-	}
-
-	output, err := e.ecrPublic.GetAuthorizationToken(&ecrpublic.GetAuthorizationTokenInput{})
+func (e *ecrPlugin) GetCredentials(ctx context.Context, image string, args []string) (*v1.CredentialProviderResponse, error) {
+	registryID, region, registry, err := parseRepoURL(image)
 	if err != nil {
 		return nil, err
 	}
 
-	if output == nil {
-		return nil, errors.New("response output from ECR was nil")
-	}
-
-	if output.AuthorizationData == nil {
-		return nil, errors.New("authorization data was empty")
-	}
-
-	return &credsData{
-		authToken: output.AuthorizationData.AuthorizationToken,
-		expiresAt: output.AuthorizationData.ExpiresAt,
-	}, nil
-}
-
-func (e *ecrPlugin) getPrivateCredsData(imageHost string, image string) (*credsData, error) {
-	klog.Infof("Getting creds for private image %s", image)
-	region, err := parseRegionFromECRPrivateHost(imageHost)
-	if err != nil {
-		return nil, err
-	}
 	if e.ecr == nil {
-		e.ecr, err = defaultECRProvider(region)
+		e.ecr, err = defaultECRProvider(region, registryID)
 		if err != nil {
 			return nil, err
 		}
 	}
-	output, err := e.ecr.GetAuthorizationToken(&ecr.GetAuthorizationTokenInput{})
+
+	output, err := e.ecr.GetAuthorizationToken(&ecr.GetAuthorizationTokenInput{
+		RegistryIds: []*string{aws.String(registryID)},
+	})
 	if err != nil {
 		return nil, err
 	}
+
 	if output == nil {
 		return nil, errors.New("response output from ECR was nil")
 	}
+
 	if len(output.AuthorizationData) == 0 {
 		return nil, errors.New("authorization data was empty")
 	}
-	return &credsData{
-		authToken: output.AuthorizationData[0].AuthorizationToken,
-		expiresAt: output.AuthorizationData[0].ExpiresAt,
-	}, nil
-}
-
-func (e *ecrPlugin) GetCredentials(ctx context.Context, image string, args []string) (*v1.CredentialProviderResponse, error) {
-	var creds *credsData
-	var err error
-
-	imageHost, err := parseHostFromImageReference(image)
-	if err != nil {
-		return nil, err
-	}
-
-	if imageHost == ecrPublicHost {
-		creds, err = e.getPublicCredsData()
-	} else {
-		creds, err = e.getPrivateCredsData(imageHost, image)
-	}
-
-	if err != nil {
-		return nil, err
-	}
 
-	if creds.authToken == nil {
+	data := output.AuthorizationData[0]
+	if data.AuthorizationToken == nil {
 		return nil, errors.New("authorization token in response was nil")
 	}
 
-	decodedToken, err := base64.StdEncoding.DecodeString(aws.StringValue(creds.authToken))
+	decodedToken, err := base64.StdEncoding.DecodeString(aws.StringValue(data.AuthorizationToken))
 	if err != nil {
 		return nil, err
 	}
@@ -181,13 +102,13 @@ func (e *ecrPlugin) GetCredentials(ctx context.Context, image string, args []str
 		return nil, errors.New("error parsing username and password from authorization token")
 	}
 
-	cacheDuration := getCacheDuration(creds.expiresAt)
+	cacheDuration := getCacheDuration(data.ExpiresAt)
 
 	return &v1.CredentialProviderResponse{
 		CacheKeyType:  v1.RegistryPluginCacheKeyType,
 		CacheDuration: cacheDuration,
 		Auth: map[string]v1.AuthConfig{
-			imageHost: {
+			registry: {
 				Username: parts[0],
 				Password: parts[1],
 			},
@@ -214,50 +135,30 @@ func getCacheDuration(expiresAt *time.Time) *metav1.Duration {
 	return cacheDuration
 }
 
-// parseHostFromImageReference parses the hostname from an image reference
-func parseHostFromImageReference(image string) (string, error) {
-	// a URL needs a scheme to be parsed correctly
-	if !strings.Contains(image, "://") {
+// parseRepoURL parses and splits the registry URL
+// returns (registryID, region, registry).
+// <registryID>.dkr.ecr(-fips).<region>.amazonaws.com(.cn)
+func parseRepoURL(image string) (string, string, string, error) {
+	if !strings.Contains(image, "https://") {
 		image = "https://" + image
 	}
 	parsed, err := url.Parse(image)
 	if err != nil {
-		return "", fmt.Errorf("error parsing image reference %s: %v", image, err)
+		return "", "", "", fmt.Errorf("error parsing image %s: %v", image, err)
 	}
-	return parsed.Hostname(), nil
-}
 
-func parseRegionFromECRPrivateHost(host string) (string, error) {
-	splitHost := ecrPrivateHostPattern.FindStringSubmatch(host)
-	if len(splitHost) != 5 {
-		return "", fmt.Errorf("invalid private ECR host: %s", host)
+	splitURL := ecrPattern.FindStringSubmatch(parsed.Hostname())
+	if len(splitURL) < 4 {
+		return "", "", "", fmt.Errorf("%s is not a valid ECR repository URL", parsed.Hostname())
 	}
-	return splitHost[3], nil
+
+	return splitURL[1], splitURL[3], parsed.Hostname(), nil
 }
 
 func main() {
-	logs.InitLogs()
-	defer logs.FlushLogs()
-
-	if err := newCredentialProviderCommand().Execute(); err != nil {
+	p := NewCredentialProvider(&ecrPlugin{})
+	if err := p.Run(context.TODO()); err != nil {
+		klog.Errorf("Error running credential provider plugin: %v", err)
 		os.Exit(1)
 	}
 }
-
-var gitVersion string
-
-func newCredentialProviderCommand() *cobra.Command {
-	cmd := &cobra.Command{
-		Use:     "ecr-credential-provider",
-		Short:   "ECR credential provider for kubelet",
-		Version: gitVersion,
-		Run: func(cmd *cobra.Command, args []string) {
-			p := NewCredentialProvider(&ecrPlugin{})
-			if err := p.Run(context.TODO()); err != nil {
-				klog.Errorf("Error running credential provider plugin: %v", err)
-				os.Exit(1)
-			}
-		},
-	}
-	return cmd
-}