add sts management and mfa

Author: yuxianrz

package-lock.json

@@ -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'
@@ -164,6 +166,9 @@ export async function startLanguageServer(
             },
             credentials: {
                 providesBearerToken: true,
+                // Add IAM credentials support
+                providesIamCredentials: true,
+                supportsAssumeRole: true,
             },
         },
         /**
@@ -211,9 +216,10 @@ export async function startLanguageServer(
 
         /** All must be setup before {@link AuthUtil.restore} otherwise they may not trigger when expected */
         AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(async () => {
+            const activeProfile = AuthUtil.instance.regionProfileManager.activeRegionProfile
             void pushConfigUpdate(client, {
                 type: 'profile',
-                profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn,
+                profileArn: activeProfile?.arn,
             })
         })
 
@@ -333,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/amazonq/src/lsp/client.ts

@@ -435,9 +435,50 @@ describe('AuthUtil', async function () {
 
             assert.ok(mockIamLogin.login.calledOnce)
             assert.ok(
-                mockIamLogin.login.calledWith({
+                mockIamLogin.login.calledWithMatch({
                     accessKey: 'accessKey',
                     secretKey: 'secretKey',
+                    sessionToken: 'sessionToken',
+                })
+            )
+            assert.strictEqual(response, mockResponse)
+        })
+
+        it('creates IAM session with role ARN', async function () {
+            const mockResponse = {
+                id: 'test-credential-id',
+                credentials: {
+                    accessKeyId: 'encrypted-access-key',
+                    secretAccessKey: 'encrypted-secret-key',
+                    sessionToken: 'encrypted-session-token',
+                    roleArn: 'arn:aws:iam::123456789012:role/TestRole',
+                },
+                updateCredentialsParams: {
+                    data: 'credential-data',
+                },
+            }
+
+            const mockIamLoginArn = {
+                login: sinon.stub().resolves(mockResponse),
+                loginType: 'iam',
+            }
+
+            sinon.stub(auth2, 'IamLogin').returns(mockIamLoginArn as any)
+
+            const response = await auth.loginIam(
+                'accessKey',
+                'secretKey',
+                'sessionToken',
+                'arn:aws:iam::123456789012:role/TestRole'
+            )
+
+            assert.ok(mockIamLoginArn.login.calledOnce)
+            assert.ok(
+                mockIamLoginArn.login.calledWith({
+                    accessKey: 'accessKey',
+                    secretKey: 'secretKey',
+                    sessionToken: 'sessionToken',
+                    roleArn: 'arn:aws:iam::123456789012:role/TestRole',
                 })
             )
             assert.strictEqual(response, mockResponse)

packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts

@@ -443,8 +443,8 @@
         "@aws-sdk/types": "^3.13.1",
         "@aws/chat-client": "^0.1.4",
         "@aws/chat-client-ui-types": "^0.1.53",
-        "@aws/language-server-runtimes": "^0.2.111",
-        "@aws/language-server-runtimes-types": "^0.1.47",
+        "@aws/language-server-runtimes": "^0.2.120",
+        "@aws/language-server-runtimes-types": "^0.1.52",
         "@cspotcode/source-map-support": "^0.8.1",
         "@sinonjs/fake-timers": "^10.0.2",
         "@types/adm-zip": "^0.4.34",