add sts management and mfa

Author: yuxianrz

packages/amazonq/src/lsp/client.ts

@@ -30,6 +30,8 @@ import {
     ResponseError,
     LSPErrorCodes,
     updateConfigurationRequestType,
+    GetMfaCodeParams,
+    GetMfaCodeResult,
 } from '@aws/language-server-runtimes/protocol'
 import {
     AuthUtil,
@@ -56,7 +58,7 @@ import { processUtils } from 'aws-core-vscode/shared'
 import { activate as activateChat } from './chat/activation'
 import { activate as activeInlineChat } from '../inlineChat/activation'
 import { AmazonQResourcePaths } from './lspInstaller'
-import { auth2 } from 'aws-core-vscode/auth'
+import { auth2, getMfaTokenFromUser, getMfaSerialFromUser } from 'aws-core-vscode/auth'
 import { ConfigSection, isValidConfigSection, pushConfigUpdate, toAmazonQLSPLogLevel } from './config'
 import { telemetry } from 'aws-core-vscode/telemetry'
 import { SessionManager } from '../app/inline/sessionManager'
@@ -337,6 +339,24 @@ async function postStartLanguageServer(
         }
     )
 
+    // Handler for when Flare needs to assume a role with MFA code
+    client.onRequest(
+        auth2.notificationTypes.getMfaCode.method,
+        async (params: GetMfaCodeParams): Promise<GetMfaCodeResult> => {
+            if (params.mfaSerial) {
+                globals.globalState.update('recentMfaSerial', { mfaSerial: params.mfaSerial })
+            }
+            const defaultMfaSerial = globals.globalState.tryGet('recentMfaSerial', Object, {
+                mfaSerial: '',
+            }).mfaSerial
+            let mfaSerial = await getMfaSerialFromUser(defaultMfaSerial, params.profileName)
+            mfaSerial = mfaSerial.trim()
+            await globals.globalState.update('recentMfaSerial', { mfaSerial: mfaSerial })
+            const mfaCode = await getMfaTokenFromUser(mfaSerial, params.profileName)
+            return { code: mfaCode ?? '', mfaSerial: mfaSerial ?? '' }
+        }
+    )
+
     const sendProfileToLsp = async () => {
         try {
             const result = await client.sendRequest(updateConfigurationRequestType.method, {

packages/core/src/auth/auth2.ts

@@ -57,7 +57,7 @@ import { LanguageClient } from 'vscode-languageclient'
 import { getLogger } from '../shared/logger/logger'
 import { ToolkitError } from '../shared/errors'
 import { useDeviceFlow } from './sso/ssoAccessTokenProvider'
-import { getCacheDir, getCacheFileWatcher, getFlareCacheFileName } from './sso/cache'
+import { getCacheDir, getCacheFileWatcher, getFlareCacheFileName, getStsCacheDir } from './sso/cache'
 import { VSCODE_EXTENSION_ID } from '../shared/extensions'
 import { IamCredentials } from '@aws/language-server-runtimes-types'
 import globals from '../shared/extensionGlobals'
@@ -75,6 +75,7 @@ export const notificationTypes = {
     getConnectionMetadata: new RequestType<undefined, ConnectionMetadata, Error>(
         getConnectionMetadataRequestType.method
     ),
+    getMfaCode: new RequestType<GetMfaCodeParams, ResponseMessage, Error>(getMfaCodeRequestType.method),
 }
 
 export type AuthState = 'notConnected' | 'connected' | 'expired'
@@ -89,6 +90,8 @@ export type LoginType = (typeof LoginTypes)[keyof typeof LoginTypes]
 
 export type cacheChangedEvent = 'delete' | 'create'
 
+export type stsCacheChangedEvent = 'delete' | 'create'
+
 export type Login = SsoLogin | IamLogin
 
 export type TokenSource = IamIdentityCenterSsoTokenSource | AwsBuilderIdSsoTokenSource
@@ -114,6 +117,10 @@ const IamProfileOptionsDefaults = {
  */
 export class LanguageClientAuth {
     readonly #ssoCacheWatcher = getCacheFileWatcher(getCacheDir(), getFlareCacheFileName(VSCODE_EXTENSION_ID.amazonq))
+    readonly #stsCacheWatcher = getCacheFileWatcher(
+        getStsCacheDir(),
+        getFlareCacheFileName(VSCODE_EXTENSION_ID.amazonq)
+    )
 
     constructor(
         private readonly client: LanguageClient,
@@ -125,6 +132,10 @@ export class LanguageClientAuth {
         return this.#ssoCacheWatcher
     }
 
+    public get stsCacheWatcher() {
+        return this.#stsCacheWatcher
+    }
+
     getSsoToken(
         tokenSource: TokenSource,
         login: boolean = false,
@@ -281,6 +292,11 @@ export class LanguageClientAuth {
         this.cacheWatcher.onDidCreate(() => cacheChangedHandler('create'))
         this.cacheWatcher.onDidDelete(() => cacheChangedHandler('delete'))
     }
+
+    registerStsCacheWatcher(stsCacheChangedHandler: (event: stsCacheChangedEvent) => any) {
+        this.stsCacheWatcher.onDidCreate(() => stsCacheChangedHandler('create'))
+        this.stsCacheWatcher.onDidDelete(() => stsCacheChangedHandler('delete'))
+    }
 }
 
 /**
@@ -357,13 +373,8 @@ export abstract class BaseLogin {
      * Decrypts an encrypted string, removes its quotes, and returns the resulting string
      */
     protected async decrypt(encrypted: string): Promise<string> {
-        try {
-            const decrypted = await jose.compactDecrypt(encrypted, this.lspAuth.encryptionKey)
-            return decrypted.plaintext.toString().replaceAll('"', '')
-        } catch (e) {
-            getLogger().error(`Failed to decrypt: ${encrypted}`)
-            return encrypted
-        }
+        const decrypted = await jose.compactDecrypt(encrypted, this.lspAuth.encryptionKey)
+        return decrypted.plaintext.toString().replaceAll('"', '')
     }
 }
 
@@ -575,14 +586,6 @@ export class IamLogin extends BaseLogin {
      * Restore the connection state and connection details to memory, if they exist.
      */
     async restore() {
-        const sessionData = await this.getProfile()
-        const credentials = sessionData?.profile?.settings
-        if (credentials?.aws_access_key_id && credentials?.aws_secret_access_key) {
-            this._data = {
-                accessKey: credentials.aws_access_key_id,
-                secretKey: credentials.aws_secret_access_key,
-            }
-        }
         try {
             await this._getIamCredential(false)
         } catch (err) {

packages/core/src/auth/sso/cache.ts

@@ -36,7 +36,9 @@ export interface SsoCache {
 }
 
 const defaultCacheDir = () => path.join(fs.getUserHomeDir(), '.aws/sso/cache')
+const defaultStsCacheDir = () => path.join(fs.getUserHomeDir(), '.aws/cli/cache')
 export const getCacheDir = () => DevSettings.instance.get('ssoCacheDirectory', defaultCacheDir())
+export const getStsCacheDir = () => DevSettings.instance.get('stsCacheDirectory', defaultStsCacheDir())
 
 export function getCache(directory = getCacheDir()): SsoCache {
     return {

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

@@ -33,6 +33,7 @@ import { telemetry } from '../../shared/telemetry/telemetry'
 import {
     AuthStateEvent,
     cacheChangedEvent,
+    stsCacheChangedEvent,
     LanguageClientAuth,
     Login,
     SsoLogin,
@@ -116,6 +117,7 @@ export class AuthUtil implements IAuthProvider {
             await this.setVscodeContextProps()
         })
         lspAuth.registerCacheWatcher(async (event: cacheChangedEvent) => await this.cacheChangedHandler(event))
+        lspAuth.registerStsCacheWatcher(async (event: stsCacheChangedEvent) => await this.stsCacheChangedHandler(event))
     }
 
     // Do NOT use this in production code, only used for testing
@@ -148,12 +150,13 @@ export class AuthUtil implements IAuthProvider {
             this.session = new SsoLogin(this.profileName, this.lspAuth, this.eventEmitter)
             await this.session.restore()
             if (!this.isConnected()) {
+                await this.session?.logout()
                 // Try to restore an IAM session
                 this.session = new IamLogin(this.profileName, this.lspAuth, this.eventEmitter)
                 await this.session.restore()
                 if (!this.isConnected()) {
                     // If both fail, reset the session
-                    this.session = undefined
+                    await this.session?.logout()
                 }
             }
         }
@@ -262,10 +265,6 @@ export class AuthUtil implements IAuthProvider {
         return Boolean(this.connection?.startUrl && this.connection?.startUrl !== builderIdStartUrl)
     }
 
-    isIamConnection() {
-        return Boolean(this.connection?.accessKey && this.connection?.secretKey)
-    }
-
     isInternalAmazonUser(): boolean {
         return this.isConnected() && this.connection?.startUrl === internalStartUrl
     }
@@ -365,6 +364,15 @@ export class AuthUtil implements IAuthProvider {
         }
     }
 
+    private async stsCacheChangedHandler(event: stsCacheChangedEvent) {
+        this.logger.debug(`Sts Cache change event received: ${event}`)
+        if (event === 'delete') {
+            await this.logout()
+        } else if (event === 'create') {
+            await this.restore()
+        }
+    }
+
     private async stateChangeHandler(e: AuthStateEvent) {
         if (e.state === 'refreshed') {
             if (this.isSsoSession()) {

packages/core/src/login/webview/vue/backend.ts

@@ -187,7 +187,7 @@ export abstract class CommonAuthWebview extends VueWebview {
     abstract fetchConnections(): Promise<AwsConnection[] | undefined>
 
     async errorNotification(e: AuthError) {
-        await showMessage('error', e.text)
+        showMessage('error', e.text)
     }
 
     abstract quitLoginScreen(): Promise<void>

packages/core/src/test/testAuthUtil.ts

@@ -58,6 +58,7 @@ export async function createTestAuthUtil() {
         updateIamProfile: sinon.stub().resolves(),
         invalidateSsoToken: sinon.stub().resolves(),
         registerCacheWatcher: sinon.stub().resolves(),
+        registerStsCacheWatcher: sinon.stub().resolves(),
         encryptionKey,
     }