Handle out of sync extension activations for encryption keys updated event (#22415)

extensions/azurecore/src/account-provider/azureAccountProviderService.ts

@@ -81,6 +81,13 @@ export class AzureAccountProviderService implements vscode.Disposable {
 		return this._onEncryptionKeysUpdated;
 	}
 
+	public async getEncryptionKeys(): Promise<CacheEncryptionKeys> {
+		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: {

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({
+			this._onEncryptionKeysUpdated.fire(this.getEncryptionKeys());
-				iv: this._ivBuffer.toString(this._bufferEncoding),
+			Logger.verbose('FileEncryptionHelper: Fired encryption keys updated event.');
-				key: this._keyBuffer.toString(this._bufferEncoding)
+		}
-			});
+	}
+
+	/**
+	 * Provides encryption keys in use for instant access.
+	 */
+	public getEncryptionKeys(): CacheEncryptionKeys {
+		return {
+			iv: this._ivBuffer!.toString(this._bufferEncoding),
+			key: this._keyBuffer!.toString(this._bufferEncoding)
 		}
 	}
 

extensions/azurecore/src/account-provider/utils/msalCachePlugin.ts

@@ -33,6 +33,14 @@ export class MsalCachePluginProvider {
 		return this._msalFilePath + '.lockfile';
 	}
 
+	public async init(): Promise<void> {
+		await this._fileEncryptionHelper.init();
+	}
+
+	public getCacheEncryptionKeys(): CacheEncryptionKeys {
+		return this._fileEncryptionHelper.getEncryptionKeys();
+	}
+
 	public getCachePlugin(): ICachePlugin {
 		const lockFilePath = this.getLockfilePath();
 		const beforeCacheAccess = async (cacheContext: TokenCacheContext): Promise<void> => {

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<CacheEncryptionKeys>;
+		/**
+		 * Fetches MSAL cache encryption keys currently in use.
+		 */
+		getEncryptionKeys(): Promise<CacheEncryptionKeys>;
 	}
 
 	export type GetSubscriptionsResult = { subscriptions: azureResource.AzureResourceSubscription[], errors: Error[] };

extensions/azurecore/src/extension.ts

@@ -241,7 +241,13 @@ export async function activate(context: vscode.ExtensionContext): Promise<azurec
 			query: string): Promise<azurecore.ResourceQueryResult<T>> {
 			return azureResourceUtils.runResourceQuery(account, subscriptions, ignoreErrors, query);
 		},
-		onEncryptionKeysUpdated: eventEmitter!.event
+		onEncryptionKeysUpdated: eventEmitter!.event,
+		async getEncryptionKeys(): Promise<azurecore.CacheEncryptionKeys> {
+			if (!providerService) {
+				throw new Error("Failed to initialize Azure account provider.");
+			}
+			return await providerService!.getEncryptionKeys();
+		}
 	};
 }
 

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<azurecore.CacheEncryptionKeys> {
+		throw new Error('Method not implemented.');
+	}
 }

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;
+			let azureCoreApi = await this.getAzureCoreAPI();
-			// Register event listener from Azure Core extension
+			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, keys);
-				client.sendNotification(EncryptionKeysChangedNotification.type,
-					<DidChangeEncryptionIVKeyParams>{
-						key: keys.key,
-						iv: keys.iv
-					});
 			});
+
+			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.');
 		}
 	}