feat(smus): Add credential providers for SMUS (aws#2187)

**Description**

Added credential providers for DER, Project role credentials and
Connections credentials. Added in memory caching for now.

The credential providers currently are structured so that it can be used
with CredentialStore but a lot of things to do with credential store is
really old. In the coming days, while we do bug bashes and tests, if the
in memory caching fares well enough, will simplify the credential
provider to implement the AWS SDK CredentialProvider directly.


**Motivation**

Support auth for SMUS and ensure all clients get credentials and don't
fail due to lack of credentials.

**Testing Done**

Updated unit tests and tested flow locally. 

---

- Treat all work as PUBLIC. Private `feature/x` branches will not be
squash-merged at release time.
- Your code changes must meet the guidelines in
[CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines).
- License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

---------

Co-authored-by: Bhargava Varadharajan <[email protected]>

Author: vpbhargav

packages/core/src/auth/auth.ts

@@ -219,6 +219,29 @@ export class Auth implements AuthService, ConnectionManager {
         }
     }
 
+    /**
+     * Gets the SSO access token for a connection
+     * @param connection The SSO connection to get the token for
+     * @returns Promise resolving to the access token string
+     */
+    @withTelemetryContext({ name: 'getSsoAccessToken', class: authClassName })
+    public async getSsoAccessToken(connection: Pick<SsoConnection, 'id'>): Promise<string> {
+        const profile = this.store.getProfileOrThrow(connection.id)
+
+        if (profile.type !== 'sso') {
+            throw new Error(`Connection ${connection.id} is not an SSO connection`)
+        }
+
+        const provider = this.getSsoTokenProvider(connection.id, profile)
+        const token = await provider.getToken()
+
+        if (!token?.accessToken) {
+            throw new Error(`No access token available for connection ${connection.id}`)
+        }
+
+        return token.accessToken
+    }
+
     public async useConnection({ id }: Pick<SsoConnection, 'id'>): Promise<SsoConnection>
     public async useConnection({ id }: Pick<IamConnection, 'id'>): Promise<IamConnection>
     @withTelemetryContext({ name: 'useConnection', class: authClassName })

packages/core/src/sagemakerunifiedstudio/auth/providers/connectionCredentialsProvider.ts

