fix: list-roles empty for single-role users (#14)

aws_account.go

@@ -2,10 +2,12 @@ package saml2aws
 
 import (
 	"bytes"
+	"encoding/base64"
 	"fmt"
 	"io"
 	"net/http"
 	"net/url"
+	"strings"
 
 	"github.com/PuerkitoBio/goquery"
 	"github.com/pkg/errors"
@@ -19,7 +21,21 @@ type AWSAccount struct {
 
 // ParseAWSAccounts extract the aws accounts from the saml assertion
 func ParseAWSAccounts(audience string, samlAssertion string) ([]*AWSAccount, error) {
-	res, err := http.PostForm(audience, url.Values{"SAMLResponse": {samlAssertion}})
+	// A user with a single role is redirected to the console instead of being
+	// shown the role chooser, so note any redirect for the fallback below.
+	redirected := false
+	client := &http.Client{
+		CheckRedirect: func(req *http.Request, via []*http.Request) error {
+			// Keep the same cap as Go's default client.
+			if len(via) >= 10 {
+				return errors.New("stopped after 10 redirects")
+			}
+			redirected = true
+			return nil
+		},
+	}
+
+	res, err := client.PostForm(audience, url.Values{"SAMLResponse": {samlAssertion}})
 	if err != nil {
 		return nil, errors.Wrap(err, "error retrieving AWS login form")
 	}
@@ -30,7 +46,74 @@ func ParseAWSAccounts(audience string, samlAssertion string) ([]*AWSAccount, err
 		return nil, errors.Wrap(err, "error retrieving AWS login body")
 	}
 
-	return ExtractAWSAccounts(data)
+	accounts, err := ExtractAWSAccounts(data)
+	if err != nil {
+		return nil, err
+	}
+
+	// Only a successful sign-in can be a single-role redirect; an error page
+	// after a redirect is not.
+	if len(accounts) == 0 && redirected && res.StatusCode < 400 {
+		return accountsFromAssertion(samlAssertion)
+	}
+
+	return accounts, nil
+}
+
+// accountsFromAssertion builds the accounts from the roles in the SAML
+// assertion. It is the fallback for when AWS redirects instead of serving the
+// role chooser. The assertion carries no account alias, so accounts are named
+// by id, as the chooser names accounts that have no alias.
+func accountsFromAssertion(samlAssertion string) ([]*AWSAccount, error) {
+	data, err := base64.StdEncoding.DecodeString(samlAssertion)
+	if err != nil {
+		return nil, errors.Wrap(err, "error decoding saml assertion")
+	}
+
+	roles, err := ExtractAwsRoles(data)
+	if err != nil {
+		return nil, errors.Wrap(err, "error parsing aws roles")
+	}
+
+	awsRoles, err := ParseAWSRoles(roles)
+	if err != nil {
+		return nil, errors.Wrap(err, "error parsing aws roles")
+	}
+
+	accounts := []*AWSAccount{}
+	byID := make(map[string]*AWSAccount)
+	for _, awsRole := range awsRoles {
+		id := accountIDFromRoleARN(awsRole.RoleARN)
+
+		account, ok := byID[id]
+		if !ok {
+			account = &AWSAccount{Name: "Account: " + id}
+			byID[id] = account
+			accounts = append(accounts, account)
+		}
+
+		account.Roles = append(account.Roles, &AWSRole{
+			RoleARN: awsRole.RoleARN,
+			Name:    roleNameFromRoleARN(awsRole.RoleARN),
+		})
+	}
+
+	return accounts, nil
+}
+
+// accountIDFromRoleARN returns the account id from an ARN of the form
+// arn:partition:service:region:account-id:resource
+func accountIDFromRoleARN(roleARN string) string {
+	parts := strings.Split(roleARN, ":")
+	if len(parts) < 5 {
+		return ""
+	}
+	return parts[4]
+}
+
+// roleNameFromRoleARN returns the role name, the last path segment of the ARN
+func roleNameFromRoleARN(roleARN string) string {
+	return roleARN[strings.LastIndex(roleARN, "/")+1:]
 }
 
 // ExtractAWSAccounts extract the accounts from the AWS html page

aws_account_test.go

@@ -1,12 +1,74 @@
 package saml2aws
 
 import (
+	"encoding/base64"
+	"net/http"
+	"net/http/httptest"
 	"os"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
+// singleRoleAssertion returns testdata/assertion.xml, base64 encoded, with the
+// second role removed, so the assertion carries exactly one role.
+func singleRoleAssertion(t *testing.T) string {
+	t.Helper()
+	data, err := os.ReadFile("testdata/assertion.xml")
+	assert.Nil(t, err)
+	xml := strings.Replace(string(data),
+		"\n        <AttributeValue>arn:aws:iam::123123123123:saml-provider/ExampleADFS,arn:aws:iam::123123123123:role/AWS-Admin-CloudOPSNonProd</AttributeValue>", "", 1)
+	assert.NotEqual(t, string(data), xml)
+	return base64.StdEncoding.EncodeToString([]byte(xml))
+}
+
+// A single-role user: AWS answers the sign-in POST with a 302 to the console
+// instead of the role chooser, so the chooser page is never served.
+func TestParseAWSAccountsSingleRoleRedirect(t *testing.T) {
+	assertion := singleRoleAssertion(t)
+
+	mux := http.NewServeMux()
+	mux.HandleFunc("/saml", func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, http.MethodPost, r.Method)
+		assert.Equal(t, assertion, r.FormValue("SAMLResponse"))
+		http.Redirect(w, r, "/console", http.StatusFound)
+	})
+	mux.HandleFunc("/console", func(w http.ResponseWriter, r *http.Request) {
+		_, _ = w.Write([]byte("<html><body>AWS Management Console</body></html>"))
+	})
+	srv := httptest.NewServer(mux)
+	defer srv.Close()
+
+	accounts, err := ParseAWSAccounts(srv.URL+"/saml", assertion)
+	require.NoError(t, err)
+	require.Len(t, accounts, 1)
+
+	account := accounts[0]
+	assert.Equal(t, "Account: 123123123123", account.Name)
+	require.Len(t, account.Roles, 1)
+	assert.Equal(t, "arn:aws:iam::123123123123:role/AWS-Admin-CloudOPSBuild", account.Roles[0].RoleARN)
+	assert.Equal(t, "AWS-Admin-CloudOPSBuild", account.Roles[0].Name)
+}
+
+// Regression guard: a multi-role user gets the chooser page, which is parsed
+// as before.
+func TestParseAWSAccountsChooserPage(t *testing.T) {
+	data, err := os.ReadFile("testdata/saml.html")
+	assert.Nil(t, err)
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_, _ = w.Write(data)
+	}))
+	defer srv.Close()
+
+	accounts, err := ParseAWSAccounts(srv.URL, singleRoleAssertion(t))
+	require.NoError(t, err)
+	require.Len(t, accounts, 2)
+	assert.Equal(t, "Account: account-alias (000000000001)", accounts[0].Name)
+}
+
 func TestExtractAWSAccounts(t *testing.T) {
 	data, err := os.ReadFile("testdata/saml.html")
 	assert.Nil(t, err)
@@ -76,3 +138,47 @@ func TestLocateRole(t *testing.T) {
 
 	assert.Equal(t, "arn:aws:iam::000000000001:role/Development", role.RoleARN)
 }
+
+// A redirect that lands on an error page must not be read as a single-role
+// user: the fallback only applies to a successful sign-in.
+func TestParseAWSAccountsRedirectToErrorPage(t *testing.T) {
+	assertion := singleRoleAssertion(t)
+
+	mux := http.NewServeMux()
+	mux.HandleFunc("/saml", func(w http.ResponseWriter, r *http.Request) {
+		http.Redirect(w, r, "/forbidden", http.StatusFound)
+	})
+	mux.HandleFunc("/forbidden", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusForbidden)
+		_, _ = w.Write([]byte("<html><body>Access denied</body></html>"))
+	})
+	srv := httptest.NewServer(mux)
+	defer srv.Close()
+
+	accounts, err := ParseAWSAccounts(srv.URL+"/saml", assertion)
+	require.NoError(t, err)
+	assert.Empty(t, accounts)
+}
+
+// A redirect loop must stop after Go's default cap of 10 redirects, not
+// follow forever.
+func TestParseAWSAccountsRedirectLoop(t *testing.T) {
+	assertion := singleRoleAssertion(t)
+
+	hits := 0
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		hits++
+		if hits > 50 {
+			// Bound the test if the cap is missing; the assertion below then fails.
+			_, _ = w.Write([]byte("<html></html>"))
+			return
+		}
+		http.Redirect(w, r, "/saml", http.StatusFound)
+	}))
+	defer srv.Close()
+
+	_, err := ParseAWSAccounts(srv.URL+"/saml", assertion)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "stopped after 10 redirects")
+	assert.LessOrEqual(t, hits, 11)
+}

