From afbccd04ad824383e27477a2b4483d93e2ba8ead Mon Sep 17 00:00:00 2001
From: Aaron Turner <synfinatic@gmail.com>
Date: Fri, 6 May 2022 17:51:19 -0700
Subject: [PATCH] Support multiple SSO Instance sessions for the same instance

This allows you to specify the same AWS SSO Instance multiple times
and them to be handled seperately so you can authenticate multple
times to the same instance.

Useful for people who have multiple user accounts in their SSO provider.

Fixes: #347
---
 CHANGELOG.md            |  2 ++
 docs/config.md          | 21 ++++++++++++++-------
 sso/awssso.go           |  2 ++
 sso/awssso_auth.go      |  2 +-
 sso/awssso_auth_test.go |  3 ++-
 sso/settings.go         |  6 ++++++
 6 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d81140a8..6b23f32a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -18,6 +18,8 @@
  * Added `Profile` to the list of default fields for the `list` command
  * `list` command can now generate a CSV
  * Renamed the `config` command to `config-profiles` to be less confusing
+ * You can now specify the same StartURL in multiple SSOConfig blocks so you
+    can authenticate as different users at the same time.
 
 ## [v1.8.1] - 2022-05-02
 
diff --git a/docs/config.md b/docs/config.md
index e55dfad1..7d8bc886 100644
--- a/docs/config.md
+++ b/docs/config.md
@@ -83,23 +83,30 @@ EnvVarTags:
 
 ## SSOConfig
 
-This is the top level block for your AWS SSO instances.  Typically an organization
-will have a single AWS SSO instance for all of their accounts under a single AWS master
-payer account.  If you have more than one AWS SSO instance, then `Default` will be
-the default unless overridden with `DefaultSSO`.
+This is the top level block for your AWS SSO instances.  Typically an
+organization will have a single AWS SSO instance for all of their accounts
+under a single AWS master payer account.  If you have more than one AWS SSO
+instance, then `Default` will be the default unless overridden with
+`DefaultSSO`.
 
 The `SSOConfig` config block is required.
 
 ### StartUrl
 
-Each AWS SSO instance has a unique start URL hosted by AWS for interacting with your
-SSO provider (Okta/OneLogin/etc).  Should be in the format of `https://xxxxxxx.awsapps.com/start`.
+Each AWS SSO instance start URL hosted by AWS for interacting with your
+SSO provider (Okta/OneLogin/etc).  Should be in the format of
+`https://xxxxxxx.awsapps.com/start`.
+
+**Note:** If you need to authenticate as different users to the same
+AWS SSO Instance, you can specify multiple [SSOConfig](#ssoconfig) blocks
+with the same `StartUrl`.
 
 The `StartUrl` is required.
 
 ### SSORegion
 
-Each AWS SSO instance is configured in a specific AWS region which needs to be set here.
+Each AWS SSO instance is configured in a specific AWS region which needs to be
+set here.
 
 The `SSORegion` is required.
 
diff --git a/sso/awssso.go b/sso/awssso.go
index 20ee7bb9..e70d7b5d 100644
--- a/sso/awssso.go
+++ b/sso/awssso.go
@@ -51,6 +51,7 @@ type SsoAPI interface {
 }
 
 type AWSSSO struct {
+	key            string // key in the settings file that names us
 	sso            SsoAPI
 	ssooidc        SsoOidcAPI
 	store          storage.SecureStorage
@@ -79,6 +80,7 @@ func NewAWSSSO(s *SSOConfig, store *storage.SecureStorage) *AWSSSO {
 	})
 
 	as := AWSSSO{
+		key:            s.key,
 		sso:            ssoSession,
 		ssooidc:        oidcSession,
 		store:          *store,
diff --git a/sso/awssso_auth.go b/sso/awssso_auth.go
index 841ec5c4..746181af 100644
--- a/sso/awssso_auth.go
+++ b/sso/awssso_auth.go
@@ -67,7 +67,7 @@ func (as *AWSSSO) Authenticate(urlAction, browser string) error {
 
 // StoreKey returns the key in the cache for this AWSSSO instance
 func (as *AWSSSO) StoreKey() string {
-	return fmt.Sprintf("%s|%s", as.SsoRegion, as.StartUrl)
+	return as.key
 }
 
 // reauthenticate talks to AWS SSO to generate a new AWS SSO AccessToken
diff --git a/sso/awssso_auth_test.go b/sso/awssso_auth_test.go
index 559853e4..e12187cd 100644
--- a/sso/awssso_auth_test.go
+++ b/sso/awssso_auth_test.go
@@ -92,11 +92,12 @@ func (m *mockSsoOidcAPI) CreateToken(ctx context.Context, params *ssooidc.Create
 
 func TestStoreKey(t *testing.T) {
 	as := &AWSSSO{
+		key:       "atest",
 		SsoRegion: "us-west-1",
 		StartUrl:  "https://testing.awsapps.com/start",
 	}
 
-	assert.Equal(t, "us-west-1|https://testing.awsapps.com/start", as.StoreKey())
+	assert.Equal(t, "atest", as.StoreKey())
 }
 
 func TestAuthenticateSteps(t *testing.T) {
diff --git a/sso/settings.go b/sso/settings.go
index 99a26fdd..d30198ff 100644
--- a/sso/settings.go
+++ b/sso/settings.go
@@ -71,6 +71,7 @@ type Settings struct {
 
 type SSOConfig struct {
 	settings      *Settings              // pointer back up
+	key           string                 // our key in Settings.SSO[]
 	SSORegion     string                 `koanf:"SSORegion" yaml:"SSORegion"`
 	StartUrl      string                 `koanf:"StartUrl" yaml:"StartUrl"`
 	Accounts      map[string]*SSOAccount `koanf:"Accounts" yaml:"Accounts,omitempty"` // key must be a string to avoid parse errors!
@@ -178,6 +179,11 @@ func LoadSettings(configFile, cacheFile string, defaults map[string]interface{},
 
 	s.setOverrides(override)
 
+	// set our SSO names
+	for k, v := range s.SSO {
+		v.key = k
+	}
+
 	if _, ok := s.SSO[s.DefaultSSO]; !ok {
 		// Select our SSO Provider
 		if len(s.SSO) == 0 {