Move encryption to LanguageClientAuth and encrypt profile data

Author: liramon1

packages/core/src/auth/auth2.ts

@@ -19,7 +19,6 @@ import {
     invalidateSsoTokenRequestType,
     invalidateStsCredentialRequestType,
     ProfileKind,
-    UpdateProfileParams,
     updateProfileRequestType,
     SsoTokenChangedParams,
     StsCredentialChangedParams,
@@ -50,8 +49,8 @@ import {
     Profile,
     SsoSession,
     GetMfaCodeParams,
-    GetMfaCodeResult,
     getMfaCodeRequestType,
+    GetMfaCodeResult,
 } from '@aws/language-server-runtimes/protocol'
 import { LanguageClient } from 'vscode-languageclient'
 import { getLogger } from '../shared/logger/logger'
@@ -125,12 +124,31 @@ export class LanguageClientAuth {
         return this.#ssoCacheWatcher
     }
 
-    getSsoToken(
+    /**
+     * Encrypts an object
+     */
+    private async encrypt<T>(request: T): Promise<string> {
+        const payload = new TextEncoder().encode(JSON.stringify(request))
+        const encrypted = await new jose.CompactEncrypt(payload)
+            .setProtectedHeader({ alg: 'dir', enc: 'A256GCM' })
+            .encrypt(this.encryptionKey)
+        return encrypted
+    }
+
+    /**
+     * Decrypts an object
+     */
+    private async decrypt<T>(request: string): Promise<T> {
+        const result = await jose.compactDecrypt(request, this.encryptionKey)
+        return JSON.parse(new TextDecoder().decode(result.plaintext)) as T
+    }
+
+    async getSsoToken(
         tokenSource: TokenSource,
         login: boolean = false,
         cancellationToken?: CancellationToken
     ): Promise<GetSsoTokenResult> {
-        return this.client.sendRequest(
+        const response: GetSsoTokenResult = await this.client.sendRequest(
             getSsoTokenRequestType.method,
             {
                 clientName: this.clientName,
@@ -142,14 +160,17 @@ export class LanguageClientAuth {
             } satisfies GetSsoTokenParams,
             cancellationToken
         )
+        // Decrypt the access token
+        response.ssoToken.accessToken = await this.decrypt(response.ssoToken.accessToken)
+        return response
     }
 
-    getIamCredential(
+    async getIamCredential(
         profileName: string,
         login: boolean = false,
         cancellationToken?: CancellationToken
     ): Promise<GetIamCredentialResult> {
-        return this.client.sendRequest(
+        const response: GetIamCredentialResult = await this.client.sendRequest(
             getIamCredentialRequestType.method,
             {
                 profileName: profileName,
@@ -159,16 +180,25 @@ export class LanguageClientAuth {
             } satisfies GetIamCredentialParams,
             cancellationToken
         )
+        // Decrypt the response credentials
+        const { accessKeyId, secretAccessKey, sessionToken, expiration } = response.credential.credentials
+        response.credential.credentials = {
+            accessKeyId: await this.decrypt(accessKeyId),
+            secretAccessKey: await this.decrypt(secretAccessKey),
+            sessionToken: sessionToken ? await this.decrypt(sessionToken) : undefined,
+            expiration: expiration,
+        }
+        return response
     }
 
-    updateSsoProfile(
+    async updateSsoProfile(
         profileName: string,
         startUrl: string,
         region: string,
         scopes: string[]
     ): Promise<UpdateProfileResult> {
         // Add SSO settings and delete credentials from profile
-        return this.client.sendRequest(updateProfileRequestType.method, {
+        const params = await this.encrypt({
             profile: {
                 kinds: [ProfileKind.SsoTokenProfile],
                 name: profileName,
@@ -188,10 +218,11 @@ export class LanguageClientAuth {
                     sso_registration_scopes: scopes,
                 },
             },
-        } satisfies UpdateProfileParams)
+        })
+        return this.client.sendRequest(updateProfileRequestType.method, params)
     }
 
-    updateIamProfile(profileName: string, opts: IamProfileOptions): Promise<UpdateProfileResult> {
+    async updateIamProfile(profileName: string, opts: IamProfileOptions): Promise<UpdateProfileResult> {
         // Substitute missing fields for defaults
         const fields = { ...IamProfileOptionsDefaults, ...opts }
         // Get the profile kind matching the provided fields
@@ -204,7 +235,7 @@ export class LanguageClientAuth {
             kind = ProfileKind.Unknown
         }
 
-        return this.client.sendRequest(updateProfileRequestType.method, {
+        const params = await this.encrypt({
             profile: {
                 kinds: [kind],
                 name: profileName,
@@ -217,10 +248,12 @@ export class LanguageClientAuth {
                 },
             },
         })
+        return this.client.sendRequest(updateProfileRequestType.method, params)
     }
 
-    listProfiles() {
-        return this.client.sendRequest(listProfilesRequestType.method, {}) as Promise<ListProfilesResult>
+    async listProfiles() {
+        const response: string = await this.client.sendRequest(listProfilesRequestType.method, {})
+        return await this.decrypt<ListProfilesResult>(response)
     }
 
     /**
@@ -352,19 +385,6 @@ export abstract class BaseLogin {
             this.eventEmitter.fire({ id: this.profileName, state: this.connectionState })
         }
     }
-
-    /**
-     * Decrypts an encrypted string, removes its quotes, and returns the resulting string
-     */
-    protected async decrypt(encrypted: string): Promise<string> {
-        try {
-            const decrypted = await jose.compactDecrypt(encrypted, this.lspAuth.encryptionKey)
-            return decrypted.plaintext.toString().replaceAll('"', '')
-        } catch (e) {
-            getLogger().error(`Failed to decrypt: ${encrypted}`)
-            return encrypted
-        }
-    }
 }
 
 /**
@@ -436,9 +456,8 @@ export class SsoLogin extends BaseLogin {
      */
     async getCredential() {
         const response = await this._getSsoToken(false)
-        const accessToken = await this.decrypt(response.ssoToken.accessToken)
         return {
-            credential: accessToken,
+            credential: response.ssoToken.accessToken,
             updateCredentialsParams: response.updateCredentialsParams,
         }
     }
@@ -562,7 +581,7 @@ export class IamLogin extends BaseLogin {
                 sourceProfile: sourceProfile,
             })
         } else {
-            // Create the target profile
+            // Create the credentials profile
             await this.lspAuth.updateIamProfile(this.profileName, {
                 accessKey: opts.accessKey,
                 secretKey: opts.secretKey,
@@ -575,14 +594,6 @@ export class IamLogin extends BaseLogin {
      * Restore the connection state and connection details to memory, if they exist.
      */
     async restore() {
-        const sessionData = await this.getProfile()
-        const credentials = sessionData?.profile?.settings
-        if (credentials?.aws_access_key_id && credentials?.aws_secret_access_key) {
-            this._data = {
-                accessKey: credentials.aws_access_key_id,
-                secretKey: credentials.aws_secret_access_key,
-            }
-        }
         try {
             await this._getIamCredential(false)
         } catch (err) {
@@ -596,15 +607,8 @@ export class IamLogin extends BaseLogin {
      */
     async getCredential() {
         const response = await this._getIamCredential(false)
-        const credentials: IamCredentials = {
-            accessKeyId: await this.decrypt(response.credential.credentials.accessKeyId),
-            secretAccessKey: await this.decrypt(response.credential.credentials.secretAccessKey),
-            sessionToken: response.credential.credentials.sessionToken
-                ? await this.decrypt(response.credential.credentials.sessionToken)
-                : undefined,
-        }
         return {
-            credential: credentials,
+            credential: response.credential.credentials,
             updateCredentialsParams: response.updateCredentialsParams,
         }
     }
@@ -639,7 +643,7 @@ export class IamLogin extends BaseLogin {
         }
 
         // Update cached credentials and credential ID
-        if (response.credential?.credentials?.accessKeyId && response.credential?.credentials?.secretAccessKey) {
+        if (response.credential.credentials.accessKeyId && response.credential.credentials.secretAccessKey) {
             this._data = {
                 accessKey: response.credential.credentials.accessKeyId,
                 secretKey: response.credential.credentials.secretAccessKey,