fix: simplify role merge

update tests

Author: dnitsch

cmd/saml.go

@@ -69,7 +69,7 @@ func getSaml(cmd *cobra.Command, args []string) error {
 	if err != nil {
 		return err
 	}
-	allRoles := credentialexchange.InsertRoleIntoChain(role, roleChain)
+	allRoles := credentialexchange.MergeRoleChain(role, roleChain, isSso)
 	conf := credentialexchange.CredentialConfig{
 		ProviderUrl:  providerUrl,
 		PrincipalArn: principalArn,

cmd/specific.go

@@ -60,7 +60,7 @@ func specific(cmd *cobra.Command, args []string) error {
 			StoreInProfile: storeInProfile,
 			Username:       user.Username,
 			Role:           role,
-			RoleChain:      credentialexchange.InsertRoleIntoChain(role, roleChain),
+			RoleChain:      credentialexchange.MergeRoleChain(role, roleChain, false),
 		},
 	}
 

internal/cmdutils/cmdutils.go

@@ -48,6 +48,9 @@ func GetCredsWebUI(ctx context.Context, svc credentialexchange.AuthSamlApi, secr
 	return credentialexchange.SetCredentials(storedCreds, conf)
 }
 
+// refreshAwsSsoCreds uses the temp user credentials returned via AWS SSO,
+// upon successful auth from the IDP.
+// Once credentials are captured they are used in the role assumption process.
 func refreshAwsSsoCreds(ctx context.Context, conf credentialexchange.CredentialConfig, secretStore SecretStorageImpl, svc credentialexchange.AuthSamlApi, webConfig *web.WebConfig) error {
 	webBrowser := web.New(webConfig)
 	capturedCreds, err := webBrowser.GetSSOCredentials(conf)

internal/credentialexchange/credentialexchange.go

@@ -128,6 +128,7 @@ type authWebTokenApi interface {
 	AssumeRoleWithWebIdentity(ctx context.Context, params *sts.AssumeRoleWithWebIdentityInput, optFns ...func(*sts.Options)) (*sts.AssumeRoleWithWebIdentityOutput, error)
 }
 
+// LoginAwsWebToken
 func LoginAwsWebToken(ctx context.Context, username string, svc authWebTokenApi) (*AWSCredentials, error) {
 	// var role string
 	r, exists := os.LookupEnv(AWS_ROLE_ARN)

internal/credentialexchange/helper.go

@@ -40,17 +40,20 @@ func SessionName(username, selfName string) string {
 	return fmt.Sprintf("%s-%s", strings.ReplaceAll(username, `\`, "--"), selfName)
 }
 
-func InsertRoleIntoChain(role string, roleChain []string) []string {
+// MergeRoleChain inserts the main role into the role chain.
+//
+// This is mainly used with AWS SSO flow where
+// the SSO user credentials are used to assume the target role(s).
+func MergeRoleChain(role string, roleChain []string, insertRoleIntoChain bool) []string {
 	// IF role is provided it can be assumed from the WEB_ID credentials
 	// this is to maintain the old implementation
-	rc := []string{}
-	for _, r := range roleChain {
-		if r != role {
-			rc = append(rc, r)
+	if insertRoleIntoChain {
+		if role != "" {
+			return append([]string{role}, roleChain...)
 		}
+		return roleChain
 	}
-
-	return rc
+	return roleChain
 }
 
 func SetCredentials(creds *AWSCredentials, config CredentialConfig) error {
@@ -100,6 +103,8 @@ func returnStdOutAsJson(creds AWSCredentials) error {
 	return nil
 }
 
+// GetWebIdTokenFileContents reads the contents of the `AWS_WEB_IDENTITY_TOKEN_FILE` environment variable.
+// Used only with specific assume
 func GetWebIdTokenFileContents() (string, error) {
 	// var content *string
 	file, exists := os.LookupEnv(WEB_ID_TOKEN_VAR)

internal/credentialexchange/helper_test.go

@@ -106,21 +106,31 @@ func Test_InsertIntoRoleSlice_with(t *testing.T) {
 	ttests := map[string]struct {
 		role      string
 		roleChain []string
+		isSso     bool
 		expect    []string
 	}{
-		"chain empty and role specified": {
-			"role", []string{}, []string{},
+		"chain empty and role specified with insertRoleIntoChain as false": {
+			"role", []string{}, false, []string{},
 		},
-		"chain set and role empty": {
-			"", []string{"rolec1"}, []string{"rolec1"},
+		"chain set and role empty with insertRoleIntoChain as false": {
+			"", []string{"rolec1"}, false, []string{"rolec1"},
 		},
-		"both set and role is always first in list": {
-			"role", []string{"rolec1"}, []string{"rolec1"},
+		"both set and role is always first in list with insertRoleIntoChain as false": {
+			"role", []string{"rolec1"}, false, []string{"rolec1"},
+		},
+		"chain empty and role specified with insertRoleIntoChain as true": {
+			"role", []string{}, true, []string{"role"},
+		},
+		"chain set and role empty with insertRoleIntoChain as true": {
+			"", []string{"rolec1"}, true, []string{"rolec1"},
+		},
+		"both set and role is always first in list with insertRoleIntoChain as true": {
+			"role", []string{"rolec1"}, true, []string{"role", "rolec1"},
 		},
 	}
 	for name, tt := range ttests {
 		t.Run(name, func(t *testing.T) {
-			if got := credentialexchange.InsertRoleIntoChain(tt.role, tt.roleChain); len(got) != len(tt.expect) {
+			if got := credentialexchange.MergeRoleChain(tt.role, tt.roleChain, tt.isSso); len(got) != len(tt.expect) {
 				t.Errorf("expected: %v, got: %v", tt.expect, got)
 			}
 		})