Merge branch 'feature/v2-to-v3-migration' into migrate-schemas

Author: ctlai95

package-lock.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Amazon Q automatically refreshes expired IAM Credentials in Sagemaker instances"
+}

packages/amazonq/.changes/next-release/Bug Fix-da0e805d-faab-4274-b37a-943c7263e42b.json

@@ -17,9 +17,9 @@ import * as crypto from 'crypto'
 import { LanguageClient } from 'vscode-languageclient'
 import { AuthUtil } from 'aws-core-vscode/codewhisperer'
 import { Writable } from 'stream'
-import { onceChanged } from 'aws-core-vscode/utils'
+import { onceChanged, onceChangedWithComparator } from 'aws-core-vscode/utils'
 import { getLogger, oneMinute, isSageMaker } from 'aws-core-vscode/shared'
-import { isSsoConnection, isIamConnection } from 'aws-core-vscode/auth'
+import { isSsoConnection, isIamConnection, areCredentialsEqual } from 'aws-core-vscode/auth'
 
 export const encryptionKey = crypto.randomBytes(32)
 
@@ -108,7 +108,10 @@ export class AmazonQLspAuth {
         this.client.info(`UpdateBearerToken: ${JSON.stringify(request)}`)
     }
 
-    public updateIamCredentials = onceChanged(this._updateIamCredentials.bind(this))
+    public updateIamCredentials = onceChangedWithComparator(
+        this._updateIamCredentials.bind(this),
+        ([prevCreds], [currentCreds]) => areCredentialsEqual(prevCreds, currentCreds)
+    )
     private async _updateIamCredentials(credentials: any) {
         getLogger().info(
             `[SageMaker Debug] Updating IAM credentials - credentials received: ${credentials ? 'YES' : 'NO'}`

packages/amazonq/src/lsp/auth.ts

@@ -539,6 +539,7 @@
         "@types/whatwg-url": "^11.0.4",
         "@types/xml2js": "^0.4.11",
         "@vue/compiler-sfc": "^3.3.2",
+        "aws-sdk-client-mock": "^4.1.0",
         "c8": "^9.0.0",
         "circular-dependency-plugin": "^5.2.2",
         "css-loader": "^6.10.0",
@@ -580,6 +581,8 @@
         "@aws-sdk/client-ec2": "<3.731.0",
         "@aws-sdk/client-glue": "^3.852.0",
         "@aws-sdk/client-iam": "<3.731.0",
+        "@aws-sdk/client-iot": "~3.693.0",
+        "@aws-sdk/client-iotsecuretunneling": "~3.693.0",
         "@aws-sdk/client-lambda": "<3.731.0",
         "@aws-sdk/client-s3": "<3.731.0",
         "@aws-sdk/client-s3-control": "^3.830.0",

packages/core/package.json

@@ -862,6 +862,7 @@ export class Auth implements AuthService, ConnectionManager {
 
     private async createCachedCredentials(provider: CredentialsProvider) {
         const providerId = provider.getCredentialsId()
+        getLogger().debug(`credentials: create cache credentials for ${provider.getProviderType()}`)
         globals.loginManager.store.invalidateCredentials(providerId)
         const { credentials, endpointUrl } = await globals.loginManager.store.upsertCredentials(providerId, provider)
         await globals.loginManager.validateCredentials(credentials, endpointUrl, provider.getDefaultRegion())

packages/core/src/auth/auth.ts

@@ -71,6 +71,18 @@ export const isBuilderIdConnection = (conn?: Connection): conn is SsoConnection
 export const isValidCodeCatalystConnection = (conn?: Connection): conn is SsoConnection =>
     isSsoConnection(conn) && hasScopes(conn, scopesCodeCatalyst)
 
+export const areCredentialsEqual = (creds1: any, creds2: any): boolean => {
+    if (!creds1 || !creds2) {
+        return creds1 === creds2
+    }
+
+    return (
+        creds1.accessKeyId === creds2.accessKeyId &&
+        creds1.secretAccessKey === creds2.secretAccessKey &&
+        creds1.sessionToken === creds2.sessionToken
+    )
+}
+
 export function hasScopes(target: SsoConnection | SsoProfile | string[], scopes: string[]): boolean {
     return scopes?.every((s) => (Array.isArray(target) ? target : target.scopes)?.includes(s))
 }

packages/core/src/auth/connection.ts

@@ -31,11 +31,16 @@ export class CredentialsStore {
      * If the expiration property does not exist, it is assumed to never expire.
      */
     public isValid(key: string): boolean {
+        // Apply 60-second buffer similar to SSO token expiry logic
+        const expirationBufferMs = 60000
+
         if (this.credentialsCache[key]) {
             const expiration = this.credentialsCache[key].credentials.expiration
-            return expiration !== undefined ? expiration >= new globals.clock.Date() : true
+            const now = new globals.clock.Date()
+            const bufferedNow = new globals.clock.Date(now.getTime() + expirationBufferMs)
+            return expiration !== undefined ? expiration >= bufferedNow : true
         }
-
+        getLogger().debug(`credentials: no credentials found for ${key}`)
         return false
     }
 

packages/core/src/auth/credentials/store.ts

@@ -19,6 +19,7 @@ export {
     getTelemetryMetadataForConn,
     isIamConnection,
     isSsoConnection,
+    areCredentialsEqual,
 } from './connection'
 export { Auth } from './auth'
 export { CredentialsStore } from './credentials/store'

packages/core/src/auth/index.ts

@@ -12,7 +12,7 @@ import { createQuickPick, DataQuickPickItem } from '../../../shared/ui/pickerPro
 import { PromptResult } from '../../../shared/ui/prompter'
 import { IotClient } from '../../../shared/clients/iotClient'
 import { isValidResponse } from '../../../shared/wizards/wizard'
-import { Iot } from 'aws-sdk'
+import { Certificate, ListCertificatesResponse } from '@aws-sdk/client-iot'
 
 export type CertGen = typeof getCertList
 
@@ -53,8 +53,8 @@ export async function attachCertificateCommand(node: IotThingNode, promptFun = p
 /**
  * Prompts the user to pick a certificate to attach.
  */
-async function promptForCert(iot: IotClient, certFetch: CertGen): Promise<PromptResult<Iot.Certificate>> {
-    const placeHolder: DataQuickPickItem<Iot.Certificate> = {
+async function promptForCert(iot: IotClient, certFetch: CertGen): Promise<PromptResult<Certificate>> {
+    const placeHolder: DataQuickPickItem<Certificate> = {
         label: 'No certificates found',
         data: undefined,
     }
@@ -71,10 +71,10 @@ async function promptForCert(iot: IotClient, certFetch: CertGen): Promise<Prompt
  */
 async function* getCertList(iot: IotClient) {
     let marker: string | undefined = undefined
-    let filteredCerts: Iot.Certificate[]
+    let filteredCerts: Certificate[]
     do {
         try {
-            const certResponse: Iot.ListCertificatesResponse = await iot.listCertificates({ marker })
+            const certResponse: ListCertificatesResponse = await iot.listCertificates({ marker })
             marker = certResponse.nextMarker
 
             /* These fields should always be defined when using the above API,

packages/core/src/awsService/iot/commands/attachCertificate.ts

@@ -4,7 +4,14 @@
  */
 
 import * as vscode from 'vscode'
-import { IoTSecureTunneling, Lambda } from 'aws-sdk'
+import { Lambda } from 'aws-sdk'
+import {
+    CloseTunnelCommand,
+    IoTSecureTunnelingClient,
+    ListTunnelsCommand,
+    OpenTunnelCommand,
+    RotateTunnelAccessTokenCommand,
+} from '@aws-sdk/client-iotsecuretunneling'
 import { getClientId } from '../../shared/telemetry/util'
 import { DefaultLambdaClient } from '../../shared/clients/lambdaClient'
 import { LocalProxy } from './localProxy'
@@ -61,7 +68,7 @@ export class LdkClient {
     private localProxy: LocalProxy | undefined
     private static instanceCreating = false
     private lambdaClientCache: Map<string, DefaultLambdaClient> = new Map()
-    private iotSTClientCache: Map<string, IoTSecureTunneling> = new Map()
+    private iotSTClientCache: Map<string, IoTSecureTunnelingClient> = new Map()
 
     constructor() {}
 
@@ -97,9 +104,9 @@ export class LdkClient {
         return this.lambdaClientCache.get(region)!
     }
 
-    private async getIoTSTClient(region: string): Promise<IoTSecureTunneling> {
+    private getIoTSTClient(region: string): IoTSecureTunnelingClient {
         if (!this.iotSTClientCache.has(region)) {
-            this.iotSTClientCache.set(region, await getIoTSTClientWithAgent(region))
+            this.iotSTClientCache.set(region, getIoTSTClientWithAgent(region))
         }
         return this.iotSTClientCache.get(region)!
     }
@@ -124,13 +131,13 @@ export class LdkClient {
             const vscodeUuid = getClientId(globals.globalState)
 
             // Create IoTSecureTunneling client
-            const iotSecureTunneling = await this.getIoTSTClient(region)
+            const iotSecureTunneling = this.getIoTSTClient(region)
 
             // Define tunnel identifier
             const tunnelIdentifier = `RemoteDebugging+${vscodeUuid}`
             const timeoutInMinutes = 720
             // List existing tunnels
-            const listTunnelsResponse = await iotSecureTunneling.listTunnels({}).promise()
+            const listTunnelsResponse = await iotSecureTunneling.send(new ListTunnelsCommand({}))
 
             // Find tunnel with our identifier
             const existingTunnel = listTunnelsResponse.tunnelSummaries?.find(
@@ -150,20 +157,20 @@ export class LdkClient {
                     return rotateResponse
                 } else {
                     // Close tunnel if less than 15 minutes remaining
-                    await iotSecureTunneling
-                        .closeTunnel({
+                    await iotSecureTunneling.send(
+                        new CloseTunnelCommand({
                             tunnelId: existingTunnel.tunnelId,
                             delete: false,
                         })
-                        .promise()
+                    )
 
                     getLogger().info(`Closed tunnel ${existingTunnel.tunnelId} with less than 15 minutes remaining`)
                 }
             }
 
             // Create new tunnel
-            const openTunnelResponse = await iotSecureTunneling
-                .openTunnel({
+            const openTunnelResponse = await iotSecureTunneling.send(
+                new OpenTunnelCommand({
                     description: tunnelIdentifier,
                     timeoutConfig: {
                         maxLifetimeTimeoutMinutes: timeoutInMinutes, // 12 hours
@@ -172,7 +179,7 @@ export class LdkClient {
                         services: ['WSS'],
                     },
                 })
-                .promise()
+            )
 
             getLogger().info(`Created new tunnel with ID: ${openTunnelResponse.tunnelId}`)
 
@@ -189,13 +196,13 @@ export class LdkClient {
     // Refresh tunnel tokens
     async refreshTunnelTokens(tunnelId: string, region: string): Promise<TunnelInfo | undefined> {
         try {
-            const iotSecureTunneling = await this.getIoTSTClient(region)
-            const rotateResponse = await iotSecureTunneling
-                .rotateTunnelAccessToken({
+            const iotSecureTunneling = this.getIoTSTClient(region)
+            const rotateResponse = await iotSecureTunneling.send(
+                new RotateTunnelAccessTokenCommand({
                     tunnelId: tunnelId,
                     clientMode: 'ALL',
                 })
-                .promise()
+            )
 
             return {
                 tunnelID: tunnelId,