From f8bc45756f0d1c194fde21652a8ba3ca272d7794 Mon Sep 17 00:00:00 2001
From: deviantony <anthony.lapenna@portainer.io>
Date: Tue, 25 Jul 2023 15:12:58 +0000
Subject: [PATCH] feat: update AWS auth logic to ECR

---
 edge/aws/aws.go | 40 ++++++++++++++++++++++++++++------------
 1 file changed, 28 insertions(+), 12 deletions(-)

diff --git a/edge/aws/aws.go b/edge/aws/aws.go
index 347e11d5b..eeab515ce 100644
--- a/edge/aws/aws.go
+++ b/edge/aws/aws.go
@@ -2,17 +2,19 @@ package aws
 
 import (
 	"context"
+	"encoding/base64"
 	"errors"
+	"strings"
 
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/credentials"
+	"github.com/aws/aws-sdk-go-v2/service/ecr"
 	iamra "github.com/aws/rolesanywhere-credential-helper/aws_signing_helper"
-	"github.com/awslabs/amazon-ecr-credential-helper/ecr-login/api"
 	"github.com/portainer/agent"
 	"github.com/rs/zerolog/log"
 )
 
-var ErrNoCredentials = errors.New("No credentials found")
+var ErrNoCredentials = errors.New("no credentials found")
 
 func DoAWSIAMRolesAnywhereAuthAndGetECRCredentials(serverURL string, awsConfig *agent.AWSConfig) (*agent.RegistryCredentials, error) {
 	if serverURL == "" || awsConfig == nil {
@@ -20,6 +22,7 @@ func DoAWSIAMRolesAnywhereAuthAndGetECRCredentials(serverURL string, awsConfig *
 			Str("server_url", serverURL).
 			Str("aws configuration region", awsConfig.Region).
 			Msg("incomplete information when using local AWS config for credential lookup")
+
 		return nil, errors.New("invalid ecr configuration")
 	}
 
@@ -28,8 +31,6 @@ func DoAWSIAMRolesAnywhereAuthAndGetECRCredentials(serverURL string, awsConfig *
 		return nil, err
 	}
 
-	factory := api.DefaultClientFactory{}
-
 	cfg, err := config.LoadDefaultConfig(
 		context.TODO(),
 		config.WithRegion(awsConfig.Region),
@@ -40,20 +41,35 @@ func DoAWSIAMRolesAnywhereAuthAndGetECRCredentials(serverURL string, awsConfig *
 		return nil, err
 	}
 
-	client := factory.NewClient(cfg)
+	client := ecr.NewFromConfig(cfg)
 
-	creds, err := client.GetCredentials(serverURL)
+	output, err := client.GetAuthorizationToken(context.TODO(), &ecr.GetAuthorizationTokenInput{})
 	if err != nil {
-		// This might not be an ECR registry
-		// Therefore we deliberately not return an error here so that the upstream logic can fallback to other credential providers
-		log.Warn().Str("server_url", serverURL).Err(err).Msg("unable to retrieve credentials from server")
-		return nil, ErrNoCredentials
+		log.Err(err).Msg("unable to get ECR authorization token")
+		return nil, err
 	}
 
+	if len(output.AuthorizationData) == 0 {
+		log.Err(err).Msg("unable to find ECR authorization token associated with the AWS account")
+		return nil, errors.New("no ECR authorization token associated with AWS account")
+	}
+
+	data := output.AuthorizationData[0]
+
+	token, err := base64.StdEncoding.DecodeString(*data.AuthorizationToken)
+	if err != nil {
+		log.Err(err).Msg("unable to decode ECR authorization token")
+		return nil, err
+	}
+
+	tokenParts := strings.Split(string(token), ":")
+	username := tokenParts[0]
+	password := tokenParts[1]
+
 	return &agent.RegistryCredentials{
 		ServerURL: serverURL,
-		Username:  creds.Username,
-		Secret:    creds.Password,
+		Username:  username,
+		Secret:    password,
 	}, nil
 }