Merge pull request aws#7997 from valerena/endpoint-url-plus-localstack

Author: rli

packages/core/package.nls.json

@@ -478,7 +478,7 @@
     "AWS.toolkit.lambda.walkthrough.title": "Get started building your application",
     "AWS.toolkit.lambda.walkthrough.description": "Your quick guide to build an application visually, iterate locally, and deploy to the cloud!",
     "AWS.toolkit.lambda.walkthrough.toolInstall.title": "Complete installation",
-    "AWS.toolkit.lambda.walkthrough.toolInstall.description": "The AWS Command Line Interface (AWS CLI) is an open source tool that enables you to interact with AWS services using commands in your command-line shell. It is required to create and interact with AWS resources. \n\n[Install AWS CLI](command:aws.toolkit.installAWSCLI)\n\n Use the Serverless Application Model (SAM) CLI to locally build, invoke, and deploy your functions. Version 1.98+ is required. \n\n[Install SAM CLI](command:aws.toolkit.installSAMCLI)\n\n Use Docker to locally emulate a Lambda environment. Docker is optional. However, if you want to invoke locally, Docker is required so Lambda can locally emulate the execution environment. \n\n[Install Docker (optional)](command:aws.toolkit.installDocker)",
+    "AWS.toolkit.lambda.walkthrough.toolInstall.description": "Manage your AWS services and resources with the AWS Command Line Interface (AWS CLI). \n\n[Install AWS CLI](command:aws.toolkit.installAWSCLI)\n\nBuild locally, invoke, and deploy your functions with the Serverless Application Model (SAM) CLI. \n\n[Install SAM CLI](command:aws.toolkit.installSAMCLI)\n\nDocker is an optional, third party tool that assists with local AWS Lambda runtime emulation. Docker is required to invoke Lambda functions on your local machine. \n\n[Install Docker (optional)](command:aws.toolkit.installDocker)\n\nEmulate your AWS cloud services locally with LocalStack to streamline testing in VS Code and CI environments. [Learn more](https://docs.localstack.cloud/aws/). \n\n[Install LocalStack (optional)](command:aws.toolkit.installLocalStack)",
     "AWS.toolkit.lambda.walkthrough.chooseTemplate.title": "Choose your application template",
     "AWS.toolkit.lambda.walkthrough.chooseTemplate.description": "Select a starter application, visually compose an application from scratch, open an existing application, or browse more application examples. \n\nInfrastructure Composer allows you to visually compose modern applications in the cloud. It will define the necessary permissions between resources when you drag a connection between them. \n\n[Initialize your project](command:aws.toolkit.lambda.initializeWalkthroughProject)",
     "AWS.toolkit.lambda.walkthrough.step1.title": "Iterate locally",

packages/core/src/auth/auth.ts

@@ -215,7 +215,7 @@ export class Auth implements AuthService, ConnectionManager {
             const provider = await this.getCredentialsProvider(id, profile)
             await this.authenticate(id, () => this.createCachedCredentials(provider), shouldInvalidate)
 
-            return this.getIamConnection(id, profile)
+            return await this.getIamConnection(id, profile)
         }
     }
 
@@ -253,7 +253,8 @@ export class Auth implements AuthService, ConnectionManager {
         if (profile === undefined) {
             throw new Error(`Connection does not exist: ${id}`)
         }
-        const conn = profile.type === 'sso' ? this.getSsoConnection(id, profile) : this.getIamConnection(id, profile)
+        const conn =
+            profile.type === 'sso' ? this.getSsoConnection(id, profile) : await this.getIamConnection(id, profile)
 
         this.#activeConnection = conn
         this.#onDidChangeActiveConnection.fire(conn)
@@ -705,7 +706,7 @@ export class Auth implements AuthService, ConnectionManager {
         if (profile.type === 'sso') {
             return this.getSsoConnection(id, profile)
         } else {
-            return this.getIamConnection(id, profile)
+            return await this.getIamConnection(id, profile)
         }
     }
 
@@ -805,17 +806,21 @@ export class Auth implements AuthService, ConnectionManager {
         )
     }
 
