feat(auth): several UX improvements for browser login flows (#3198)

## Problems
* It's not always obvious that the extension is waiting for the user to login
* Some CodeCatalyst logic doesn't check for 'valid' connections prior to using them, creating noise
* The "Copy Code" modal is sometimes shown twice when refreshing expired builder ID connections

## Solutions
* Restrict connections to 1 re-auth flow at a time
* Add some extra checks in CC logic to reduce noise
* Show a progress notification while the user is connecting via the browser

Author: JadenSimon

.changes/next-release/Bug Fix-b901b228-a954-413d-a9c5-b47c61ec5799.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "auth: \"Copy Code\" modal is shown twice when refreshing expired IAM Identity Center connections"
+}

.changes/next-release/Feature-6c58a802-1a1a-4819-bb6b-b95e66145a4c.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "auth: verification codes for browser logins are now shown in a notification after opening the login URL"
+}

src/codecatalyst/explorer.ts

@@ -148,9 +148,11 @@ export class CodeCatalystRootNode implements RootNode {
         try {
             await this.authProvider.restore()
             const conn = this.authProvider.activeConnection
-            const client = conn !== undefined ? await createClient(conn) : undefined
+            if (conn !== undefined && this.authProvider.auth.getConnectionState(conn) === 'valid') {
+                const client = await createClient(conn)
 
-            return client !== undefined ? await getConnectedDevEnv(client) : undefined
+                return await getConnectedDevEnv(client)
+            }
         } catch (err) {
             getLogger().warn(`codecatalyst: failed to get Dev Environment: ${UnknownError.cast(err).message}`)
         }

src/codecatalyst/reconnect.ts

@@ -25,7 +25,7 @@ const maxReconnectTime = 10 * 60 * 1000
 export function watchRestartingDevEnvs(ctx: ExtContext, authProvider: CodeCatalystAuthenticationProvider) {
     let restartHandled = false
     authProvider.onDidChangeActiveConnection(async conn => {
-        if (restartHandled || conn === undefined) {
+        if (restartHandled || conn === undefined || authProvider.auth.getConnectionState(conn) !== 'valid') {
             return
         }
 

src/credentials/auth.ts

@@ -331,6 +331,8 @@ export class Auth implements AuthService, ConnectionManager {
         }
     }
 
+    public async reauthenticate({ id }: Pick<SsoConnection, 'id'>): Promise<SsoConnection>
+    public async reauthenticate({ id }: Pick<IamConnection, 'id'>): Promise<IamConnection>
     public async reauthenticate({ id }: Pick<Connection, 'id'>): Promise<Connection> {
         const profile = this.store.getProfileOrThrow(id)
         if (profile.type === 'sso') {
@@ -347,7 +349,7 @@ export class Auth implements AuthService, ConnectionManager {
     }
 
     public async useConnection({ id }: Pick<SsoConnection, 'id'>): Promise<SsoConnection>
-    public async useConnection({ id }: Pick<Connection, 'id'>): Promise<Connection>
+    public async useConnection({ id }: Pick<IamConnection, 'id'>): Promise<IamConnection>
     public async useConnection({ id }: Pick<Connection, 'id'>): Promise<Connection> {
         const profile = this.store.getProfile(id)
         if (profile === undefined) {
@@ -569,7 +571,7 @@ export class Auth implements AuthService, ConnectionManager {
             type: 'iam',
             state: profile.metadata.connectionState,
             label: profile.metadata.label ?? id,
-            getCredentials: () => this.debouncedGetCredentials(id, provider),
+            getCredentials: () => this.getCredentials(id, provider),
         }
     }
 
@@ -591,11 +593,12 @@ export class Auth implements AuthService, ConnectionManager {
             startUrl: profile.startUrl,
             state: profile.metadata.connectionState,
             label: profile.metadata?.label ?? label,
-            getToken: () => this.debouncedGetToken(id, provider),
+            getToken: () => this.getToken(id, provider),
         }
     }
 
-    private async authenticate<T>(id: Connection['id'], callback: () => Promise<T>): Promise<T> {
+    private readonly authenticate = keyedDebounce(this._authenticate.bind(this))
+    private async _authenticate<T>(id: Connection['id'], callback: () => Promise<T>): Promise<T> {
         await this.updateConnectionState(id, 'authenticating')
 
         try {
@@ -625,15 +628,15 @@ export class Auth implements AuthService, ConnectionManager {
         }
     }
 
-    private readonly debouncedGetToken = keyedDebounce(Auth.prototype.getToken.bind(this))
-    private async getToken(id: Connection['id'], provider: SsoAccessTokenProvider): Promise<SsoToken> {
+    private readonly getToken = keyedDebounce(this._getToken.bind(this))
+    private async _getToken(id: Connection['id'], provider: SsoAccessTokenProvider): Promise<SsoToken> {
         const token = await provider.getToken()
 
         return token ?? this.handleInvalidCredentials(id, () => provider.createToken())
     }
 
-    private readonly debouncedGetCredentials = keyedDebounce(Auth.prototype.getCredentials.bind(this))
-    private async getCredentials(id: Connection['id'], provider: CredentialsProvider): Promise<Credentials> {
+    private readonly getCredentials = keyedDebounce(this._getCredentials.bind(this))
+    private async _getCredentials(id: Connection['id'], provider: CredentialsProvider): Promise<Credentials> {
         const credentials = await this.getCachedCredentials(provider)
         if (credentials !== undefined) {
             return credentials

src/credentials/providers/sharedCredentialsProvider.ts

@@ -220,7 +220,13 @@ export class SharedCredentialsProvider implements CredentialsProvider {
         const loadedCreds = await this.patchSourceCredentials()
 
         const provider = chain(this.makeCredentialsProvider(loadedCreds))
-        return resolveProviderWithCancel(this.profileName, provider())
+
+        // SSO profiles already show a notification, no need to show another
+        if (isSsoProfile(this.profile)) {
+            return provider()
+        } else {
+            return resolveProviderWithCancel(this.profileName, provider())
+        }
     }
 
     /**

src/credentials/sso/clients.ts

@@ -16,29 +16,19 @@ import {
     SSOServiceException,
 } from '@aws-sdk/client-sso'
 import {
-    AuthorizationPendingException,
     CreateTokenRequest,
     RegisterClientRequest,
-    SlowDownException,
     SSOOIDC,
     StartDeviceAuthorizationRequest,
 } from '@aws-sdk/client-sso-oidc'
 import { AsyncCollection } from '../../shared/utilities/asyncCollection'
 import { pageableToCollection } from '../../shared/utilities/collectionUtils'
-import {
-    assertHasProps,
-    hasStringProps,
-    isNonNullable,
-    RequiredProps,
-    selectFrom,
-} from '../../shared/utilities/tsUtils'
+import { assertHasProps, isNonNullable, RequiredProps, selectFrom } from '../../shared/utilities/tsUtils'
 import { getLogger } from '../../shared/logger'
 import { SsoAccessTokenProvider } from './ssoAccessTokenProvider'
-import { sleep } from '../../shared/utilities/timeoutUtils'
 import { isClientFault } from '../../shared/errors'
 import { DevSettings } from '../../shared/settings'
 
-const backoffDelayMs = 5000
 export class OidcClient {
     public constructor(private readonly client: SSOOIDC, private readonly clock: { Date: typeof Date }) {}
 
@@ -75,28 +65,6 @@ export class OidcClient {
         }
     }
 
-    public async pollForToken(request: CreateTokenRequest, timeout: number, interval = backoffDelayMs) {
-        while (this.clock.Date.now() + interval <= timeout) {
-            try {
-                return await this.createToken(request)
-            } catch (err) {
-                if (!hasStringProps(err, 'name')) {
-                    throw err
-                }
-
-                if (err instanceof SlowDownException) {
-                    interval += backoffDelayMs
-                } else if (!(err instanceof AuthorizationPendingException)) {
-                    throw err
-                }
-            }
-
-            await sleep(interval)
-        }
-
-        throw new Error('Timed-out waiting for authentication token')
-    }
-
     public static create(region: string) {
         return new this(
             new SSOOIDC({

src/credentials/sso/ssoAccessTokenProvider.ts

@@ -3,16 +3,19 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import globals from '../../shared/extensionGlobals'
+import * as nls from 'vscode-nls'
+const localize = nls.loadMessageBundle()
 
-import { SSOOIDCServiceException } from '@aws-sdk/client-sso-oidc'
+import globals from '../../shared/extensionGlobals'
+import * as vscode from 'vscode'
+import { AuthorizationPendingException, SlowDownException, SSOOIDCServiceException } from '@aws-sdk/client-sso-oidc'
 import { openSsoPortalLink, SsoToken, ClientRegistration, isExpired, SsoProfile } from './model'
 import { getCache } from './cache'
-import { hasProps, RequiredProps, selectFrom } from '../../shared/utilities/tsUtils'
-import { CancellationError } from '../../shared/utilities/timeoutUtils'
+import { hasProps, hasStringProps, RequiredProps, selectFrom } from '../../shared/utilities/tsUtils'
+import { CancellationError, sleep } from '../../shared/utilities/timeoutUtils'
 import { OidcClient } from './clients'
 import { loadOr } from '../../shared/utilities/cacheUtils'
-import { isClientFault } from '../../shared/errors'
+import { isClientFault, ToolkitError } from '../../shared/errors'
 
 const clientRegistrationType = 'public'
 const deviceGrantType = 'urn:ietf:params:oauth:grant-type:device_code'
@@ -156,12 +159,7 @@ export class SsoAccessTokenProvider {
             grantType: deviceGrantType,
         }
 
-        const token = await this.oidc.pollForToken(
-            tokenRequest,
-            registration.expiresAt.getTime(),
-            authorization.interval
-        )
-
+        const token = await pollForTokenWithProgress(() => this.oidc.createToken(tokenRequest), authorization)
         return this.formatToken(token, registration)
     }
 
@@ -177,3 +175,56 @@ export class SsoAccessTokenProvider {
         return new this(profile)
     }
 }
+
+const backoffDelayMs = 5000
+async function pollForTokenWithProgress<T>(
+    fn: () => Promise<T>,
+    authorization: Awaited<ReturnType<OidcClient['startDeviceAuthorization']>>,
+    interval = authorization.interval ?? backoffDelayMs
+) {
+    async function poll(token: vscode.CancellationToken) {
+        while (
+            authorization.expiresAt.getTime() - globals.clock.Date.now() > interval &&
+            !token.isCancellationRequested
+        ) {
+            try {
+                return await fn()
+            } catch (err) {
+                if (!hasStringProps(err, 'name')) {
+                    throw err
+                }
+
+                if (err instanceof SlowDownException) {
+                    interval += backoffDelayMs
+                } else if (!(err instanceof AuthorizationPendingException)) {
+                    throw err
+                }
+            }
+
+            await sleep(interval)
+        }
+
+        throw new ToolkitError('Timed-out waiting for browser login flow to complete', {
+            code: 'TimedOut',
+        })
+    }
+
+    return vscode.window.withProgress(
+        {
+            title: localize(
+                'AWS.auth.loginWithBrowser.messageDetail',
+                'Login page opened in browser. When prompted, provide this code: {0}',
+                authorization.userCode
+            ),
+            cancellable: true,
+            location: vscode.ProgressLocation.Notification,
+        },
+        (_, token) =>
+            Promise.race([
+                poll(token),
+                new Promise<never>((_, reject) =>
+                    token.onCancellationRequested(() => reject(new CancellationError('user')))
+                ),
+            ])
+    )
+}

src/test/credentials/auth.test.ts

@@ -34,10 +34,11 @@ describe('Auth', function () {
 
     function createTestTokenProvider() {
         let token: SsoToken | undefined
+        let counter = 0
         const provider = stub(SsoAccessTokenProvider)
         provider.getToken.callsFake(async () => token)
         provider.createToken.callsFake(
-            async () => (token = { accessToken: '123', expiresAt: new Date(Date.now() + 1000000) })
+            async () => (token = { accessToken: String(++counter), expiresAt: new Date(Date.now() + 1000000) })
         )
         provider.invalidate.callsFake(async () => (token = undefined))
 
@@ -185,6 +186,15 @@ describe('Auth', function () {
         assert.strictEqual(auth.activeConnection?.state, 'invalid')
     })
 
+    it('prevents concurrent `reauthenticate` operations on the same connection', async function () {
+        const conn = await setupInvalidSsoConnection(auth, ssoProfile)
+        await Promise.all([auth.reauthenticate(conn), auth.reauthenticate(conn)])
+        const t1 = await conn.getToken()
+        assert.strictEqual(t1.accessToken, '2', 'Only two tokens should have been created')
+        const t3 = await auth.reauthenticate(conn).then(c => c.getToken())
+        assert.notStrictEqual(t1.accessToken, t3.accessToken, 'Access tokens should change after `reauthenticate`')
+    })
+
     describe('SSO Connections', function () {
         async function runExpiredGetTokenFlow(conn: SsoConnection, selection: string | RegExp) {
             const token = conn.getToken()