Notify STS when encryption keys are updated in azurecore (#22384)

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

@@ -12,7 +12,7 @@ import { promises as fsPromises } from 'fs';
 import { SimpleTokenCache } from './utils/simpleTokenCache';
 import providerSettings from './providerSettings';
 import { AzureAccountProvider as AzureAccountProvider } from './azureAccountProvider';
-import { AzureAccountProviderMetadata } from 'azurecore';
+import { AzureAccountProviderMetadata, CacheEncryptionKeys } from 'azurecore';
 import { ProviderSettings } from './interfaces';
 import { MsalCachePluginProvider } from './utils/msalCachePlugin';
 import * as loc from '../localizedConstants';
@@ -41,10 +41,12 @@ export class AzureAccountProviderService implements vscode.Disposable {
 	private _event: events.EventEmitter = new events.EventEmitter();
 	private readonly _uriEventHandler: UriEventHandler = new UriEventHandler();
 	public clientApplication!: PublicClientApplication;
+	private _onEncryptionKeysUpdated: vscode.EventEmitter<CacheEncryptionKeys>;
 
 	constructor(private _context: vscode.ExtensionContext,
 		private _userStoragePath: string,
 		private _authLibrary: string) {
+		this._onEncryptionKeysUpdated = new vscode.EventEmitter<CacheEncryptionKeys>();
 		this._disposables.push(vscode.window.registerUriHandler(this._uriEventHandler));
 	}
 
@@ -75,6 +77,10 @@ export class AzureAccountProviderService implements vscode.Disposable {
 			});
 	}
 
