Move MFA handler to auth2

Author: yuxianrz

packages/amazonq/src/lsp/client.ts

@@ -30,8 +30,6 @@ import {
     ResponseError,
     LSPErrorCodes,
     updateConfigurationRequestType,
-    GetMfaCodeParams,
-    GetMfaCodeResult,
 } from '@aws/language-server-runtimes/protocol'
 import {
     AuthUtil,
@@ -58,7 +56,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, getMfaTokenFromUser, getMfaSerialFromUser } from 'aws-core-vscode/auth'
+import { auth2 } 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'
@@ -339,24 +337,6 @@ 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) {
-                await 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

@@ -50,6 +50,7 @@ import {
     Profile,
     SsoSession,
     GetMfaCodeParams,
+    GetMfaCodeResult,
     getMfaCodeRequestType,
 } from '@aws/language-server-runtimes/protocol'
 import { LanguageClient } from 'vscode-languageclient'
@@ -59,6 +60,8 @@ import { useDeviceFlow } from './sso/ssoAccessTokenProvider'
 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'
+import { getMfaSerialFromUser, getMfaTokenFromUser } from './credentials/utils'
 
 export const notificationTypes = {
     updateIamCredential: new RequestType<UpdateCredentialsParams, ResponseMessage, Error>(
@@ -72,7 +75,6 @@ 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'
@@ -291,6 +293,10 @@ export class LanguageClientAuth {
         this.client.onNotification(stsCredentialChangedRequestType.method, stsCredentialChangedHandler)
     }
 
+    registerGetMfaCodeHandler(getMfaCodeHandler: (params: GetMfaCodeParams) => Promise<GetMfaCodeResult>) {
+        this.client.onRequest(getMfaCodeRequestType.method, getMfaCodeHandler)
+    }
+
     registerCacheWatcher(cacheChangedHandler: (event: cacheChangedEvent) => any) {
         this.cacheWatcher.onDidCreate(() => cacheChangedHandler('create'))
         this.cacheWatcher.onDidDelete(() => cacheChangedHandler('delete'))
@@ -541,6 +547,7 @@ export class IamLogin extends BaseLogin {
         lspAuth.registerStsCredentialChangedHandler((params: StsCredentialChangedParams) =>
             this.stsCredentialChangedHandler(params)
         )
+        lspAuth.registerGetMfaCodeHandler((params: GetMfaCodeParams) => this.getMfaCodeHandler(params))
     }
 
     async login(opts: { accessKey: string; secretKey: string; sessionToken?: string; roleArn?: string }) {
@@ -680,4 +687,18 @@ export class IamLogin extends BaseLogin {
             }
         }
     }
+
+    private async getMfaCodeHandler(params: GetMfaCodeParams): Promise<GetMfaCodeResult> {
+        if (params.mfaSerial) {
+            await 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 ?? '' }
+    }
 }