test/e2e/README.md

@@ -92,7 +92,7 @@ driver from npm (`playwright-core@1.47.2` + Node 20) at
 |---|---|
 | `TestLoginWritesCredentials` | `login --skip-prompt` writes the `[e2e]` profile with key, secret and session token; `x_principal_arn` is the assumed-role ARN; expiry is about 1h out; the credentials work: `GetCallerIdentity` returns the assumed role |
 | `TestScriptPrintsLoggedInCredentials` | `script --shell=env` prints the stored credentials |
-| `TestListRolesSingleRole` | `list-roles` shows the role from a one-role assertion. **Known bug, skipped** unless `E2E_KNOWN_BUGS=1` (see below) |
+| `TestListRolesSingleRole` | `list-roles` shows the role from a one-role assertion (fallback to the assertion's roles when sign-in redirects) |
 | `TestWrongPasswordFailsWithErrorElementConfig` | a wrong password fails fast and writes no credentials, with `kc_auth_error_element` set for Keycloak 26 |
 | `TestWrongPasswordFailsDefaultConfig` | the same with the default config. **Known bug, skipped** unless `E2E_KNOWN_BUGS=1` (see below) |
 
@@ -154,7 +154,8 @@ to it.
    but with one valid role AWS answers `302` to the console, so the chooser
    never appears. `ParseAWSAccounts` follows the redirect, finds no
    `fieldset > div.saml-account`, and saml2aws exits 0 with no output.
-   Verified against real AWS. Reproduce with `E2E_KNOWN_BUGS=1`.
+   Verified against real AWS. Fixed: when sign-in redirects and no accounts
+   are parsed, `list-roles` now builds the accounts from the assertion's roles.
 6. `script` validates the IdP account (it needs `--url` and the provider)
    even though it only prints stored credentials.
 

test/e2e/e2e_test.go

@@ -143,13 +143,15 @@ func TestScriptPrintsLoggedInCredentials(t *testing.T) {
 	require.Contains(t, out, "AWS_ACCESS_KEY_ID="+key)
 }
 
-// Known bug: list-roles reads account names from the AWS sign-in role
-// chooser, but with a single role AWS skips the chooser and redirects to the
-// console. saml2aws follows the redirect, finds no accounts and prints
-// nothing, without an error. Set E2E_KNOWN_BUGS=1 to run it.
+// list-roles reads account names from the AWS sign-in role chooser, but with a
+// single role AWS skips the chooser and redirects to the console. Once the
+// redirect finds no accounts, list-roles falls back to the roles in the
+// assertion (fork issue aspratt/saml2aws#14).
 func TestListRolesSingleRole(t *testing.T) {
-	if os.Getenv("E2E_KNOWN_BUGS") == "" {
-		t.Skip("known bug: list-roles prints nothing when the assertion has one role")
+	// list-roles always posts to the real AWS sign-in page, which moto doesn't
+	// stand in for and the offline lane can't reach, so this runs in e2e-aws only.
+	if mustEnv(t, "E2E_TARGET") != "aws" {
+		t.Skip("needs the real AWS sign-in page; runs in the e2e-aws lane only")
 	}
 	home := t.TempDir()
 	out, err := run(t, home, append(append([]string{"list-roles"}, idpArgs(t)...),