feat(auth): request necessary scopes only (#4916)

* feat(auth): request AWS account scopes only when necessary

Applies to both IdC and BuilderID
- Don't request if we are signing into Amazon Q
- Request if we are in toolkit re-using a connection from Amazon Q
  - via login page
  - via quick pick menu
- Don't invalidate the connection in Amazon Q if we cancel adding scopes in toolkit.

TODO:
- Cancelling adding scopes in Q still invalidates the connection in toolkit.

* remove invalidate option from addScopes, fix telemetry

* restore original scopes if re-using a connection in Q fails

* allow reauth without invalidation

* fix tests

* re-add and update test

* docstrings

Author: hayemaxi

packages/core/src/auth/auth.ts

@@ -92,8 +92,6 @@ interface AuthService {
     /**
      * Replaces the profile for a connection with a new one.
      *
-     * This will invalidate the connection, potentially requiring a re-authentication.
-     *
      * **IAM connections are not implemented**
      */
     updateConnection(connection: Pick<Connection, 'id'>, profile: Profile): Promise<Connection>
@@ -166,9 +164,10 @@ 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> {
+    public async reauthenticate({ id }: Pick<SsoConnection, 'id'>, invalidate?: boolean): Promise<SsoConnection>
+    public async reauthenticate({ id }: Pick<IamConnection, 'id'>, invalidate?: boolean): Promise<IamConnection>
+    public async reauthenticate({ id }: Pick<Connection, 'id'>, invalidate?: boolean): Promise<Connection> {
+        const shouldInvalidate = invalidate ?? true
         const profile = this.store.getProfileOrThrow(id)
         if (profile.type === 'sso') {
             const provider = this.getSsoTokenProvider(id, profile)
@@ -177,13 +176,13 @@ export class Auth implements AuthService, ConnectionManager {
             // so we need to set it here.
             await telemetry.aws_loginWithBrowser.run(async span => {
                 span.record({ isReAuth: true, credentialStartUrl: profile.startUrl })
-                await this.authenticate(id, () => provider.createToken())
+                await this.authenticate(id, () => provider.createToken(), shouldInvalidate)
             })
 
             return this.getSsoConnection(id, profile)
         } else {
             const provider = await this.getCredentialsProvider(id, profile)
-            await this.authenticate(id, () => this.createCachedCredentials(provider))
+            await this.authenticate(id, () => this.createCachedCredentials(provider), shouldInvalidate)
 
             return this.getIamConnection(id, profile)
         }
@@ -383,14 +382,24 @@ export class Auth implements AuthService, ConnectionManager {
         await this.validateConnection(connection.id, profile)
     }
 
-    public async updateConnection(connection: Pick<SsoConnection, 'id'>, profile: SsoProfile): Promise<SsoConnection>
-    public async updateConnection(connection: Pick<Connection, 'id'>, profile: Profile): Promise<Connection> {
+    public async updateConnection(
+        connection: Pick<SsoConnection, 'id'>,
+        profile: SsoProfile,
+        invalidate?: boolean
+    ): Promise<SsoConnection>
+    public async updateConnection(
+        connection: Pick<Connection, 'id'>,
+        profile: Profile,
+        invalidate?: boolean
+    ): Promise<Connection> {
         getLogger().info(`auth: Updating connection ${connection.id}`)
         if (profile.type === 'iam') {
             throw new Error('Updating IAM connections is not supported')
         }
 
-        await this.invalidateConnection(connection.id, { skipGlobalLogout: true })
+        if (invalidate ?? false) {
+            await this.invalidateConnection(connection.id, { skipGlobalLogout: true })
+        }
 
         const newProfile = await this.store.updateProfile(connection.id, profile)
         const updatedConn = this.getSsoConnection(connection.id, newProfile as StoredProfile<SsoProfile>)
@@ -554,15 +563,17 @@ export class Auth implements AuthService, ConnectionManager {
             }
         }
 
-        return runCheck().catch(err => this.handleSsoTokenError(id, err))
+        return runCheck().catch(async err => {
+            await this.handleSsoTokenError(id, err) // may throw without setting state to invalid - this is intended.
+            await this.updateConnectionState(id, 'invalid')
+        })
     }
 
     private async handleSsoTokenError(id: Connection['id'], err: unknown) {
         this.throwOnNetworkError(err)
 
         this.#validationErrors.set(id, UnknownError.cast(err))
         getLogger().info(`auth: Handling validation error of connection: ${id}`)
-        return this.updateConnectionState(id, 'invalid')
     }
 
     private async getConnectionFromStoreEntry([id, profile]: readonly [Connection['id'], StoredProfile<Profile>]) {
@@ -699,7 +710,8 @@ export class Auth implements AuthService, ConnectionManager {
     }
 
     private readonly authenticate = keyedDebounce(this._authenticate.bind(this))
-    private async _authenticate<T>(id: Connection['id'], callback: () => Promise<T>): Promise<T> {
+    private async _authenticate<T>(id: Connection['id'], callback: () => Promise<T>, invalidate?: boolean): Promise<T> {
+        const originalState = this.getConnectionState({ id }) ?? 'unauthenticated'
         await this.updateConnectionState(id, 'authenticating')
 
         try {
@@ -709,6 +721,7 @@ export class Auth implements AuthService, ConnectionManager {
             return result
         } catch (err) {
             await this.handleSsoTokenError(id, err)
+            await this.updateConnectionState(id, invalidate ? 'invalid' : originalState)
             throw err
         }
     }

packages/core/src/auth/secondaryAuth.ts

@@ -218,30 +218,7 @@ export class SecondaryAuth<T extends Connection = Connection> {
     }
 
     public async addScopes(conn: T & SsoConnection, extraScopes: string[]) {
-        const oldScopes = conn.scopes ?? []
-        const newScopes = Array.from(new Set([...oldScopes, ...extraScopes]))
-
-        const updateConnectionScopes = (scopes: string[]) => {
-            return this.auth.updateConnection(conn, {
-                type: 'sso',
-                scopes,
-                startUrl: conn.startUrl,
-                ssoRegion: conn.ssoRegion,
-            })
-        }
-
-        const updatedConn = await updateConnectionScopes(newScopes)
-
-        try {
-            return await this.auth.reauthenticate(updatedConn)
-        } catch (e) {
-            // We updated the connection scopes pre-emptively, but if there is some issue (e.g. user cancels,
-            // InvalidGrantException, etc), then we need to revert to the old connection scopes. Otherwise,
-            // this could soft-lock users into a broken connection that cannot be re-authenticated without
-            // first deleting the connection.
-            await updateConnectionScopes(oldScopes)
-            throw e
-        }
+        return await addScopes(conn, extraScopes, this.auth)
     }
 
     // Used to lazily restore persisted connections.
@@ -288,3 +265,43 @@ export class SecondaryAuth<T extends Connection = Connection> {
         }
     }
 }
+
+/**
+ * Used to add scopes to a connection, usually when re-using a connection across extensions.
+ * It does not invalidate or otherwise change the state of the connection, but it does
+ * trigger listeners for connection updates.
+ *
+ * How connections and scopes currently work for both this quickpick and the common Login page:
+ * - Don't request AWS scopes if we are signing into Amazon Q
+ * - Request AWS scopes for explorer sign in. Request AWS + CodeCatalyst scopes for CC sign in.
+ * - Request scope difference if re-using a connection. Cancelling or otherwise failing to get the new scopes does NOT invalidate the connection.
+ * - Adding scopes updates the connection profile, but does not change its state.
+ *
+ * Note: This should exist in connection.ts or utils.ts, but due to circular dependencies, it must go here.
+ */
+export async function addScopes(conn: SsoConnection, extraScopes: string[], auth = Auth.instance) {
+    const oldScopes = conn.scopes ?? []
+    const newScopes = Array.from(new Set([...oldScopes, ...extraScopes]))
+
+    const updateConnectionScopes = (scopes: string[]) => {
+        return auth.updateConnection(conn, {
+            type: 'sso',
+            scopes,
+            startUrl: conn.startUrl,
+            ssoRegion: conn.ssoRegion,
+        })
+    }
+
+    const updatedConn = await updateConnectionScopes(newScopes)
+
+    try {
+        return await auth.reauthenticate(updatedConn, false)
+    } catch (e) {
+        // We updated the connection scopes pre-emptively, but if there is some issue (e.g. user cancels,
+        // InvalidGrantException, etc), then we need to revert to the old connection scopes. Otherwise,
+        // this could soft-lock users into a broken connection that cannot be re-authenticated without
+        // first deleting the connection.
+        await updateConnectionScopes(oldScopes)
+        throw e
+    }
+}

packages/core/src/auth/utils.ts

@@ -20,8 +20,8 @@ import { TreeNode } from '../shared/treeview/resourceTreeDataProvider'
 import { createInputBox } from '../shared/ui/inputPrompter'
 import { telemetry } from '../shared/telemetry/telemetry'
 import { createCommonButtons, createExitButton, createHelpButton, createRefreshButton } from '../shared/ui/buttons'
-import { getIdeProperties, isCloud9 } from '../shared/extensionUtilities'
-import { getDependentAuths } from './secondaryAuth'
+import { getIdeProperties, isAmazonQ, isCloud9 } from '../shared/extensionUtilities'
+import { addScopes, getDependentAuths } from './secondaryAuth'
 import { DevSettings } from '../shared/settings'
 import { createRegionPrompter } from '../shared/ui/common/region'
 import { saveProfileToCredentials } from './credentials/sharedCredentials'
@@ -40,6 +40,7 @@ import {
     createSsoProfile,
     hasScopes,
     scopesSsoAccountAccess,
+    isSsoConnection,
 } from './connection'
 import { Commands, placeholder, RegisteredCommand, vscodeComponent } from '../shared/vscode/commands2'
 import { Auth } from './auth'
@@ -91,14 +92,29 @@ export async function promptForConnection(auth: Auth, type?: 'iam' | 'iam-only'
     return resp
 }
 
+/**
+ * Creates the 'Switch Connection' quickpick for the Toolkit.
+ * See {@link addScopes} for details about how scopes are requested for new and existing connections.
+ */
 export async function promptAndUseConnection(...[auth, type]: Parameters<typeof promptForConnection>) {
     return telemetry.aws_setCredentials.run(async span => {
-        const resp = await promptForConnection(auth, type)
-        if (!resp) {
+        let conn = await promptForConnection(auth, type)
+        if (!conn) {
             throw new CancellationError('user')
         }
 
-        await auth.useConnection(resp)
+        // HACK: We assume that if we are toolkit we want AWS account scopes.
+        // TODO: Although, Q shouldn't enter this codepath anyways.
+        // We should deprecate any codepath in Q that may enter this.
+        if (!isAmazonQ() && isSsoConnection(conn) && !hasScopes(conn, scopesSsoAccountAccess)) {
+            try {
+                conn = await addScopes(conn, scopesSsoAccountAccess)
+            } catch (e: any) {
+                throw new ToolkitError(`Failed to use connection${conn.label ? `: ${conn.label}` : '.'}`, e)
+            }
+        }
+
+        await auth.useConnection(conn)
     })
 }
 

packages/core/src/codecatalyst/auth.ts

@@ -356,7 +356,8 @@ export class CodeCatalystAuthenticationProvider {
             await this.promptOnboarding()
         }
 
-        return this.secondaryAuth.useNewConnection(conn)
+        await this.secondaryAuth.useNewConnection(conn)
+        return conn
     }
 
     /**

packages/core/src/codewhisperer/util/authUtil.ts

@@ -20,7 +20,6 @@ import {
     isIamConnection,
     isSsoConnection,
     isBuilderIdConnection,
-    scopesSsoAccountAccess,
     scopesCodeWhispererChat,
     scopesFeatureDev,
     scopesGumby,
@@ -36,7 +35,7 @@ import { showReauthenticateMessage } from '../../shared/utilities/messages'
 import { showAmazonQWalkthroughOnce } from '../../amazonq/onboardingPage/walkthrough'
 
 /** Backwards compatibility for connections w pre-chat scopes */
-export const codeWhispererCoreScopes = [...scopesSsoAccountAccess, ...scopesCodeWhispererCore]
+export const codeWhispererCoreScopes = [...scopesCodeWhispererCore]
 export const codeWhispererChatScopes = [...codeWhispererCoreScopes, ...scopesCodeWhispererChat]
 export const amazonQScopes = [...codeWhispererChatScopes, ...scopesGumby, ...scopesFeatureDev]
 

packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts

@@ -100,16 +100,23 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
                                     // if failed, connection is set to invalid
                                     const oldScopes = conn?.scopes ? conn.scopes : []
                                     const newScopes = Array.from(new Set([...oldScopes, ...amazonQScopes]))
-                                    newConn = await Auth.instance.createConnectionFromApi({
+                                    const payload: AwsConnection = {
                                         type: conn.type,
                                         ssoRegion: conn.ssoRegion,
                                         scopes: newScopes,
                                         startUrl: conn.startUrl,
                                         state: conn.state,
                                         id: conn.id,
                                         label: conn.label,
-                                    })
-                                    await Auth.instance.reauthenticate(newConn)
+                                    }
+                                    newConn = await Auth.instance.createConnectionFromApi(payload)
+                                    try {
+                                        await Auth.instance.reauthenticate(newConn, false)
+                                    } catch (e) {
+                                        // Restore original scopes as to not soft-lock connections.
+                                        await Auth.instance.createConnectionFromApi({ ...payload, scopes: oldScopes })
+                                        throw e
+                                    }
                                     await AuthUtil.instance.secondaryAuth.useNewConnection(newConn)
                                 }
                                 if (!auto) {

packages/core/src/login/webview/vue/toolkit/backend_toolkit.ts

@@ -7,10 +7,18 @@ import * as vscode from 'vscode'
 import { tryAddCredentials } from '../../../../auth/utils'
 import { getLogger } from '../../../../shared/logger'
 import { CommonAuthWebview } from '../backend'
-import { AwsConnection, Connection, createSsoProfile } from '../../../../auth/connection'
+import {
+    AwsConnection,
+    Connection,
+    createSsoProfile,
+    hasScopes,
+    isIdcSsoConnection,
+    scopesSsoAccountAccess,
+} from '../../../../auth/connection'
 import { Auth } from '../../../../auth/auth'
 import { CodeCatalystAuthenticationProvider } from '../../../../codecatalyst/auth'
 import { AuthError, AuthFlowState, TelemetryMetadata } from '../types'
+import { addScopes } from '../../../../auth/secondaryAuth'
 
 export class ToolkitLoginWebview extends CommonAuthWebview {
     public override id: string = 'aws.toolkit.AmazonCommonAuth'
@@ -36,20 +44,26 @@ export class ToolkitLoginWebview extends CommonAuthWebview {
 
         if (this.isCodeCatalystLogin) {
             return this.ssoSetup('startCodeCatalystSSOSetup', async () => {
-                this.storeMetricMetadata({ ...metadata, authEnabledFeatures: 'codecatalyst' })
+                this.storeMetricMetadata({ ...metadata })
+
+                const conn = await this.codeCatalystAuth.connectToEnterpriseSso(startUrl, region)
+
+                this.storeMetricMetadata({ authEnabledFeatures: this.getAuthEnabledFeatures(conn) })
 
-                await this.codeCatalystAuth.connectToEnterpriseSso(startUrl, region)
                 await vscode.commands.executeCommand('setContext', 'aws.explorer.showAuthView', false)
                 await this.showResourceExplorer()
             })
         }
 
         return this.ssoSetup('createIdentityCenterConnection', async () => {
-            this.storeMetricMetadata({ ...metadata, authEnabledFeatures: 'awsExplorer' })
+            this.storeMetricMetadata({ ...metadata })
 
             const ssoProfile = createSsoProfile(startUrl, region)
             const conn = await Auth.instance.createConnection(ssoProfile)
             await Auth.instance.useConnection(conn)
+
+            this.storeMetricMetadata({ authEnabledFeatures: this.getAuthEnabledFeatures(conn) })
+
             await vscode.commands.executeCommand('setContext', 'aws.explorer.showAuthView', false)
             void vscode.window.showInformationMessage('Toolkit: Successfully connected to AWS IAM Identity Center')
             void this.showResourceExplorer()
@@ -126,7 +140,7 @@ export class ToolkitLoginWebview extends CommonAuthWebview {
      */
     async useConnection(connectionId: string, auto: boolean): Promise<AuthError | undefined> {
         return this.ssoSetup('useConnection', async () => {
-            const conn = await Auth.instance.getConnection({ id: connectionId })
+            let conn = await Auth.instance.getConnection({ id: connectionId })
             if (conn === undefined || conn.type !== 'sso') {
                 return
             }
@@ -136,6 +150,9 @@ export class ToolkitLoginWebview extends CommonAuthWebview {
             if (this.isCodeCatalystLogin) {
                 await this.codeCatalystAuth.tryUseConnection(conn)
             } else {
+                if (isIdcSsoConnection(conn) && !hasScopes(conn, scopesSsoAccountAccess)) {
+                    conn = await addScopes(conn, scopesSsoAccountAccess)
+                }
                 await Auth.instance.useConnection({ id: connectionId })
             }
 

packages/core/src/test/credentials/auth.test.ts

@@ -154,11 +154,12 @@ describe('Auth', function () {
             assert.deepStrictEqual(updated.scopes, updatedProfile.scopes)
         })
 
-        it('invalidates the connection', async function () {
+        it('does not change the connection state', async function () {
             const conn = await auth.createConnection(ssoProfile)
+            const originalState = auth.getConnectionState(conn)
             const updated = await auth.updateConnection(conn, updatedProfile)
 
-            assert.strictEqual(auth.getConnectionState(updated), 'invalid')
+            assert.strictEqual(auth.getConnectionState(updated), originalState)
         })
 
         it('fires an event when updating', async function () {