+	public getEncryptionKeysEmitter(): vscode.EventEmitter<CacheEncryptionKeys> {
+		return this._onEncryptionKeysUpdated;
+	}
+
 	public dispose() {
 		while (this._disposables.length) {
 			const item = this._disposables.pop();
@@ -155,10 +161,12 @@ export class AzureAccountProviderService implements vscode.Disposable {
 
 			// ADAL Token Cache
 			let simpleTokenCache = new SimpleTokenCache(tokenCacheKey, this._userStoragePath, noSystemKeychain, this._credentialProvider);
-			await simpleTokenCache.init();
+			if (this._authLibrary === Constants.AuthLibrary.ADAL) {
+				await simpleTokenCache.init();
+			}
 
 			// MSAL Cache Plugin
-			this._cachePluginProvider = new MsalCachePluginProvider(tokenCacheKeyMsal, this._userStoragePath, this._credentialProvider);
+			this._cachePluginProvider = new MsalCachePluginProvider(tokenCacheKeyMsal, this._userStoragePath, this._credentialProvider, this._onEncryptionKeysUpdated);
 
 			const msalConfiguration: Configuration = {
 				auth: {

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

@@ -9,12 +9,14 @@ import * as vscode from 'vscode';
 import { AuthLibrary } from '../../constants';
 import * as LocalizedConstants from '../../localizedConstants';
 import { Logger } from '../../utils/Logger';
+import { CacheEncryptionKeys } from 'azurecore';
 
 export class FileEncryptionHelper {
 	constructor(
 		private readonly _authLibrary: AuthLibrary,
 		private readonly _credentialService: azdata.CredentialProvider,
-		protected readonly _fileName: string
+		protected readonly _fileName: string,
+		private readonly _onEncryptionKeysUpdated?: vscode.EventEmitter<CacheEncryptionKeys>
 	) {
 		this._algorithm = this._authLibrary === AuthLibrary.MSAL ? 'aes-256-cbc' : 'aes-256-gcm';
 		this._bufferEncoding = this._authLibrary === AuthLibrary.MSAL ? 'utf16le' : 'hex';
@@ -48,6 +50,14 @@ export class FileEncryptionHelper {
 			this._ivBuffer = Buffer.from(iv, this._bufferEncoding);
 			this._keyBuffer = Buffer.from(key, this._bufferEncoding);
 		}
+
+		// 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)
+			});
+		}
 	}
 
 	fileSaver = async (content: string): Promise<string> => {

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

@@ -9,18 +9,21 @@ import { promises as fsPromises } from 'fs';
 import * as lockFile from 'lockfile';
 import * as path from 'path';
 import * as azdata from 'azdata';
+import * as vscode from 'vscode';
 import { AccountsClearTokenCacheCommand, AuthLibrary } from '../../constants';
 import { Logger } from '../../utils/Logger';
 import { FileEncryptionHelper } from './fileEncryptionHelper';
+import { CacheEncryptionKeys } from 'azurecore';
 
 export class MsalCachePluginProvider {
 	constructor(
 		private readonly _serviceName: string,
 		private readonly _msalFilePath: string,
-		private readonly _credentialService: azdata.CredentialProvider
+		private readonly _credentialService: azdata.CredentialProvider,
+		private readonly _onEncryptionKeysUpdated: vscode.EventEmitter<CacheEncryptionKeys>
 	) {
 		this._msalFilePath = path.join(this._msalFilePath, this._serviceName);
-		this._fileEncryptionHelper = new FileEncryptionHelper(AuthLibrary.MSAL, this._credentialService, this._serviceName);
+		this._fileEncryptionHelper = new FileEncryptionHelper(AuthLibrary.MSAL, this._credentialService, this._serviceName, this._onEncryptionKeysUpdated);
 	}
 
 	private _lockTaken: boolean = false;

extensions/azurecore/src/azurecore.d.ts

@@ -5,7 +5,7 @@
 
 declare module 'azurecore' {
 	import * as azdata from 'azdata';
-	import { TreeDataProvider } from 'vscode';
+	import * as vscode from 'vscode';
 	import { BlobItem } from '@azure/storage-blob';
 
 	/**
@@ -314,8 +314,13 @@ declare module 'azurecore' {
 		getRegionDisplayName(region?: string): string;
 		getProviderMetadataForAccount(account: AzureAccount): AzureAccountProviderMetadata;
 		provideResources(): azureResource.IAzureResourceProvider[];
-
 		runGraphQuery<T extends azureResource.AzureGraphResource>(account: AzureAccount, subscriptions: azureResource.AzureResourceSubscription[], ignoreErrors: boolean, query: string): Promise<ResourceQueryResult<T>>;
+		/**
+		 * Event emitted when MSAL cache encryption keys are updated in credential store.
+		 * Returns encryption keys used for encryption/decryption of MSAL cache that can be used
+		 * by connection providers to read/write to the same access token cache for stable connectivity.
+		 */
+		onEncryptionKeysUpdated: vscode.Event<CacheEncryptionKeys>;
 	}
 
 	export type GetSubscriptionsResult = { subscriptions: azureResource.AzureResourceSubscription[], errors: Error[] };
@@ -333,6 +338,7 @@ declare module 'azurecore' {
 	export type AzureRestResponse = { response: any, errors: Error[] };
 	export type GetBlobsResult = { blobs: azureResource.Blob[], errors: Error[] };
 	export type GetStorageAccountAccessKeyResult = { keyName1: string, keyName2: string, errors: Error[] };
+	export type CacheEncryptionKeys = { key: string; iv: string; }
 
 	export namespace azureResource {
 

extensions/azurecore/src/extension.ts

@@ -99,38 +99,43 @@ export async function activate(context: vscode.ExtensionContext): Promise<azurec
 
 	updatePiiLoggingLevel();
 
+	let eventEmitter: vscode.EventEmitter<azurecore.CacheEncryptionKeys>;
 	// Create the provider service and activate
-	initAzureAccountProvider(extensionContext, storagePath, authLibrary!).catch((err) => Logger.error(err));
+	let providerService = await initAzureAccountProvider(extensionContext, storagePath, authLibrary!).catch((err) => Logger.error(err));
-
+	if (providerService) {
-	registerAzureServices(appContext);
+		eventEmitter = providerService.getEncryptionKeysEmitter();
-	const azureResourceTree = new AzureResourceTreeProvider(appContext, authLibrary);
-	const connectionDialogTree = new ConnectionDialogTreeProvider(appContext, authLibrary);
-	pushDisposable(vscode.window.registerTreeDataProvider('azureResourceExplorer', azureResourceTree));
-	pushDisposable(vscode.window.registerTreeDataProvider('connectionDialog/azureResourceExplorer', connectionDialogTree));
-	pushDisposable(vscode.workspace.onDidChangeConfiguration(e => onDidChangeConfiguration(e)));
-	registerAzureResourceCommands(appContext, azureResourceTree, connectionDialogTree, authLibrary);
-	azdata.dataprotocol.registerDataGridProvider(new AzureDataGridProvider(appContext, authLibrary));
-	vscode.commands.registerCommand('azure.dataGrid.openInAzurePortal', async (item: azdata.DataGridItem) => {
-		const portalEndpoint = item.portalEndpoint;
-		const subscriptionId = item.subscriptionId;
-		const resourceGroup = item.resourceGroup;
-		const type = item.type;
-		const name = item.name;
-		if (portalEndpoint && subscriptionId && resourceGroup && type && name) {
-			await vscode.env.openExternal(vscode.Uri.parse(`${portalEndpoint}/#resource/subscriptions/${subscriptionId}/resourceGroups/${resourceGroup}/providers/${type}/${name}`));
-		} else {
-			Logger.error(`Missing required values - subscriptionId : ${subscriptionId} resourceGroup : ${resourceGroup} type: ${type} name: ${name}`);
-			void vscode.window.showErrorMessage(loc.unableToOpenAzureLink);
-		}
-	});
 
+		registerAzureServices(appContext);
+		const azureResourceTree = new AzureResourceTreeProvider(appContext, authLibrary);
+		const connectionDialogTree = new ConnectionDialogTreeProvider(appContext, authLibrary);
+		pushDisposable(vscode.window.registerTreeDataProvider('azureResourceExplorer', azureResourceTree));
+		pushDisposable(vscode.window.registerTreeDataProvider('connectionDialog/azureResourceExplorer', connectionDialogTree));
+		pushDisposable(vscode.workspace.onDidChangeConfiguration(e => onDidChangeConfiguration(e)));
+		registerAzureResourceCommands(appContext, azureResourceTree, connectionDialogTree, authLibrary);
+		azdata.dataprotocol.registerDataGridProvider(new AzureDataGridProvider(appContext, authLibrary));
+		vscode.commands.registerCommand('azure.dataGrid.openInAzurePortal', async (item: azdata.DataGridItem) => {
+			const portalEndpoint = item.portalEndpoint;
+			const subscriptionId = item.subscriptionId;
+			const resourceGroup = item.resourceGroup;
+			const type = item.type;
+			const name = item.name;
+			if (portalEndpoint && subscriptionId && resourceGroup && type && name) {
+				await vscode.env.openExternal(vscode.Uri.parse(`${portalEndpoint}/#resource/subscriptions/${subscriptionId}/resourceGroups/${resourceGroup}/providers/${type}/${name}`));
+			} else {
+				Logger.error(`Missing required values - subscriptionId : ${subscriptionId} resourceGroup : ${resourceGroup} type: ${type} name: ${name}`);
+				void vscode.window.showErrorMessage(loc.unableToOpenAzureLink);
+			}
+		});
+	}
 	return {
 		getSubscriptions(account?: azurecore.AzureAccount, ignoreErrors?: boolean, selectedOnly: boolean = false): Promise<azurecore.GetSubscriptionsResult> {
 			return selectedOnly
 				? azureResourceUtils.getSelectedSubscriptions(appContext, account, ignoreErrors)
 				: azureResourceUtils.getSubscriptions(appContext, account, ignoreErrors);
 		},
-		getResourceGroups(account?: azurecore.AzureAccount, subscription?: azurecore.azureResource.AzureResourceSubscription, ignoreErrors?: boolean): Promise<azurecore.GetResourceGroupsResult> { return azureResourceUtils.getResourceGroups(appContext, account, subscription, ignoreErrors); },
+		getResourceGroups(account?: azurecore.AzureAccount, subscription?: azurecore.azureResource.AzureResourceSubscription, ignoreErrors?: boolean): Promise<azurecore.GetResourceGroupsResult> {
+			return azureResourceUtils.getResourceGroups(appContext, account, subscription, ignoreErrors);
+		},
 		getLocations(account?: azurecore.AzureAccount,
 			subscription?: azurecore.azureResource.AzureResourceSubscription,
 			ignoreErrors?: boolean): Promise<azurecore.GetLocationsResult> {
@@ -235,7 +240,8 @@ export async function activate(context: vscode.ExtensionContext): Promise<azurec
 			ignoreErrors: boolean,
 			query: string): Promise<azurecore.ResourceQueryResult<T>> {
 			return azureResourceUtils.runResourceQuery(account, subscriptions, ignoreErrors, query);
-		}
+		},
+		onEncryptionKeysUpdated: eventEmitter!.event
 	};
 }
 
@@ -267,13 +273,15 @@ async function findOrMakeStoragePath() {
 	return storagePath;
 }
 
-async function initAzureAccountProvider(extensionContext: vscode.ExtensionContext, storagePath: string, authLibrary: string): Promise<void> {
+async function initAzureAccountProvider(extensionContext: vscode.ExtensionContext, storagePath: string, authLibrary: string): Promise<AzureAccountProviderService | undefined> {
 	try {
 		const accountProviderService = new AzureAccountProviderService(extensionContext, storagePath, authLibrary);
 		extensionContext.subscriptions.push(accountProviderService);
 		await accountProviderService.activate();
+		return accountProviderService;
 	} catch (err) {
 		Logger.error('Unexpected error starting account provider: ' + err.message);
+		return undefined;
 	}
 }
 

extensions/machine-learning/src/test/stubs.ts

@@ -61,5 +61,6 @@ export class AzurecoreApiStub implements azurecore.IExtension {
 	provideResources(): azurecore.azureResource.IAzureResourceProvider[] {
 		throw new Error('Method not implemented.');
 	}
+	onEncryptionKeysUpdated: any
 
 }

extensions/mssql/src/contracts.ts

@@ -1604,3 +1604,25 @@ export namespace DropObjectRequest {
 	export const type = new RequestType<DropObjectRequestParams, void, void, void>('objectManagement/drop');
 }
 // ------------------------------- < Object Management > ------------------------------------
+
+// ------------------------------- < Encryption IV/KEY updation Event > ------------------------------------
+/**
+ * Parameters for the MSAL cache encryption key notification
+ */
+export class DidChangeEncryptionIVKeyParams {
+	/**
+	 * Buffer encoded IV string for MSAL cache encryption
+	 */
+	public iv: string;
+	/**
+	 * Buffer encoded Key string for MSAL cache encryption
+	 */
+	public key: string;
+}
+
+/**
+ * Notification sent when the encryption keys are changed.
+ */
+export namespace EncryptionKeysChangedNotification {
+	export const type = new NotificationType<DidChangeEncryptionIVKeyParams, void>('connection/encryptionKeysChanged');
+}

extensions/mssql/src/sqlToolsServer.ts

@@ -9,6 +9,7 @@ import * as Constants from './constants';
 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 { TelemetryReporter, LanguageClientErrorHandler } from './telemetry';
 import { SqlOpsDataClient, ClientOptions } from 'dataprotocol-client';
@@ -19,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 } from './contracts';
+import { CompletionExtensionParams, CompletionExtLoadRequest, DidChangeEncryptionIVKeyParams, EncryptionKeysChangedNotification } from './contracts';
 import { promises as fs } from 'fs';
 import * as nls from 'vscode-nls';
 import { LanguageExtensionService } from './languageExtension/languageExtensionService';
@@ -82,6 +83,7 @@ export class SqlToolsServer {
 			statusView.text = localize('startingServiceStatusMsg', "Starting {0}", Constants.serviceName);
 			this.client.start();
 			await Promise.all([this.activateFeatures(context), clientReadyPromise]);
+			await this.handleEncryptionKeyEventNotification(this.client);
 			return this.client;
 		} catch (e) {
 			TelemetryReporter.sendTelemetryEvent('ServiceInitializingFailed');
@@ -90,6 +92,35 @@ export class SqlToolsServer {
 		}
 	}
 
+	/**
+	 * This is a hop notification handler to send Encryption Key and Iv information from Azure Core extension to backend
+	 * SqlToolsService. This notification is needed for Azure authentication flows to be able to read/write into
+	 * shared MSAL cache.
+	 * @param client SqlOpsDataClient instance
+	 */
+	private async handleEncryptionKeyEventNotification(client: SqlOpsDataClient) {
+		if (getAzureAuthenticationLibraryConfig() === 'MSAL' && getEnableSqlAuthenticationProviderConfig()) {
+			let onDidEncryptionKeysChanged = (await this.getAzureCoreAPI()).onEncryptionKeysUpdated;
+			// Register event listener from Azure Core extension
+			onDidEncryptionKeysChanged((keys: azurecore.CacheEncryptionKeys) => {
+				// Send client notification for updated encryption keys
+				client.sendNotification(EncryptionKeysChangedNotification.type,
+					<DidChangeEncryptionIVKeyParams>{
+						key: keys.key,
+						iv: keys.iv
+					});
+			});
+		}
+	}
+
+	private async getAzureCoreAPI(): Promise<azurecore.IExtension> {
+		const api = (await vscode.extensions.getExtension(azurecore.extension.name)?.activate()) as azurecore.IExtension;
+		if (!api) {
+			throw new Error('Azure core extension could not be activated.');
+		}
+		return api;
+	}
+
 	private async download(context: AppContext): Promise<string> {
 		const configDir = context.extensionContext.extensionPath;
 		const rawConfig = await fs.readFile(path.join(configDir, 'config.json'));