Move IAM profile options into typed object

liramon1 · liramon1 · commit 5a71c09d7f71 · 2025-08-05T14:01:12.000-04:00
diff --git a/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts
@@ -11,6 +11,7 @@ import { createTestAuthUtil, TestFolder } from 'aws-core-vscode/test'
 import { constants, cache } from 'aws-core-vscode/auth'
 import { auth2 } from 'aws-core-vscode/auth'
 import { mementoUtils, fs } from 'aws-core-vscode/shared'
+import { GetIamCredentialResult } from '@aws/language-server-runtimes/protocol'
 
 describe('AuthUtil', async function () {
     let auth: any
@@ -431,7 +432,11 @@ describe('AuthUtil', async function () {
 
             sinon.stub(auth2, 'IamLogin').returns(mockIamLogin as any)
 
-            const response = await auth.loginIam('accessKey', 'secretKey', 'sessionToken')
+            const response = await auth.loginIam({
+                accessKey: 'accessKey',
+                secretKey: 'secretKey',
+                sessionToken: 'sessionToken',
+            })
 
             assert.ok(mockIamLogin.login.calledOnce)
             assert.ok(
@@ -445,13 +450,15 @@ describe('AuthUtil', async function () {
         })
 
         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',
+            const mockResponse: GetIamCredentialResult = {
+                credential: {
+                    id: 'test-credential-id',
+                    kinds: [],
+                    credentials: {
+                        accessKeyId: 'encrypted-access-key',
+                        secretAccessKey: 'encrypted-secret-key',
+                        sessionToken: 'encrypted-session-token',
+                    },
                 },
                 updateCredentialsParams: {
                     data: 'credential-data',
@@ -465,22 +472,17 @@ describe('AuthUtil', async function () {
 
             sinon.stub(auth2, 'IamLogin').returns(mockIamLoginArn as any)
 
-            const response = await auth.loginIam(
-                'accessKey',
-                'secretKey',
-                'sessionToken',
-                'arn:aws:iam::123456789012:role/TestRole'
-            )
+            const opts: auth2.IamProfileOptions = {
+                accessKey: 'myAccessKey',
+                secretKey: 'mySecretKey',
+                sessionToken: 'mySessionToken',
+                roleArn: 'arn:aws:iam::123456789012:role/MyTestRole',
+            }
+
+            const response = await auth.loginIam(opts)
 
             assert.ok(mockIamLoginArn.login.calledOnce)
-            assert.ok(
-                mockIamLoginArn.login.calledWith({
-                    accessKey: 'accessKey',
-                    secretKey: 'secretKey',
-                    sessionToken: 'sessionToken',
-                    roleArn: 'arn:aws:iam::123456789012:role/TestRole',
-                })
-            )
+            assert.ok(mockIamLoginArn.login.calledWith(opts))
             assert.strictEqual(response, mockResponse)
         })
     })
diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts
@@ -95,6 +95,22 @@ export type Login = SsoLogin | IamLogin
 
 export type TokenSource = IamIdentityCenterSsoTokenSource | AwsBuilderIdSsoTokenSource
 
+export type IamProfileOptions = {
+    accessKey?: string
+    secretKey?: string
+    sessionToken?: string
+    roleArn?: string
+    sourceProfile?: string
+}
+
+const IamProfileOptionsDefaults = {
+    accessKey: '',
+    secretKey: '',
+    sessionToken: '',
+    roleArn: '',
+    sourceProfile: '',
+} satisfies IamProfileOptions
+
 /**
  * Handles auth requests to the Identity Server in the Amazon Q LSP.
  */
@@ -185,58 +201,32 @@ export class LanguageClientAuth {
         } satisfies UpdateProfileParams)
     }
 
-    updateIamProfile(
-        profileName: string,
-        accessKey: string,
-        secretKey: string,
-        sessionToken?: string,
-        roleArn?: string,
-        sourceProfile?: string
-    ): Promise<UpdateProfileResult> {
-        // Add credentials and delete SSO settings from profile
-        let profile: Profile
-        if (roleArn && sourceProfile) {
-            profile = {
-                kinds: [ProfileKind.IamSourceProfileProfile],
-                name: profileName,
-                settings: {
-                    sso_session: '',
-                    aws_access_key_id: '',
-                    aws_secret_access_key: '',
-                    aws_session_token: '',
-                    role_arn: roleArn,
-                    source_profile: sourceProfile,
-                },
-            }
-        } else if (accessKey && secretKey) {
-            profile = {
-                kinds: [ProfileKind.IamCredentialsProfile],
-                name: profileName,
-                settings: {
-                    sso_session: '',
-                    aws_access_key_id: accessKey,
-                    aws_secret_access_key: secretKey,
-                    aws_session_token: sessionToken,
-                    role_arn: '',
-                    source_profile: '',
-                },
-            }
+    updateIamProfile(profileName: string, opts: IamProfileOptions): Promise<UpdateProfileResult> {
+        // Substitute missing fields for defaults
+        const fields = { ...IamProfileOptionsDefaults, ...opts }
+        // Get the profile kind matching the provided fields
+        let kind: ProfileKind
+        if (fields.roleArn && fields.sourceProfile) {
+            kind = ProfileKind.IamSourceProfileProfile
+        } else if (fields.accessKey && fields.secretKey) {
+            kind = ProfileKind.IamCredentialsProfile
         } else {
-            profile = {
-                kinds: [ProfileKind.Unknown],
+            kind = ProfileKind.Unknown
+        }
+
+        return this.client.sendRequest(updateProfileRequestType.method, {
+            profile: {
+                kinds: [kind],
                 name: profileName,
                 settings: {
-                    aws_access_key_id: '',
-                    aws_secret_access_key: '',
-                    aws_session_token: '',
-                    role_arn: '',
-                    source_profile: '',
+                    aws_access_key_id: fields.accessKey,
+                    aws_secret_access_key: fields.secretKey,
+                    aws_session_token: fields.sessionToken,
+                    role_arn: fields.roleArn,
+                    source_profile: fields.sourceProfile,
                 },
-            }
-        }
-        return this.client.sendRequest(updateProfileRequestType.method, {
-            profile: profile,
-        } satisfies UpdateProfileParams)
+            },
+        })
     }
 
     listProfiles() {
@@ -550,7 +540,7 @@ export class IamLogin extends BaseLogin {
         lspAuth.registerGetMfaCodeHandler((params: GetMfaCodeParams) => this.getMfaCodeHandler(params))
     }
 
-    async login(opts: { accessKey: string; secretKey: string; sessionToken?: string; roleArn?: string }) {
+    async login(opts: IamProfileOptions) {
         await this.updateProfile(opts)
         return this._getIamCredential(true)
     }
@@ -566,34 +556,33 @@ export class IamLogin extends BaseLogin {
         if (this.iamCredentialId) {
             await this.lspAuth.invalidateStsCredential(this.iamCredentialId)
         }
-        await this.lspAuth.updateIamProfile(this.profileName, '', '', '', '', '')
-        await this.lspAuth.updateIamProfile(this.profileName + '-source', '', '', '', '', '')
+        await this.lspAuth.updateIamProfile(this.profileName, {})
+        await this.lspAuth.updateIamProfile(this.profileName + '-source', {})
         this.updateConnectionState('notConnected')
         this._data = undefined
         // TODO: DeleteProfile api in Identity Service (this doesn't exist yet)
     }
 
-    async updateProfile(opts: { accessKey: string; secretKey: string; sessionToken?: string; roleArn?: string }) {
+    async updateProfile(opts: IamProfileOptions) {
         if (opts.roleArn) {
+            // Create the source and target profiles
             const sourceProfile = this.profileName + '-source'
-            await this.lspAuth.updateIamProfile(
-                sourceProfile,
-                opts.accessKey,
-                opts.secretKey,
-                opts.sessionToken,
-                '',
-                ''
-            )
-            await this.lspAuth.updateIamProfile(this.profileName, '', '', '', opts.roleArn, sourceProfile)
+            await this.lspAuth.updateIamProfile(sourceProfile, {
+                accessKey: opts.accessKey,
+                secretKey: opts.secretKey,
+                sessionToken: opts.sessionToken,
+            })
+            await this.lspAuth.updateIamProfile(this.profileName, {
+                roleArn: opts.roleArn,
+                sourceProfile: sourceProfile,
+            })
         } else {
-            await this.lspAuth.updateIamProfile(
-                this.profileName,
-                opts.accessKey,
-                opts.secretKey,
-                opts.sessionToken,
-                '',
-                ''
-            )
+            // Create the target profile
+            await this.lspAuth.updateIamProfile(this.profileName, {
+                accessKey: opts.accessKey,
+                secretKey: opts.secretKey,
+                sessionToken: opts.sessionToken,
+            })
         }
     }
 
diff --git a/packages/core/src/auth/credentials/utils.ts b/packages/core/src/auth/credentials/utils.ts
@@ -105,11 +105,8 @@ const errorMessageUserCancelled = localize('AWS.error.mfa.userCancelled', 'User
 /**
  * @description Prompts user for MFA serial number
  *
- * Entered token is passed to the callback.
- * If user cancels out, the callback is passed an error with a fixed message string.
- *
+ * @param defaultSerial Default value for the serial number input
  * @param profileName Name of Credentials profile we are asking an MFA Token for
- * @param callback tokens/errors are passed through here
  */
 export async function getMfaSerialFromUser(defaultSerial: string, profileName: string): Promise<string> {
     const inputBox = createInputBox({
@@ -133,12 +130,8 @@ export async function getMfaSerialFromUser(defaultSerial: string, profileName: s
 /**
  * @description Prompts user for MFA token
  *
- * Entered token is passed to the callback.
- * If user cancels out, the callback is passed an error with a fixed message string.
- *
  * @param mfaSerial Serial arn of MFA device
  * @param profileName Name of Credentials profile we are asking an MFA Token for
- * @param callback tokens/errors are passed through here
  */
 export async function getMfaTokenFromUser(mfaSerial: string, profileName: string): Promise<string> {
     const inputBox = createInputBox({
diff --git a/packages/core/src/auth/sso/cache.ts b/packages/core/src/auth/sso/cache.ts
@@ -36,7 +36,7 @@ export interface SsoCache {
 }
 
 const defaultCacheDir = () => path.join(fs.getUserHomeDir(), '.aws/sso/cache')
-const defaultStsCacheDir = () => path.join(fs.getUserHomeDir(), '.aws/cli/cache')
+const defaultStsCacheDir = () => path.join(fs.getUserHomeDir(), '.aws/flare/cache')
 export const getCacheDir = () => DevSettings.instance.get('ssoCacheDirectory', defaultCacheDir())
 export const getStsCacheDir = () => DevSettings.instance.get('stsCacheDirectory', defaultStsCacheDir())
 
diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts
@@ -40,6 +40,7 @@ import {
     IamLogin,
     AuthState,
     LoginTypes,
+    IamProfileOptions,
 } from '../../auth/auth2'
 import { builderIdStartUrl, internalStartUrl } from '../../auth/sso/constants'
 import { VSCODE_EXTENSION_ID } from '../../shared/extensions'
@@ -193,22 +194,12 @@ export class AuthUtil implements IAuthProvider {
     }
 
     // Log in using IAM or STS credentials
-    async loginIam(
-        accessKey: string,
-        secretKey: string,
-        sessionToken?: string,
-        roleArn?: string
-    ): Promise<GetIamCredentialResult | undefined> {
+    async loginIam(opts: IamProfileOptions): Promise<GetIamCredentialResult | undefined> {
         // Create IAM login session
         if (!this.isIamSession()) {
             this.session = new IamLogin(this.profileName, this.lspAuth, this.eventEmitter)
         }
-        const response = await (this.session as IamLogin).login({
-            accessKey: accessKey,
-            secretKey: secretKey,
-            sessionToken: sessionToken,
-            roleArn: roleArn,
-        })
+        const response = await (this.session as IamLogin).login(opts)
         await showAmazonQWalkthroughOnce()
         return response
     }
diff --git a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts
@@ -206,7 +206,7 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
         await globals.globalState.update('recentRoleArn', { roleArn: roleArn })
         const runAuth = async (): Promise<AuthError | undefined> => {
             try {
-                await AuthUtil.instance.loginIam(accessKey, secretKey, sessionToken, roleArn)
+                await AuthUtil.instance.loginIam({ accessKey, secretKey, sessionToken, roleArn })
             } catch (e) {
                 getLogger().error('Failed submitting credentials %O', e)
                 const message = e instanceof Error ? e.message : (e as string)
diff --git a/packages/core/src/test/credentials/auth2.test.ts b/packages/core/src/test/credentials/auth2.test.ts
@@ -104,7 +104,11 @@ describe('LanguageClientAuth', () => {
         })
 
         it('sends correct IAM profile update parameters', async () => {
-            await auth.updateIamProfile(profileName, 'accessKey', 'secretKey', 'sessionToken')
+            await auth.updateIamProfile(profileName, {
+                accessKey: 'myAccessKey',
+                secretKey: 'mySecretKey',
+                sessionToken: 'mySessionToken',
+            })
 
             sinon.assert.calledOnce(client.sendRequest)
             const requestParams = client.sendRequest.firstCall.args[1]
@@ -113,9 +117,11 @@ describe('LanguageClientAuth', () => {
                 kinds: [ProfileKind.IamCredentialsProfile],
             })
             sinon.assert.match(requestParams.profile.settings, {
-                aws_access_key_id: 'accessKey',
-                aws_secret_access_key: 'secretKey',
-                aws_session_token: 'sessionToken',
+                aws_access_key_id: 'myAccessKey',
+                aws_secret_access_key: 'mySecretKey',
+                aws_session_token: 'mySessionToken',
+                role_arn: '',
+                source_profile: '',
             })
         })
     })
@@ -674,7 +680,7 @@ describe('IamLogin', () => {
             const response = await iamLogin.login(loginOpts)
 
             sinon.assert.calledOnce(lspAuth.updateIamProfile)
-            sinon.assert.calledWith(lspAuth.updateIamProfile, profileName, loginOpts.accessKey, loginOpts.secretKey)
+            sinon.assert.calledWith(lspAuth.updateIamProfile, profileName, loginOpts)
             sinon.assert.calledOnce(lspAuth.getIamCredential)
             sinon.assert.match(iamLogin.getConnectionState(), 'connected')
             sinon.assert.match(response.credential.id, 'test-credential-id')

Original file line number	Diff line number	Diff line change
`@@ -36,7 +36,7 @@ export interface SsoCache {`
`36`	`36`	`}`
`37`	`37`
`38`	`38`	`const defaultCacheDir = () => path.join(fs.getUserHomeDir(), '.aws/sso/cache')`
`39`		`-const defaultStsCacheDir = () => path.join(fs.getUserHomeDir(), '.aws/cli/cache')`
	`39`	`+const defaultStsCacheDir = () => path.join(fs.getUserHomeDir(), '.aws/flare/cache')`
`40`	`40`	`export const getCacheDir = () => DevSettings.instance.get('ssoCacheDirectory', defaultCacheDir())`
`41`	`41`	`export const getStsCacheDir = () => DevSettings.instance.get('stsCacheDirectory', defaultStsCacheDir())`
`42`	`42`