-    private getIamConnection(
+    private async getIamConnection(
         id: Connection['id'],
         profile: StoredProfile<IamProfile>
-    ): IamConnection & StatefulConnection {
+    ): Promise<IamConnection & StatefulConnection> {
+        // Get the provider to extract the endpoint URL
+        const provider = await this.getCredentialsProvider(id, profile)
+        const endpointUrl = provider.getEndpointUrl?.()
         return {
             id,
             type: 'iam',
             state: profile.metadata.connectionState,
             label:
                 profile.metadata.label ?? (profile.type === 'iam' && profile.subtype === 'linked' ? profile.name : id),
             getCredentials: async () => this.getCredentials(id, await this.getCredentialsProvider(id, profile)),
+            endpointUrl,
         }
     }
 
@@ -832,6 +837,8 @@ export class Auth implements AuthService, ConnectionManager {
             label: profile.metadata?.label ?? this.getSsoProfileLabel(profile),
             getToken: () => this.getToken(id, provider),
             getRegistration: () => provider.getClientRegistration(),
+            // SsoConnection is managed internally in the AWS Toolkit, so the endpointUrl can't be configured
+            endpointUrl: undefined,
         }
     }
 
@@ -856,8 +863,8 @@ export class Auth implements AuthService, ConnectionManager {
     private async createCachedCredentials(provider: CredentialsProvider) {
         const providerId = provider.getCredentialsId()
         globals.loginManager.store.invalidateCredentials(providerId)
-        const { credentials } = await globals.loginManager.store.upsertCredentials(providerId, provider)
-        await globals.loginManager.validateCredentials(credentials, provider.getDefaultRegion())
+        const { credentials, endpointUrl } = await globals.loginManager.store.upsertCredentials(providerId, provider)
+        await globals.loginManager.validateCredentials(credentials, endpointUrl, provider.getDefaultRegion())
 
         return credentials
     }

packages/core/src/auth/connection.ts

@@ -111,6 +111,7 @@ export function createSsoProfile(
 export interface SsoConnection extends SsoProfile {
     readonly id: string
     readonly label: string
+    readonly endpointUrl?: string | undefined
 
     /**
      * Retrieves a bearer token, refreshing or re-authenticating as-needed.
@@ -129,6 +130,7 @@ export interface IamConnection {
     // This may change in the future after refactoring legacy implementations
     readonly id: string
     readonly label: string
+    readonly endpointUrl: string | undefined
     getCredentials(): Promise<Credentials>
 }
 

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

@@ -12,6 +12,7 @@ import { CredentialsProviderManager } from '../providers/credentialsProviderMana
 export interface CachedCredentials {
     credentials: AWS.Credentials
     credentialsHashCode: string
+    endpointUrl?: string
 }
 
 /**
@@ -92,6 +93,7 @@ export class CredentialsStore {
         const credentials = {
             credentials: await credentialsProvider.getCredentials(),
             credentialsHashCode: credentialsProvider.getHashCode(),
+            endpointUrl: credentialsProvider.getEndpointUrl?.(),
         }
 
         this.credentialsCache[asString(credentialsId)] = credentials

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

@@ -12,6 +12,7 @@ export const SharedCredentialsKeys = {
     AWS_SESSION_TOKEN: 'aws_session_token',
     CREDENTIAL_PROCESS: 'credential_process',
     CREDENTIAL_SOURCE: 'credential_source',
+    ENDPOINT_URL: 'endpoint_url',
     REGION: 'region',
     ROLE_ARN: 'role_arn',
     SOURCE_PROFILE: 'source_profile',

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

@@ -21,7 +21,7 @@ import { isValidResponse } from '../../shared/wizards/wizard'
 const credentialsTimeout = 300000 // 5 minutes
 const credentialsProgressDelay = 1000
 
-export function asEnvironmentVariables(credentials: Credentials): NodeJS.ProcessEnv {
+export function asEnvironmentVariables(credentials: Credentials, endpointUrl?: string): NodeJS.ProcessEnv {
     const environmentVariables: NodeJS.ProcessEnv = {}
 
     environmentVariables.AWS_ACCESS_KEY = credentials.accessKeyId
@@ -30,6 +30,9 @@ export function asEnvironmentVariables(credentials: Credentials): NodeJS.Process
     environmentVariables.AWS_SECRET_ACCESS_KEY = credentials.secretAccessKey
     environmentVariables.AWS_SESSION_TOKEN = credentials.sessionToken
     environmentVariables.AWS_SECURITY_TOKEN = credentials.sessionToken
+    if (endpointUrl !== undefined) {
+        environmentVariables.AWS_ENDPOINT_URL = endpointUrl
+    }
 
     return environmentVariables
 }

packages/core/src/auth/deprecated/loginManager.ts

@@ -30,10 +30,11 @@ import { isAutomation } from '../../shared/vscode/env'
 import { Credentials } from '@aws-sdk/types'
 import { ToolkitError } from '../../shared/errors'
 import * as localizedText from '../../shared/localizedText'
-import { DefaultStsClient } from '../../shared/clients/stsClient'
+import { DefaultStsClient, type GetCallerIdentityResponse } from '../../shared/clients/stsClient'
 import { findAsync } from '../../shared/utilities/collectionUtils'
 import { telemetry } from '../../shared/telemetry/telemetry'
 import { withTelemetryContext } from '../../shared/telemetry/util'
+import { localStackConnectionHeader, localStackConnectionString } from '../utils'
 
 const loginManagerClassName = 'LoginManager'
 /**
@@ -65,19 +66,19 @@ export class LoginManager {
 
         try {
             provider = await getProvider(args.providerId)
-
-            const credentials = (await this.store.upsertCredentials(args.providerId, provider))?.credentials
+            const { credentials, endpointUrl } = await this.store.upsertCredentials(args.providerId, provider)
             if (!credentials) {
                 throw new Error(`No credentials found for id ${asString(args.providerId)}`)
             }
 
-            const accountId = await this.validateCredentials(credentials, provider.getDefaultRegion())
+            const accountId = await this.validateCredentials(credentials, endpointUrl, provider.getDefaultRegion())
             this.awsContext.credentialsShim = createCredentialsShim(this.store, args.providerId, credentials)
             await this.awsContext.setCredentials({
                 credentials,
                 accountId: accountId,
                 credentialsId: asString(args.providerId),
                 defaultRegion: provider.getDefaultRegion(),
+                endpointUrl: provider.getEndpointUrl?.(),
             })
 
             telemetryResult = 'Succeeded'
@@ -111,16 +112,40 @@ export class LoginManager {
         }
     }
 
-    public async validateCredentials(credentials: Credentials, region = this.defaultCredentialsRegion) {
-        const stsClient = new DefaultStsClient(region, credentials)
-        const accountId = (await stsClient.getCallerIdentity()).Account
+    public async validateCredentials(
+        credentials: Credentials,
+        endpointUrl?: string,
+        region = this.defaultCredentialsRegion
+    ) {
+        const stsClient = new DefaultStsClient(region, credentials, endpointUrl)
+        const callerIdentity = await stsClient.getCallerIdentity()
+        await this.detectExternalConnection(callerIdentity)
+        // Validate presence of Account Id
+        const accountId = callerIdentity.Account
         if (!accountId) {
+            if (endpointUrl !== undefined) {
+                telemetry.auth_customEndpoint.emit({ source: 'validateCredentials', result: 'Failed' })
+            }
             throw new Error('Could not determine Account Id for credentials')
         }
+        if (endpointUrl !== undefined) {
+            telemetry.auth_customEndpoint.emit({ source: 'validateCredentials', result: 'Succeeded' })
+        }
 
         return accountId
     }
 
+    private async detectExternalConnection(callerIdentity: GetCallerIdentityResponse): Promise<void> {
+        // @ts-ignore
+        const headers = callerIdentity.$response?.httpResponse?.headers
+        if (headers !== undefined && localStackConnectionHeader in headers) {
+            await globals.globalState.update('aws.toolkit.externalConnection', localStackConnectionString)
+            telemetry.auth_localstackEndpoint.emit({ source: 'validateCredentials', result: 'Succeeded' })
+        } else {
+            await globals.globalState.update('aws.toolkit.externalConnection', undefined)
+        }
+    }
+
     /**
      * Removes Credentials from the Toolkit. Essentially the Toolkit becomes "logged out".
      *

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

@@ -112,6 +112,10 @@ export interface CredentialsProvider {
      */
     getTelemetryType(): CredentialType
     getDefaultRegion(): string | undefined
+    /**
+     * Gets the endpoint URL configured for this profile, if any.
+     */
+    getEndpointUrl?(): string | undefined
     getHashCode(): string
     getCredentials(): Promise<AWS.Credentials>
     /**

packages/core/src/auth/providers/envVarsCredentialsProvider.ts

@@ -61,4 +61,9 @@ export class EnvVarsCredentialsProvider implements CredentialsProvider {
         }
         return this.credentials
     }
+
+    public getEndpointUrl(): string | undefined {
+        const env = process.env as EnvironmentVariables
+        return env.AWS_ENDPOINT_URL?.toString()
+    }
 }

packages/core/src/auth/providers/sharedCredentialsProvider.ts

@@ -105,6 +105,10 @@ export class SharedCredentialsProvider implements CredentialsProvider {
         return this.profile[SharedCredentialsKeys.REGION]
     }
 
+    public getEndpointUrl(): string | undefined {
+        return this.profile[SharedCredentialsKeys.ENDPOINT_URL]?.trim()
+    }
+
     public async canAutoConnect(): Promise<boolean> {
         if (isSsoProfile(this.profile)) {
             const tokenProvider = SsoAccessTokenProvider.create({