chore: Implement a retry in aws auth command (#9618)

Signed-off-by: Leonardo Luz Almeida <[email protected]>

cmd/argocd-k8s-auth/commands/aws.go

@@ -1,6 +1,7 @@
 package commands
 
 import (
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
@@ -40,16 +41,7 @@ func newAWSCommand() *cobra.Command {
 	var command = &cobra.Command{
 		Use: "aws",
 		Run: func(c *cobra.Command, args []string) {
-			sess, err := session.NewSession()
-			errors.CheckError(err)
-			stsAPI := sts.New(sess)
-			if roleARN != "" {
-				creds := stscreds.NewCredentials(sess, roleARN)
-				stsAPI = sts.New(sess, &aws.Config{Credentials: creds})
-			}
-			request, _ := stsAPI.GetCallerIdentityRequest(&sts.GetCallerIdentityInput{})
-			request.HTTPRequest.Header.Add(clusterIDHeader, clusterName)
-			presignedURLString, err := request.Presign(requestPresignParam)
+			presignedURLString, err := getSignedRequestWithRetry(time.Minute, clusterName, roleARN)
 			errors.CheckError(err)
 			token := v1Prefix + base64.RawURLEncoding.EncodeToString([]byte(presignedURLString))
 			// Set token expiration to 1 minute before the presigned URL expires for some cushion
@@ -62,6 +54,41 @@ func newAWSCommand() *cobra.Command {
 	return command
 }
 
+func getSignedRequestWithRetry(timeout time.Duration, clusterName, roleARN string) (string, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+	for {
+		signed, err := getSignedRequest(clusterName, roleARN)
+		if err == nil {
+			return signed, nil
+		}
+		select {
+		case <-ctx.Done():
+			return "", fmt.Errorf("timeout while trying to get signed aws request: last error: %s", err)
+		case <-time.After(5 * time.Second):
+		}
+	}
+}
+
+func getSignedRequest(clusterName, roleARN string) (string, error) {
+	sess, err := session.NewSession()
+	if err != nil {
+		return "", fmt.Errorf("error creating new AWS session: %s", err)
+	}
+	stsAPI := sts.New(sess)
+	if roleARN != "" {
+		creds := stscreds.NewCredentials(sess, roleARN)
+		stsAPI = sts.New(sess, &aws.Config{Credentials: creds})
+	}
+	request, _ := stsAPI.GetCallerIdentityRequest(&sts.GetCallerIdentityInput{})
+	request.HTTPRequest.Header.Add(clusterIDHeader, clusterName)
+	signed, err := request.Presign(requestPresignParam)
+	if err != nil {
+		return "", fmt.Errorf("error presigning AWS request: %s", err)
+	}
+	return signed, nil
+}
+
 func formatJSON(token string, expiration time.Time) string {
 	expirationTimestamp := metav1.NewTime(expiration)
 	execInput := &clientauthv1beta1.ExecCredential{