Merge pull request #21 from mikkeloscar/truncate-roles-with-paths

Fix handling of roles with paths

credentials_getter.go

@@ -103,20 +103,37 @@ func GetBaseRoleARN(sess *session.Session) (string, error) {
 // https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRole.html
 func normalizeRoleARN(roleARN string) (string, error) {
 	parts := strings.Split(roleARN, "/")
-	if len(parts) != 2 {
+	if len(parts) < 2 {
 		return "", fmt.Errorf("invalid roleARN: %s", roleARN)
 	}
 
+	remainingChars := roleSessionNameMaxSize
+
 	accountID := strings.TrimPrefix(parts[0], arnPrefix)
 	accountID = strings.TrimSuffix(accountID, roleARNSuffix)
 
-	roleName := strings.Replace(parts[1], ":", "_", -1)
-	roleName = strings.Replace(roleName, "/", ".", -1)
+	remainingChars -= len(accountID)
+
+	return accountID + normalizePath(parts[1:], remainingChars), nil
+}
+
+// normalizePath normalizes the path levels into a roleSession valid string.
+// The last level always gets as many chars as possible leaving only a minimum
+// of one char for each of the other levels.
+// e.g. given the levels: ["aaaaa", "bbbbb", "ccccccc"], and remaining "12" it
+// would be reduced to the string: ".a.b.ccccccc"
+func normalizePath(levels []string, remaining int) string {
+	if len(levels) == 0 {
+		return ""
+	}
 
-	roleNameMaxSize := roleSessionNameMaxSize - 1 - len(accountID)
+	last := levels[len(levels)-1]
+	last = strings.Replace(last, ":", "_", -1)
+	otherLevels := len(levels[:len(levels)-1])
+	maxName := remaining - (otherLevels * 2) - 1
 
-	if len(roleName) > roleNameMaxSize {
-		roleName = roleName[:roleNameMaxSize]
+	if len(last) > maxName {
+		last = last[:maxName]
 	}
-	return accountID + "." + roleName, nil
+	return normalizePath(levels[:len(levels)-1], remaining-len(last)-1) + "." + last
 }

credentials_getter_test.go

@@ -60,17 +60,37 @@ func TestGet(t *testing.T) {
 // 	baseRole, err := GetBaseRoleARN(sess)
 // }
 
-func TestNormalizeRoleARN(t *testing.T) {
-	roleARN := "arn:aws:iam::012345678910:role/role-name"
-	expectedARN := "012345678910.role-name"
-	normalized, err := normalizeRoleARN(roleARN)
-	require.NoError(t, err)
-	require.Equal(t, expectedARN, normalized)
-
-	// truncate long role names
-	roleARN = "arn:aws:iam::012345678910:role/role-name-very-very-very-very-very-very-very-very-long"
-	expectedARN = "012345678910.role-name-very-very-very-very-very-very-very-very-l"
-	normalized, err = normalizeRoleARN(roleARN)
-	require.NoError(t, err)
-	require.Equal(t, expectedARN, normalized)
+func TestNormalizeRoleARN(tt *testing.T) {
+	for _, tc := range []struct {
+		msg         string
+		roleARN     string
+		expectedARN string
+	}{
+		{
+			msg:         "simple role",
+			roleARN:     "arn:aws:iam::012345678910:role/role-name",
+			expectedARN: "012345678910.role-name",
+		},
+		{
+			msg:         "truncate long role names",
+			roleARN:     "arn:aws:iam::012345678910:role/role-name-very-very-very-very-very-very-very-very-long",
+			expectedARN: "012345678910.role-name-very-very-very-very-very-very-very-very-l",
+		},
+		{
+			msg:         "role name with path",
+			roleARN:     "arn:aws:iam::012345678910:role/path-name/role-name",
+			expectedARN: "012345678910.path-name.role-name",
+		},
+		{
+			msg:         "truncate path for long role names",
+			roleARN:     "arn:aws:iam::012345678910:role/aaaaa/bbbbb/ccccccccccccccccccccccccccccccccccccc-role-name",
+			expectedARN: "012345678910.a.b.ccccccccccccccccccccccccccccccccccccc-role-name",
+		},
+	} {
+		tt.Run(tc.msg, func(t *testing.T) {
+			normalized, err := normalizeRoleARN(tc.roleARN)
+			require.NoError(t, err)
+			require.Equal(t, tc.expectedARN, normalized)
+		})
+	}
 }