From 43f97f4f56a52082612261b1398cbbfe153952ac Mon Sep 17 00:00:00 2001 From: Cheena Malhotra <13396919+cheenamalhotra@users.noreply.github.com> Date: Wed, 22 Mar 2023 22:23:03 -0700 Subject: [PATCH] Handle out of sync extension activations for encryption keys updated event (#22415) --- .../azureAccountProviderService.ts | 11 ++++++++ .../utils/fileEncryptionHelper.ts | 16 ++++++++--- .../account-provider/utils/msalCachePlugin.ts | 8 ++++++ extensions/azurecore/src/azurecore.d.ts | 4 +++ extensions/azurecore/src/extension.ts | 8 +++++- extensions/machine-learning/src/test/stubs.ts | 4 ++- extensions/mssql/src/sqlToolsServer.ts | 28 ++++++++++++------- 7 files changed, 63 insertions(+), 16 deletions(-) diff --git a/extensions/azurecore/src/account-provider/azureAccountProviderService.ts b/extensions/azurecore/src/account-provider/azureAccountProviderService.ts index 2a239794df..a71a0673c1 100644 --- a/extensions/azurecore/src/account-provider/azureAccountProviderService.ts +++ b/extensions/azurecore/src/account-provider/azureAccountProviderService.ts @@ -81,6 +81,13 @@ export class AzureAccountProviderService implements vscode.Disposable { return this._onEncryptionKeysUpdated; } + public async getEncryptionKeys(): Promise { + if (!this._cachePluginProvider) { + await this.onDidChangeConfiguration(); + } + return this._cachePluginProvider!.getCacheEncryptionKeys(); + } + public dispose() { while (this._disposables.length) { const item = this._disposables.pop(); @@ -167,6 +174,10 @@ export class AzureAccountProviderService implements vscode.Disposable { // MSAL Cache Plugin this._cachePluginProvider = new MsalCachePluginProvider(tokenCacheKeyMsal, this._userStoragePath, this._credentialProvider, this._onEncryptionKeysUpdated); + if (this._authLibrary === Constants.AuthLibrary.MSAL) { + // Initialize cache provider and encryption keys + await this._cachePluginProvider.init(); + } const msalConfiguration: Configuration = { auth: { diff --git a/extensions/azurecore/src/account-provider/utils/fileEncryptionHelper.ts b/extensions/azurecore/src/account-provider/utils/fileEncryptionHelper.ts index 95554e675d..c4409d5cee 100644 --- a/extensions/azurecore/src/account-provider/utils/fileEncryptionHelper.ts +++ b/extensions/azurecore/src/account-provider/utils/fileEncryptionHelper.ts @@ -53,10 +53,18 @@ export class FileEncryptionHelper { // Emit event with cache encryption keys to send notification to provider services. if (this._authLibrary === AuthLibrary.MSAL && this._onEncryptionKeysUpdated) { - this._onEncryptionKeysUpdated.fire({ - iv: this._ivBuffer.toString(this._bufferEncoding), - key: this._keyBuffer.toString(this._bufferEncoding) - }); + this._onEncryptionKeysUpdated.fire(this.getEncryptionKeys()); + Logger.verbose('FileEncryptionHelper: Fired encryption keys updated event.'); + } + } + + /** + * Provides encryption keys in use for instant access. + */ + public getEncryptionKeys(): CacheEncryptionKeys { + return { + iv: this._ivBuffer!.toString(this._bufferEncoding), + key: this._keyBuffer!.toString(this._bufferEncoding) } } diff --git a/extensions/azurecore/src/account-provider/utils/msalCachePlugin.ts b/extensions/azurecore/src/account-provider/utils/msalCachePlugin.ts index 1e1642a6d0..12b5b6442d 100644 --- a/extensions/azurecore/src/account-provider/utils/msalCachePlugin.ts +++ b/extensions/azurecore/src/account-provider/utils/msalCachePlugin.ts @@ -33,6 +33,14 @@ export class MsalCachePluginProvider { return this._msalFilePath + '.lockfile'; } + public async init(): Promise { + await this._fileEncryptionHelper.init(); + } + + public getCacheEncryptionKeys(): CacheEncryptionKeys { + return this._fileEncryptionHelper.getEncryptionKeys(); + } + public getCachePlugin(): ICachePlugin { const lockFilePath = this.getLockfilePath(); const beforeCacheAccess = async (cacheContext: TokenCacheContext): Promise => { diff --git a/extensions/azurecore/src/azurecore.d.ts b/extensions/azurecore/src/azurecore.d.ts index 3abc1dab9f..dff4de72f0 100644 --- a/extensions/azurecore/src/azurecore.d.ts +++ b/extensions/azurecore/src/azurecore.d.ts @@ -321,6 +321,10 @@ declare module 'azurecore' { * by connection providers to read/write to the same access token cache for stable connectivity. */ onEncryptionKeysUpdated: vscode.Event; + /** + * Fetches MSAL cache encryption keys currently in use. + */ + getEncryptionKeys(): Promise; } export type GetSubscriptionsResult = { subscriptions: azureResource.AzureResourceSubscription[], errors: Error[] }; diff --git a/extensions/azurecore/src/extension.ts b/extensions/azurecore/src/extension.ts index 171288290e..bd277eb043 100644 --- a/extensions/azurecore/src/extension.ts +++ b/extensions/azurecore/src/extension.ts @@ -241,7 +241,13 @@ export async function activate(context: vscode.ExtensionContext): Promise> { return azureResourceUtils.runResourceQuery(account, subscriptions, ignoreErrors, query); }, - onEncryptionKeysUpdated: eventEmitter!.event + onEncryptionKeysUpdated: eventEmitter!.event, + async getEncryptionKeys(): Promise { + if (!providerService) { + throw new Error("Failed to initialize Azure account provider."); + } + return await providerService!.getEncryptionKeys(); + } }; } diff --git a/extensions/machine-learning/src/test/stubs.ts b/extensions/machine-learning/src/test/stubs.ts index 8e7d3b4a7c..483bae7352 100644 --- a/extensions/machine-learning/src/test/stubs.ts +++ b/extensions/machine-learning/src/test/stubs.ts @@ -62,5 +62,7 @@ export class AzurecoreApiStub implements azurecore.IExtension { throw new Error('Method not implemented.'); } onEncryptionKeysUpdated: any - + getEncryptionKeys(): Promise { + throw new Error('Method not implemented.'); + } } diff --git a/extensions/mssql/src/sqlToolsServer.ts b/extensions/mssql/src/sqlToolsServer.ts index 8683fbd1c0..898dcd407b 100644 --- a/extensions/mssql/src/sqlToolsServer.ts +++ b/extensions/mssql/src/sqlToolsServer.ts @@ -10,7 +10,7 @@ import * as vscode from 'vscode'; import * as azdata from 'azdata'; import * as path from 'path'; import * as azurecore from 'azurecore'; -import { getAzureAuthenticationLibraryConfig, getCommonLaunchArgsAndCleanupOldLogFiles, getConfigTracingLevel, getEnableSqlAuthenticationProviderConfig, getOrDownloadServer, getParallelMessageProcessingConfig, TracingLevel } from './utils'; +import { getAzureAuthenticationLibraryConfig, getCommonLaunchArgsAndCleanupOldLogFiles, getConfigTracingLevel, getEnableSqlAuthenticationProviderConfig, getOrDownloadServer, getParallelMessageProcessingConfig, logDebug, TracingLevel } from './utils'; import { TelemetryReporter, LanguageClientErrorHandler } from './telemetry'; import { SqlOpsDataClient, ClientOptions } from 'dataprotocol-client'; import { TelemetryFeature, AgentServicesFeature, SerializationFeature, AccountFeature, SqlAssessmentServicesFeature, ProfilerFeature, TableDesignerFeature, ExecutionPlanServiceFeature } from './features'; @@ -20,7 +20,7 @@ import { SchemaCompareService } from './schemaCompare/schemaCompareService'; import { AppContext } from './appContext'; import { DacFxService } from './dacfx/dacFxService'; import { CmsService } from './cms/cmsService'; -import { CompletionExtensionParams, CompletionExtLoadRequest, DidChangeEncryptionIVKeyParams, EncryptionKeysChangedNotification } from './contracts'; +import { CompletionExtensionParams, CompletionExtLoadRequest, EncryptionKeysChangedNotification } from './contracts'; import { promises as fs } from 'fs'; import * as nls from 'vscode-nls'; import { LanguageExtensionService } from './languageExtension/languageExtensionService'; @@ -100,16 +100,24 @@ export class SqlToolsServer { */ private async handleEncryptionKeyEventNotification(client: SqlOpsDataClient) { if (getAzureAuthenticationLibraryConfig() === 'MSAL' && getEnableSqlAuthenticationProviderConfig()) { - let onDidEncryptionKeysChanged = (await this.getAzureCoreAPI()).onEncryptionKeysUpdated; - // Register event listener from Azure Core extension + let azureCoreApi = await this.getAzureCoreAPI(); + let onDidEncryptionKeysChanged = azureCoreApi.onEncryptionKeysUpdated; + // Register event listener from Azure Core extension and + // send client notification for updated encryption keys onDidEncryptionKeysChanged((keys: azurecore.CacheEncryptionKeys) => { - // Send client notification for updated encryption keys - client.sendNotification(EncryptionKeysChangedNotification.type, - { - key: keys.key, - iv: keys.iv - }); + client.sendNotification(EncryptionKeysChangedNotification.type, keys); }); + + try { + // Fetch encryption keys directly from AzureCore as notification event may not fire again + // if Azure Core extension was activated before. + const keys = await azureCoreApi.getEncryptionKeys(); + client.sendNotification(EncryptionKeysChangedNotification.type, keys); + } + catch (e) { + console.error(`An error occurred when fetching encryption keys: ${e}`); + } + logDebug('SqlToolsServer: Registered encryption key event handler.'); } }