Add role ARN field to IAM credentials form

Author: liramon1

packages/core/src/auth/auth2.ts

@@ -155,6 +155,7 @@ export class LanguageClientAuth {
                     sso_session: profileName,
                     aws_access_key_id: '',
                     aws_secret_access_key: '',
+                    role_arn: '',
                 },
             },
             ssoSession: {
@@ -168,9 +169,9 @@ export class LanguageClientAuth {
         } satisfies UpdateProfileParams)
     }
 
-    updateIamProfile(profileName: string, accessKey: string, secretKey: string, sessionToken?: string): Promise<UpdateProfileResult> {
-        // Use unknown profile type if invalidating all IAM fields
-        const kind = !accessKey && !secretKey && !sessionToken ? ProfileKind.EmptyProfile : ProfileKind.IamCredentialProfile
+    updateIamProfile(profileName: string, accessKey: string, secretKey: string, sessionToken?: string, roleArn?: string): Promise<UpdateProfileResult> {
+        // Use empty profile type if invalidating all IAM fields
+        const kind = !accessKey && !secretKey ? ProfileKind.EmptyProfile : ProfileKind.IamCredentialProfile
         // Add credentials and delete SSO settings from profile
         return this.client.sendRequest(updateProfileRequestType.method, {
             profile: {
@@ -182,6 +183,7 @@ export class LanguageClientAuth {
                     aws_access_key_id: accessKey,
                     aws_secret_access_key: secretKey,
                     aws_session_token: sessionToken,
+                    role_arn: roleArn,
                 },
             },
             ssoSession: {
@@ -488,7 +490,7 @@ export class IamLogin extends BaseLogin {
         // )
     }
 
-    async login(opts: { accessKey: string; secretKey: string, sessionToken?: string }) {
+    async login(opts: { accessKey: string; secretKey: string, sessionToken?: string, roleArn?: string }) {
         await this.updateProfile(opts)
         return this._getIamCredential(true)
     }
@@ -504,18 +506,18 @@ 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, '', '', '', '')
         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 }) {
-        await this.lspAuth.updateIamProfile(this.profileName, opts.accessKey, opts.secretKey, opts.sessionToken)
+    async updateProfile(opts: { accessKey: string; secretKey: string, sessionToken?: string, roleArn?: string }) {
+        await this.lspAuth.updateIamProfile(this.profileName, opts.accessKey, opts.secretKey, opts.sessionToken, opts.roleArn)
         this._data = {
             accessKey: opts.accessKey,
             secretKey: opts.secretKey,
-            sessionToken: opts.sessionToken
+            sessionToken: opts.sessionToken,
         }
     }
 
@@ -585,7 +587,9 @@ export class IamLogin extends BaseLogin {
             this.cancellationToken = undefined
         }
 
-        this.iamCredentialId = response.id
+        if (response.credentials.sessionToken) {
+            this.iamCredentialId = response.id
+        }
         this.updateConnectionState('connected')
         return response
     }

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

@@ -174,13 +174,13 @@ export class AuthUtil implements IAuthProvider {
     }
 
     // Log in using IAM or STS credentials
-    async login_iam(accessKey: string, secretKey: string, sessionToken?: string): Promise<GetIamCredentialResult | undefined> {
+    async login_iam(accessKey: string, secretKey: string, sessionToken?: string, roleArn?: string): Promise<GetIamCredentialResult | undefined> {
         let response: GetIamCredentialResult | undefined
         // Create IAM login session
         if (!this.isIamSession()) {
             this.session = new IamLogin(this.profileName, this.lspAuth, this.eventEmitter)
         }
-        response = await (this.session as IamLogin).login({ accessKey: accessKey, secretKey: secretKey, sessionToken: sessionToken })
+        response = await (this.session as IamLogin).login({ accessKey: accessKey, secretKey: secretKey, sessionToken: sessionToken, roleArn: roleArn })
         await showAmazonQWalkthroughOnce()
         return response
     }

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

@@ -197,7 +197,8 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
         profileName: string,
         accessKey: string,
         secretKey: string,
-        sessionToken?: string
+        sessionToken?: string,
+        roleArn?: string
     ): Promise<AuthError | undefined> {
         getLogger().debug(`called startIamCredentialSetup()`)
         // Defining separate auth function to emit telemetry before returning from this method
@@ -206,7 +207,7 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
         })
         const runAuth = async (): Promise<AuthError | undefined> => {
             try {
-                await AuthUtil.instance.login_iam(accessKey, secretKey, sessionToken)
+                await AuthUtil.instance.login_iam(accessKey, secretKey, sessionToken, roleArn)
             } catch (e) {
                 getLogger().error('Failed submitting credentials %O', e)
                 return { id: this.id, text: e as string }

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

@@ -176,6 +176,7 @@ export abstract class CommonAuthWebview extends VueWebview {
         accessKey: string,
         secretKey: string,
         sessionToken?: string,
+        role_arn?: string,
     ): Promise<AuthError | undefined>
 
     async showResourceExplorer(): Promise<void> {

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

@@ -289,6 +289,15 @@
                     v-model="sessionToken"
                     @keydown.enter="handleContinueClick()"
                 />
+                <div class="title">Role ARN (Optional)</div>
+                <input
+                    class="iamInput bottomMargin"
+                    type="text"
+                    id="roleArn"
+                    name="roleArn"
+                    v-model="roleArn"
+                    @keydown.enter="handleContinueClick()"
+                />
             </div>
             <button class="continue-button" :disabled="shouldDisableIamContinue()" v-on:click="handleContinueClick()">
                 Continue
@@ -379,6 +388,7 @@ export default defineComponent({
             accessKey: '',
             secretKey: '',
             sessionToken: '',
+            roleArn: '',
         }
     },
     async created() {
@@ -517,7 +527,7 @@ export default defineComponent({
                     return
                 }
                 this.stage = 'AUTHENTICATING'
-                const error = await client.startIamCredentialSetup(this.profileName, this.accessKey, this.secretKey, this.sessionToken)
+                const error = await client.startIamCredentialSetup(this.profileName, this.accessKey, this.secretKey, this.sessionToken, this.roleArn)
                 if (error) {
                     this.stage = 'START'
                     void client.errorNotification(error)