Add session tokens to auth2 logic

Author: liramon1

packages/core/src/auth/auth2.ts

@@ -273,7 +273,7 @@ export class LanguageClientAuth {
 export abstract class BaseLogin {
     protected connectionState: AuthState = 'notConnected'
     protected cancellationToken: CancellationTokenSource | undefined
-    protected _data: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string } | undefined
+    protected _data: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string; sessionToken?: string } | undefined
 
     constructor(
         public readonly profileName: string,
@@ -525,6 +525,7 @@ export class IamLogin extends BaseLogin {
         this._data = {
             accessKey: opts.accessKey,
             secretKey: opts.secretKey,
+            sessionToken: opts.sessionToken
         }
     }
 
@@ -542,6 +543,7 @@ export class IamLogin extends BaseLogin {
             this._data = {
                 accessKey: credentials.aws_access_key_id,
                 secretKey: credentials.aws_secret_access_key,
+                sessionToken: credentials.aws_session_token
             }
         }
         try {

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

@@ -64,7 +64,7 @@ export interface IAuthProvider {
     getToken(): Promise<string>
     getIamCredential(): Promise<IamCredentials>
     readonly profileName: string
-    readonly connection?: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string }
+    readonly connection?: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string; sessionToken?: string }
 }
 
 /**
@@ -132,12 +132,13 @@ export class AuthUtil implements IAuthProvider {
             this.session = new SsoLogin(this.profileName, this.lspAuth, this.eventEmitter)
             await this.session.restore()
             if (!this.isConnected()) {
+                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
+                    this.session?.logout()
                 }
             }
         }