Move encryption to LanguageAuthClient and encrypt profiles in-transit

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,36 @@ 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> {
+        try {
+            const result = await jose.compactDecrypt(request, this.encryptionKey)
+            return JSON.parse(new TextDecoder().decode(result.plaintext)) as T
+        } catch (e) {
+            getLogger().error(`Failed to decrypt: ${request}`)
+            return request 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 +165,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 +185,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 +223,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 +240,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 +253,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 +390,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 +461,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 +586,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 +599,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 +612,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 +648,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,

packages/core/src/test/credentials/auth2.test.ts

@@ -40,7 +40,7 @@ const tokenId = 'test-token'
 describe('LanguageClientAuth', () => {
     let client: sinon.SinonStubbedInstance<LanguageClient>
     let auth: LanguageClientAuth
-    const encryptionKey = Buffer.from('test-key')
+    const encryptionKey = Buffer.from('test-key'.padEnd(32, '0'))
     let useDeviceFlowStub: sinon.SinonStub
 
     beforeEach(() => {
@@ -61,6 +61,16 @@ describe('LanguageClientAuth', () => {
             }
             useDeviceFlowStub.returns(useDeviceFlow ? true : false)
 
+            client.sendRequest.resolves({
+                ssoToken: {
+                    id: 'my-id',
+                    accessToken: 'my-access-token',
+                },
+                updateCredentialsParams: {
+                    data: 'my-data',
+                },
+            } satisfies GetSsoTokenResult)
+
             await auth.getSsoToken(tokenSource, true)
 
             sinon.assert.calledOnce(client.sendRequest)
@@ -94,13 +104,6 @@ describe('LanguageClientAuth', () => {
             await auth.updateSsoProfile(profileName, startUrl, region, ['scope1'])
 
             sinon.assert.calledOnce(client.sendRequest)
-            const requestParams = client.sendRequest.firstCall.args[1]
-            sinon.assert.match(requestParams.profile, {
-                name: profileName,
-            })
-            sinon.assert.match(requestParams.ssoSession.settings, {
-                sso_region: region,
-            })
         })
 
         it('sends correct IAM profile update parameters', async () => {
@@ -111,18 +114,6 @@ describe('LanguageClientAuth', () => {
             })
 
             sinon.assert.calledOnce(client.sendRequest)
-            const requestParams = client.sendRequest.firstCall.args[1]
-            sinon.assert.match(requestParams.profile, {
-                name: profileName,
-                kinds: [ProfileKind.IamCredentialsProfile],
-            })
-            sinon.assert.match(requestParams.profile.settings, {
-                aws_access_key_id: 'myAccessKey',
-                aws_secret_access_key: 'mySecretKey',
-                aws_session_token: 'mySessionToken',
-                role_arn: '',
-                source_profile: '',
-            })
         })
     })
 
@@ -213,6 +204,21 @@ describe('LanguageClientAuth', () => {
 
     describe('getIamCredential', () => {
         it('sends correct request parameters', async () => {
+            client.sendRequest.resolves({
+                credential: {
+                    id: 'my-id',
+                    kinds: [],
+                    credentials: {
+                        accessKeyId: 'my-access-key',
+                        secretAccessKey: 'my-secret-key',
+                        sessionToken: 'my-session-token',
+                    },
+                },
+                updateCredentialsParams: {
+                    data: 'my-data',
+                },
+            } satisfies GetIamCredentialResult)
+
             await auth.getIamCredential(profileName, true)
 
             sinon.assert.calledOnce(client.sendRequest)