Persist auth in memory, removing loading credentials from SDK to support SSO

Author: satyakigh

src/auth/AuthProtocol.ts

@@ -6,6 +6,7 @@ import {
 } from 'vscode-languageserver';
 import {
     UpdateCredentialsParams,
+    UpdateCredentialsResult,
     ConnectionMetadata,
     ListProfilesParams,
     ListProfilesResult,
@@ -22,7 +23,9 @@ import {
 export const IamCredentialsUpdateRequest = Object.freeze({
     method: 'aws/credentials/iam/update' as const,
     messageDirection: MessageDirection.clientToServer,
-    type: new ProtocolRequestType<UpdateCredentialsParams, void, never, void, void>('aws/credentials/iam/update'),
+    type: new ProtocolRequestType<UpdateCredentialsParams, UpdateCredentialsResult, never, void, void>(
+        'aws/credentials/iam/update',
+    ),
 } as const);
 
 export const BearerCredentialsUpdateRequest = Object.freeze({

src/auth/AwsCredentials.ts

@@ -1,12 +1,10 @@
-import { AwsCredentialIdentity } from '@aws-sdk/types';
 import { DeepReadonly } from 'ts-essentials';
 import { LspAuthHandlers } from '../protocol/LspAuthHandlers';
 import { DefaultSettings } from '../settings/Settings';
 import { SettingsManager } from '../settings/SettingsManager';
-import { parseProfile } from '../settings/SettingsParser';
-import { ClientMessage } from '../telemetry/ClientMessage';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
 import { extractErrorMessage } from '../utils/Errors';
+import { getRegion } from '../utils/Region';
 import { parseWithPrettyError } from '../utils/ZodErrorWrapper';
 import {
     parseListProfilesResult,
@@ -30,31 +28,33 @@ import {
     InvalidateSsoTokenParams,
     InvalidateSsoTokenResult,
     SsoTokenChangedParams,
+    IamCredentials,
 } from './AwsLspAuthTypes';
-import { sdkIAMCredentials } from './AwsSdkCredentialsProvider';
 
 export class AwsCredentials {
     private readonly logger = LoggerFactory.getLogger(AwsCredentials);
-    private profileName = DefaultSettings.profile.profile;
 
+    private iamCredentials?: IamCredentials;
     private bearerCredentials?: BearerCredentials;
     private connectionMetadata?: ConnectionMetadata;
 
     constructor(
         private readonly awsHandlers: LspAuthHandlers,
         private readonly settingsManager: SettingsManager,
-        private readonly clientMessage: ClientMessage,
-        private readonly getIAMFromSdk: (
-            profile: string,
-        ) => Promise<DeepReadonly<AwsCredentialIdentity>> = sdkIAMCredentials,
     ) {}
 
-    getIAM(): Promise<DeepReadonly<AwsCredentialIdentity>> {
-        return this.getIAMFromSdk(this.profileName);
+    getIAM(): DeepReadonly<IamCredentials> {
+        if (!this.iamCredentials) {
+            throw new Error('IAM credentials not configured');
+        }
+        return structuredClone(this.iamCredentials);
     }
 
-    getBearer(): DeepReadonly<BearerCredentials | undefined> {
-        return this.bearerCredentials;
+    getBearer(): DeepReadonly<BearerCredentials> {
+        if (!this.bearerCredentials) {
+            throw new Error('Bearer credentials not configured');
+        }
+        return structuredClone(this.bearerCredentials);
     }
 
     getConnectionMetadata(): ConnectionMetadata | undefined {
@@ -137,29 +137,28 @@ export class AwsCredentials {
         }
     }
 
-    handleIamCredentialsUpdate(params: UpdateCredentialsParams) {
-        let newProfileName = DefaultSettings.profile.profile;
-        let newRegion = DefaultSettings.profile.region;
+    handleIamCredentialsUpdate(params: UpdateCredentialsParams): boolean {
         try {
             const { data } = parseWithPrettyError(parseUpdateCredentialsParams, params);
             if ('accessKeyId' in data) {
-                const profile = parseWithPrettyError(
-                    parseProfile,
-                    {
-                        profile: data.profile,
-                        region: data.region,
-                    },
-                    DefaultSettings.profile,
-                );
-
-                newProfileName = profile.profile;
-                newRegion = profile.region;
+                const region = getRegion(data.region);
+
+                this.iamCredentials = {
+                    ...data,
+                    region,
+                };
+
+                this.settingsManager.updateProfileSettings(data.profile, region);
+                return true;
             }
+
+            throw new Error('Not an IAM credential');
         } catch (error) {
-            this.logger.error(`Failed to update IAM profile: ${extractErrorMessage(error)}`);
-        } finally {
-            this.profileName = newProfileName;
-            this.settingsManager.updateProfileSettings(newProfileName, newRegion);
+            this.iamCredentials = undefined;
+
+            this.logger.error(`Failed to update IAM credentials: ${extractErrorMessage(error)}`);
+            this.settingsManager.updateProfileSettings(DefaultSettings.profile.profile, DefaultSettings.profile.region);
+            return false;
         }
     }
 
@@ -183,7 +182,7 @@ export class AwsCredentials {
 
     handleIamCredentialsDelete() {
         this.logger.info('IAM credentials deleted');
-        this.profileName = DefaultSettings.profile.profile;
+        this.iamCredentials = undefined;
     }
 
     handleBearerCredentialsDelete() {

src/auth/AwsCredentialsParser.ts

@@ -19,11 +19,11 @@ const ProfileKindSchema = z.enum([
 ]);
 
 const IamCredentialsSchema = z.object({
-    profile: z.string().optional(),
+    profile: z.string({ message: 'Profile name is required' }),
     accessKeyId: z.string({ message: 'Access key ID is required' }),
     secretAccessKey: z.string({ message: 'Secret access key is required' }),
     sessionToken: z.string().optional(),
-    region: z.string().optional(),
+    region: z.string({ message: 'Region is required' }),
 });
 
 const BearerCredentialsSchema = z.object({

src/auth/AwsLspAuthTypes.ts

@@ -1,12 +1,10 @@
 /* eslint-disable @typescript-eslint/no-empty-object-type */
 
-// Must be kept consistent with https://github.com/aws/language-server-runtimes
-export type IamCredentials = {
-    profile?: string;
-    accessKeyId: string;
-    secretAccessKey: string;
-    sessionToken?: string;
-    region?: string;
+import { AwsCredentialIdentity } from '@aws-sdk/types';
+
+export type IamCredentials = AwsCredentialIdentity & {
+    profile: string;
+    region: string;
 };
 
 export type BearerCredentials = {
@@ -28,6 +26,10 @@ export type UpdateCredentialsParams = {
     encrypted?: boolean;
 };
 
+export type UpdateCredentialsResult = {
+    success: boolean;
+};
+
 export type ProfileKind =
     | 'Unknown'
     | 'SsoTokenProfile'

src/handlers/AuthHandler.ts

@@ -1,14 +1,13 @@
 import { RequestHandler, NotificationHandler } from 'vscode-languageserver/node';
-import { UpdateCredentialsParams, SsoTokenChangedParams } from '../auth/AwsLspAuthTypes';
+import { UpdateCredentialsParams, UpdateCredentialsResult, SsoTokenChangedParams } from '../auth/AwsLspAuthTypes';
 import { ServerComponents } from '../server/ServerComponents';
 
 export function iamCredentialsUpdateHandler(
     components: ServerComponents,
-): RequestHandler<UpdateCredentialsParams, void, void> {
-    return (params: UpdateCredentialsParams) => {
-        // AwsCredentials.handleIamCredentialsUpdate already calls settingsManager.updateProfileSettings
-        // which will notify all subscribed components via the observable pattern
-        components.awsCredentials.handleIamCredentialsUpdate(params);
+): RequestHandler<UpdateCredentialsParams, UpdateCredentialsResult, void> {
+    return (params: UpdateCredentialsParams): UpdateCredentialsResult => {
+        const success = components.awsCredentials.handleIamCredentialsUpdate(params);
+        return { success };
     };
 }
 

src/protocol/LspAuthHandlers.ts

@@ -17,6 +17,7 @@ import {
     ListProfilesParams,
     SsoTokenChangedParams,
     UpdateCredentialsParams,
+    UpdateCredentialsResult,
     UpdateProfileParams,
 } from '../auth/AwsLspAuthTypes';
 
@@ -27,7 +28,7 @@ export class LspAuthHandlers {
     // ========================================
     // RECEIVE: Client → Server
     // ========================================
-    onIamCredentialsUpdate(handler: RequestHandler<UpdateCredentialsParams, void, void>) {
+    onIamCredentialsUpdate(handler: RequestHandler<UpdateCredentialsParams, UpdateCredentialsResult, void>) {
         this.connection.onRequest(IamCredentialsUpdateRequest.type, handler);
     }
 

src/server/CfnExternal.ts

@@ -63,7 +63,7 @@ export class CfnExternal implements Configurables, Closeable {
     }
 
     configurables(): Configurable[] {
-        return [this.schemaTaskManager, this.schemaRetriever, this.awsClient, this.cfnLintService, this.guardService];
+        return [this.schemaTaskManager, this.schemaRetriever, this.cfnLintService, this.guardService];
     }
 
     async close() {

src/server/CfnInfraCore.ts

@@ -51,8 +51,7 @@ export class CfnInfraCore implements Configurables, Closeable {
         this.fileContextManager = overrides.fileContextManager ?? new FileContextManager(this.documentManager);
 
         this.awsCredentials =
-            overrides.awsCredentials ??
-            new AwsCredentials(lspComponents.authHandlers, this.settingsManager, this.clientMessage);
+            overrides.awsCredentials ?? new AwsCredentials(lspComponents.authHandlers, this.settingsManager);
 
         this.diagnosticCoordinator =
             overrides.diagnosticCoordinator ?? new DiagnosticCoordinator(lspComponents.diagnostics);

src/services/AwsClient.ts

@@ -1,46 +1,32 @@
 import { CloudControlClient } from '@aws-sdk/client-cloudcontrol';
 import { CloudFormationClient } from '@aws-sdk/client-cloudformation';
-import { AwsCredentialIdentity } from '@aws-sdk/types';
 import { AwsCredentials } from '../auth/AwsCredentials';
-import { SettingsConfigurable, ISettingsSubscriber, SettingsSubscription } from '../settings/ISettingsSubscriber';
-import { DefaultSettings } from '../settings/Settings';
+import { IamCredentials } from '../auth/AwsLspAuthTypes';
 import { ExtensionId, ExtensionVersion } from '../utils/ExtensionConfig';
 
-export class AwsClient implements SettingsConfigurable {
-    constructor(
-        private readonly credentialsProvider: AwsCredentials,
-        private region: string = DefaultSettings.profile.region,
-        private settingsSubscription?: SettingsSubscription,
-    ) {}
+type IamClientConfig = {
+    region: string;
+    credentials: IamCredentials;
+    customUserAgent: string;
+};
 
-    configure(settingsManager: ISettingsSubscriber): void {
-        if (this.settingsSubscription) {
-            this.settingsSubscription.unsubscribe();
-        }
-
-        this.settingsSubscription = settingsManager.subscribe('profile', (newProfile) => {
-            this.region = newProfile.region;
-        });
-    }
+export class AwsClient {
+    constructor(private readonly credentialsProvider: AwsCredentials) {}
 
     // By default, clients will retry on throttling exceptions 3 times
-    public async getCloudFormationClient() {
-        return new CloudFormationClient(await this.iamClientConfig());
+    public getCloudFormationClient() {
+        return new CloudFormationClient(this.iamClientConfig());
     }
 
-    public async getCloudControlClient() {
-        return new CloudControlClient(await this.iamClientConfig());
+    public getCloudControlClient() {
+        return new CloudControlClient(this.iamClientConfig());
     }
 
-    private async iamClientConfig(): Promise<{
-        region: string;
-        credentials: AwsCredentialIdentity;
-        customUserAgent: string;
-    }> {
-        const data = await this.credentialsProvider.getIAM();
+    private iamClientConfig(): IamClientConfig {
+        const credential = this.credentialsProvider.getIAM();
         return {
-            region: this.region,
-            credentials: data,
+            region: credential.region,
+            credentials: this.credentialsProvider.getIAM(),
             customUserAgent: `${ExtensionId}/${ExtensionVersion}`,
         };
     }

src/services/CcapiService.ts

@@ -12,7 +12,7 @@ export class CcapiService {
     constructor(private readonly awsClient: AwsClient) {}
 
     private async withClient<T>(request: (client: CloudControlClient) => Promise<T>): Promise<T> {
-        const client = await this.awsClient.getCloudControlClient();
+        const client = this.awsClient.getCloudControlClient();
         return await request(client);
     }