@@ -0,0 +1,220 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { getLogger } from '../../../shared/logger/logger'
+import { ToolkitError } from '../../../shared/errors'
+import * as AWS from '@aws-sdk/types'
+import { CredentialsId, CredentialsProvider, CredentialsProviderType } from '../../../auth/providers/credentials'
+
+import { DataZoneClient } from '../../shared/client/datazoneClient'
+import { SmusAuthenticationProvider } from './smusAuthenticationProvider'
+import { CredentialType } from '../../../shared/telemetry/telemetry'
+import { SmusCredentialExpiry, validateCredentialFields } from '../../shared/smusUtils'
+
+/**
+ * Credentials provider for SageMaker Unified Studio Connection credentials
+ * Uses DataZone API to get connection credentials for a specific connection *
+ * This provider implements independent caching with 10-minute expiry
+ */
+export class ConnectionCredentialsProvider implements CredentialsProvider {
+    private readonly logger = getLogger()
+    private credentialCache?: {
+        credentials: AWS.Credentials
+        expiresAt: Date
+    }
+
+    constructor(
+        private readonly smusAuthProvider: SmusAuthenticationProvider,
+        private readonly connectionId: string
+    ) {}
+
+    /**
+     * Gets the connection ID
+     * @returns Connection ID
+     */
+    public getConnectionId(): string {
+        return this.connectionId
+    }
+
+    /**
+     * Gets the credentials ID
+     * @returns Credentials ID
+     */
+    public getCredentialsId(): CredentialsId {
+        return {
+            credentialSource: 'temp',
+            credentialTypeId: `${this.smusAuthProvider.getDomainId()}:${this.connectionId}`,
+        }
+    }
+
+    /**
+     * Gets the provider type
+     * @returns Provider type
+     */
+    public getProviderType(): CredentialsProviderType {
+        return 'temp'
+    }
+
+    /**
+     * Gets the telemetry type
+     * @returns Telemetry type
+     */
+    public getTelemetryType(): CredentialType {
+        return 'other'
+    }
+
+    /**
+     * Gets the default region
+     * @returns Default region
+     */
+    public getDefaultRegion(): string | undefined {
+        return this.smusAuthProvider.getDomainRegion()
+    }
+
+    /**
+     * Gets the hash code
+     * @returns Hash code
+     */
+    public getHashCode(): string {
+        const hashCode = `smus-connection:${this.smusAuthProvider.getDomainId()}:${this.connectionId}`
+        return hashCode
+    }
+
+    /**
+     * Determines if the provider can auto-connect
+     * @returns Promise resolving to boolean
+     */
+    public async canAutoConnect(): Promise<boolean> {
+        return false // SMUS requires manual authentication
+    }
+
+    /**
+     * Determines if the provider is available
+     * @returns Promise resolving to boolean
+     */
+    public async isAvailable(): Promise<boolean> {
+        try {
+            return this.smusAuthProvider.isConnected()
+        } catch (err) {
+            this.logger.error('SMUS Connection: Error checking if auth provider is connected: %s', err)
+            return false
+        }
+    }
+
+    /**
+     * Gets Connection credentials with independent caching
+     * @returns Promise resolving to credentials
+     */
+    public async getCredentials(): Promise<AWS.Credentials> {
+        this.logger.debug(`SMUS Connection: Getting credentials for connection ${this.connectionId}`)
+
+        // Check cache first (10-minute expiry)
+        if (this.credentialCache && this.credentialCache.expiresAt > new Date()) {
+            this.logger.debug(
+                `SMUS Connection: Using cached connection credentials for connection ${this.connectionId}`
+            )
+            return this.credentialCache.credentials
+        }
+
+        this.logger.debug(
+            `SMUS Connection: Calling GetConnection to fetch credentials for connection ${this.connectionId}`
+        )
+
+        try {
+            const datazoneClient = await DataZoneClient.getInstance(this.smusAuthProvider)
+            const getConnectionResponse = await datazoneClient.getConnection({
+                domainIdentifier: this.smusAuthProvider.getDomainId(),
+                identifier: this.connectionId,
+                withSecret: true,
+            })
+
+            this.logger.debug(`SMUS Connection: Successfully retrieved connection details for ${this.connectionId}`)
+
+            // Extract connection credentials
+            const connectionCredentials = getConnectionResponse.connectionCredentials
+            if (!connectionCredentials) {
+                throw new ToolkitError(
+                    `No connection credentials available in response for connection ${this.connectionId}`,
+                    {
+                        code: 'NoConnectionCredentials',
+                    }
+                )
+            }
+
+            // Validate credential fields
+            validateCredentialFields(
+                connectionCredentials,
+                'InvalidConnectionCredentials',
+                'connection credential response'
+            )
+
+            // Create AWS credentials with expiration
+            // Use the expiration from the response if available, otherwise default to 10 minutes
+            let expiresAt: Date
+            if (connectionCredentials.expiration) {
+                // The API returns expiration as a string or Date, handle both cases
+                expiresAt =
+                    connectionCredentials.expiration instanceof Date
+                        ? connectionCredentials.expiration
+                        : new Date(connectionCredentials.expiration)
+            } else {
+                expiresAt = new Date(Date.now() + SmusCredentialExpiry.connectionExpiryMs)
+            }
+
+            const awsCredentials: AWS.Credentials = {
+                accessKeyId: connectionCredentials.accessKeyId as string,
+                secretAccessKey: connectionCredentials.secretAccessKey as string,
+                sessionToken: connectionCredentials.sessionToken as string,
+                expiration: expiresAt,
+            }
+
+            // Cache connection credentials (10-minute expiry)
+            const cacheExpiresAt = new Date(Date.now() + SmusCredentialExpiry.connectionExpiryMs)
+            this.credentialCache = {
+                credentials: awsCredentials,
+                expiresAt: cacheExpiresAt,
+            }
+
+            this.logger.debug(
+                `SMUS Connection: Successfully cached connection credentials for connection ${this.connectionId}, expires in %s minutes`,
+                Math.round((cacheExpiresAt.getTime() - Date.now()) / 60000)
+            )
+
+            return awsCredentials
+        } catch (err) {
+            this.logger.error(
+                `SMUS Connection: Failed to get connection credentials for connection ${this.connectionId}: %s`,
+                err
+            )
+
+            // Re-throw ToolkitErrors with specific codes (NoConnectionCredentials, InvalidConnectionCredentials)
+            if (
+                err instanceof ToolkitError &&
+                (err.code === 'NoConnectionCredentials' || err.code === 'InvalidConnectionCredentials')
+            ) {
+                throw err
+            }
+
+            // Wrap other errors in ConnectionCredentialsFetchFailed
+            throw new ToolkitError(`Failed to get connection credentials for ${this.connectionId}: ${err}`, {
+                code: 'ConnectionCredentialsFetchFailed',
+                cause: err instanceof Error ? err : undefined,
+            })
+        }
+    }
+
+    /**
+     * Invalidates cached connection credentials
+     * Clears the internal cache without fetching new credentials
+     */
+    public invalidate(): void {
+        this.logger.debug(`SMUS Connection: Invalidating cached credentials for connection ${this.connectionId}`)
+        // Clear cache to force fresh fetch on next getCredentials() call
+        this.credentialCache = undefined
+        this.logger.debug(
+            `SMUS Connection: Successfully invalidated connection credentials cache for connection ${this.connectionId}`
+        )
+    }
+}