Add Simple Account Picker for use with Always Encrypted (#9707)

Adds the ability for the user to select from two or more linked azure accounts, using an integrated UI dialog, when executing a query that requires a Always Encrypted column master key located in Azure Key Vault.

extensions/mssql/coverConfig.json

@@ -0,0 +1,19 @@
+{
+	"enabled": true,
+	"relativeSourcePath": "..",
+	"relativeCoverageDir": "../../coverage",
+	"ignorePatterns": [
+		"**/node_modules/**",
+		"**/test/**"
+	],
+	"includePid": false,
+	"reports": [
+		"cobertura",
+		"lcov"
+	],
+	"verbose": false,
+	"remapOptions": {
+		"basePath": "..",
+		"useAbsolutePaths": true
+	}
+}

extensions/mssql/package.json

@@ -1052,10 +1052,16 @@
   },
   "devDependencies": {
     "@types/bytes": "^3.0.0",
+    "@types/chai": "^4.2.11",
     "@types/kerberos": "^1.1.0",
+    "@types/mocha": "^7.0.2",
     "@types/request": "^2.48.2",
     "@types/request-promise": "^4.1.44",
     "@types/stream-meter": "^0.0.22",
-    "@types/through2": "^2.0.34"
+    "@types/through2": "^2.0.34",
+    "chai": "^4.2.0",
+    "mocha": "^7.1.1",
+    "typemoq": "^2.1.0",
+    "vscodetestcover": "github:corivera/vscodetestcover#1.0.5"
   }
 }

extensions/mssql/src/features.ts

@@ -5,12 +5,13 @@
 import * as nls from 'vscode-nls';
 import { SqlOpsDataClient, SqlOpsFeature } from 'dataprotocol-client';
 import { ClientCapabilities, StaticFeature, RPCMessageType, ServerCapabilities } from 'vscode-languageclient';
-import { Disposable, window } from 'vscode';
+import { Disposable, window, QuickPickItem, QuickPickOptions } from 'vscode';
 import { Telemetry } from './telemetry';
 import * as contracts from './contracts';
 import * as azdata from 'azdata';
 import * as Utils from './utils';
 import * as UUID from 'vscode-languageclient/lib/utils/uuid';
+import { DataItemCache } from './util/dataCache';
 
 const localize = nls.loadMessageBundle();
 
@@ -31,46 +32,78 @@ export class TelemetryFeature implements StaticFeature {
 
 export class AccountFeature implements StaticFeature {
 
+	tokenCache: DataItemCache<contracts.RequestSecurityTokenResponse | undefined>;
+
 	constructor(private _client: SqlOpsDataClient) { }
 
 	fillClientCapabilities(capabilities: ClientCapabilities): void { }
 
 	initialize(): void {
-		this._client.onRequest(contracts.SecurityTokenRequest.type, async (e): Promise<contracts.RequestSecurityTokenResponse | undefined> => {
-			const accountList = await azdata.accounts.getAllAccounts();
-
-			if (accountList.length < 1) {
-				// TODO: Prompt user to add account
-				window.showErrorMessage(localize('mssql.missingLinkedAzureAccount', "Azure Data Studio needs to contact Azure Key Vault to access a column master key for Always Encrypted, but no linked Azure account is available. Please add a linked Azure account and retry the query."));
-				return undefined;
-			} else if (accountList.length > 1) {
-				// TODO: Prompt user to select an account
-				window.showErrorMessage(localize('mssql.multipleLinkedAzureAccount', "Azure Data Studio needs to contact Azure Key Vault to access a column master key for Always Encrypted, which is not supported if multiple linked Azure accounts are present. Make sure only one linked Azure account exists and retry the query."));
-				return undefined;
-			}
-
-			let account = accountList[0];
-			const securityToken: { [key: string]: any } = await azdata.accounts.getSecurityToken(account, azdata.AzureResource.AzureKeyVault);
-			const tenant = account.properties.tenants.find((t: { [key: string]: string }) => e.authority.includes(t.id));
-			const unauthorizedMessage = localize('mssql.insufficientlyPrivelagedAzureAccount', "The configured Azure account for {0} does not have sufficient permissions for Azure Key Vault to access a column master key for Always Encrypted.", account.key.accountId);
-			if (!tenant) {
-				window.showErrorMessage(unauthorizedMessage);
-				return undefined;
-			}
-			let tokenBundle = securityToken[tenant.id];
-			if (!tokenBundle) {
-				window.showErrorMessage(unauthorizedMessage);
-				return undefined;
-			}
-
-			let params: contracts.RequestSecurityTokenResponse = {
-				accountKey: JSON.stringify(account.key),
-				token: securityToken[tenant.id].token
-			};
-
-			return params;
+		let timeToLiveInSeconds = 10;
+		this.tokenCache = new DataItemCache(this.getToken, timeToLiveInSeconds);
+		this._client.onRequest(contracts.SecurityTokenRequest.type, async (request): Promise<contracts.RequestSecurityTokenResponse | undefined> => {
+			return this.tokenCache.getData(request);
 		});
 	}
+
+	protected async getToken(request: contracts.RequestSecurityTokenParams): Promise<contracts.RequestSecurityTokenResponse | undefined> {
+		const accountList = await azdata.accounts.getAllAccounts();
+		let account: azdata.Account;
+
+		if (accountList.length < 1) {
+			// TODO: Prompt user to add account
+			window.showErrorMessage(localize('mssql.missingLinkedAzureAccount', "Azure Data Studio needs to contact Azure Key Vault to access a column master key for Always Encrypted, but no linked Azure account is available. Please add a linked Azure account and retry the query."));
+			return undefined;
+		} else if (accountList.length > 1) {
+			let options: QuickPickOptions = {
+				ignoreFocusOut: true,
+				placeHolder: localize('mssql.chooseLinkedAzureAccount', "Please select a linked Azure account:")
+			};
+			let items = accountList.map(a => new AccountFeature.AccountQuickPickItem(a));
+			let selectedItem = await window.showQuickPick(items, options);
+			if (!selectedItem) { // The user canceled the selection.
+				window.showErrorMessage(localize('mssql.canceledLinkedAzureAccountSelection', "Azure Data Studio needs to contact Azure Key Vault to access a column master key for Always Encrypted, but no linked Azure account was selected. Please retry the query and select a linked Azure account when prompted."));
+				return undefined;
+			}
+			account = selectedItem.account;
+		} else {
+			account = accountList[0];
+		}
+
+		const securityToken: { [key: string]: any } = await azdata.accounts.getSecurityToken(account, azdata.AzureResource.AzureKeyVault);
+		const tenant = account.properties.tenants.find((t: { [key: string]: string }) => request.authority.includes(t.id));
+		const unauthorizedMessage = localize('mssql.insufficientlyPrivelagedAzureAccount', "The configured Azure account for {0} does not have sufficient permissions for Azure Key Vault to access a column master key for Always Encrypted.", account.key.accountId);
+		if (!tenant) {
+			window.showErrorMessage(unauthorizedMessage);
+			return undefined;
+		}
+		let tokenBundle = securityToken[tenant.id];
+		if (!tokenBundle) {
+			window.showErrorMessage(unauthorizedMessage);
+			return undefined;
+		}
+
+		let params: contracts.RequestSecurityTokenResponse = {
+			accountKey: JSON.stringify(account.key),
+			token: securityToken[tenant.id].token
+		};
+
+		return params;
+	}
+
+	static AccountQuickPickItem = class implements QuickPickItem {
+		account: azdata.Account;
+		label: string;
+		description?: string;
+		detail?: string;
+		picked?: boolean;
+		alwaysShow?: boolean;
+
+		constructor(account: azdata.Account) {
+			this.account = account;
+			this.label = account.key.accountId;
+		}
+	};
 }
 
 export class AgentServicesFeature extends SqlOpsFeature<undefined> {

extensions/mssql/src/test/index.ts

@@ -0,0 +1,48 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as path from 'path';
+const testRunner = require('vscodetestcover');
+
+const suite = 'mssql Extension Tests';
+
+const mochaOptions: any = {
+	ui: 'bdd',
+	useColors: true,
+	timeout: 10000
+};
+
+// set relevant mocha options from the environment
+if (process.env.ADS_TEST_GREP) {
+	mochaOptions.grep = process.env.ADS_TEST_GREP;
+	console.log(`setting options.grep to: ${mochaOptions.grep}`);
+}
+if (process.env.ADS_TEST_INVERT_GREP) {
+	mochaOptions.invert = parseInt(process.env.ADS_TEST_INVERT_GREP);
+	console.log(`setting options.invert to: ${mochaOptions.invert}`);
+}
+if (process.env.ADS_TEST_TIMEOUT) {
+	mochaOptions.timeout = parseInt(process.env.ADS_TEST_TIMEOUT);
+	console.log(`setting options.timeout to: ${mochaOptions.timeout}`);
+}
+if (process.env.ADS_TEST_RETRIES) {
+	mochaOptions.retries = parseInt(process.env.ADS_TEST_RETRIES);
+	console.log(`setting options.retries to: ${mochaOptions.retries}`);
+}
+
+if (process.env.BUILD_ARTIFACTSTAGINGDIRECTORY) {
+	mochaOptions.reporter = 'mocha-multi-reporters';
+	mochaOptions.reporterOptions = {
+		reporterEnabled: 'spec, mocha-junit-reporter',
+		mochaJunitReporterReporterOptions: {
+			testsuitesTitle: `${suite} ${process.platform}`,
+			mochaFile: path.join(process.env.BUILD_ARTIFACTSTAGINGDIRECTORY, `test-results/${process.platform}-${suite.toLowerCase().replace(/[^\w]/g, '-')}-results.xml`)
+		}
+	};
+}
+
+testRunner.configure(mochaOptions, { coverConfig: '../../coverConfig.json' });
+
+export = testRunner;

extensions/mssql/src/test/util/dataCache.test.ts

@@ -0,0 +1,70 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { DataItemCache } from '../../util/dataCache';
+import 'mocha';
+import { should } from 'chai'; should();
+import * as TypeMoq from "typemoq";
+
+describe('DataItemCache', function (): void {
+
+	const testCacheItem = 'Test Cache Item';
+	const fetchFunction = () => Promise.resolve(testCacheItem);
+	let fetchFunctionMock: TypeMoq.IMock<() => Promise<string>>;
+	let dataItemCache: DataItemCache<String>;
+
+	beforeEach(function (): void {
+		fetchFunctionMock = TypeMoq.Mock.ofInstance(fetchFunction);
+		fetchFunctionMock.setup(fx => fx()).returns(() => Promise.resolve(testCacheItem));
+		dataItemCache = new DataItemCache<string>(fetchFunctionMock.object, 1);
+	});
+
+	it('Should be initialized empty', function (): void {
+		dataItemCache.should.have.property('cachedItem').and.be.undefined;
+	});
+
+	it('Should be initialized as expired', function (): void {
+		dataItemCache.isCacheExpired().should.be.true;
+	});
+
+	it('Should not be expired immediately after first data fetch', async function (): Promise<void> {
+		await dataItemCache.getData();
+
+		dataItemCache.isCacheExpired().should.be.false;
+	});
+
+	it('Should return expected cached item from getValue()', async function (): Promise<void> {
+		let actualValue = await dataItemCache.getData();
+
+		actualValue.should.equal(testCacheItem);
+	});
+
+	it('Should be expired after data is fetched and TTL passes', async function (): Promise<void> {
+		await dataItemCache.getData();
+		await sleep(1.1);
+
+		dataItemCache.isCacheExpired().should.be.true;
+	});
+
+	it('Should call fetch function once for consecutive getValue() calls prior to expiration', async function (): Promise<void> {
+		await dataItemCache.getData();
+		await dataItemCache.getData();
+		await dataItemCache.getData();
+
+		fetchFunctionMock.verify(fx => fx() ,TypeMoq.Times.once());
+	});
+
+	it('Should call fetch function twice for consecutive getValue() calls if TTL expires in between', async function (): Promise<void> {
+		await dataItemCache.getData();
+		await sleep(1.1);
+		await dataItemCache.getData();
+
+		fetchFunctionMock.verify(fx => fx(), TypeMoq.Times.exactly(2));
+	});
+});
+
+const sleep = (seconds: number) => {
+	return new Promise(resolve => setTimeout(resolve, 1000 * seconds));
+}

extensions/mssql/src/util/dataCache.ts

@@ -0,0 +1,37 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+export class DataItemCache<T> {
+
+	millisecondsToLive: number;
+	getValueFunction: (...args: any[]) => Promise<T>;
+	cachedItem: T;
+	fetchDate: Date;
+
+	constructor(getValueFunction: (...args: any[]) => Promise<T>, secondsToLive: number) {
+		this.millisecondsToLive = secondsToLive * 1000;
+		this.getValueFunction = getValueFunction;
+		this.cachedItem = undefined;
+		this.fetchDate = new Date(0);
+	}
+
+	public isCacheExpired(): boolean {
+		return (this.fetchDate.getTime() + this.millisecondsToLive) < new Date().getTime();
+	}
+
+	public async getData(...args: any[]): Promise<T> {
+		if (!this.cachedItem || this.isCacheExpired()) {
+			let data = await this.getValueFunction(...args);
+			this.cachedItem = data;
+			this.fetchDate = new Date();
+			return data;
+		}
+		return this.cachedItem;
+	}
+
+	public resetCache(): void {
+		this.fetchDate = new Date(0);
+	}
+}

scripts/test-extensions-unit.bat

@@ -83,6 +83,11 @@ echo *** starting machine-learning-services tests ***
 echo *******************************
 call "%INTEGRATION_TEST_ELECTRON_PATH%" --extensionDevelopmentPath=%~dp0\..\extensions\machine-learning-services --extensionTestsPath=%~dp0\..\extensions\machine-learning-services\out\test --user-data-dir=%VSCODEUSERDATADIR% --extensions-dir=%VSCODEEXTENSIONSDIR% --remote-debugging-port=9222 --disable-telemetry --disable-crash-reporter --disable-updates --nogpu
 
+REM echo ******************************************
+REM echo *** starting mssql tests ***
+REM echo ******************************************
+REM call "%INTEGRATION_TEST_ELECTRON_PATH%" --extensionDevelopmentPath=%~dp0\..\extensions\mssql --extensionTestsPath=%~dp0\..\extensions\mssql\out\test --user-data-dir=%VSCODEUSERDATADIR% --extensions-dir=%VSCODEEXTENSIONSDIR% --remote-debugging-port=9222 --disable-telemetry --disable-crash-reporter --disable-updates --nogpu
+
 if %errorlevel% neq 0 exit /b %errorlevel%
 
 rmdir /s /q %VSCODEUSERDATADIR%

scripts/test-extensions-unit.sh

@@ -84,5 +84,10 @@ echo *** starting machine-learning-services tests ***
 echo ******************************************
 "$INTEGRATION_TEST_ELECTRON_PATH" $LINUX_NO_SANDBOX --extensionDevelopmentPath=$ROOT/extensions/machine-learning-services --extensionTestsPath=$ROOT/extensions/machine-learning-services/out/test --user-data-dir=$VSCODEUSERDATADIR --extensions-dir=$VSCODEEXTDIR --disable-telemetry --disable-crash-reporter --disable-updates --nogpu
 
+# echo ******************************************
+# echo *** starting mssql tests ***
+# echo ******************************************
+# "$INTEGRATION_TEST_ELECTRON_PATH" $LINUX_NO_SANDBOX --extensionDevelopmentPath=$ROOT/extensions/mssql --extensionTestsPath=$ROOT/extensions/mssql/out/test --user-data-dir=$VSCODEUSERDATADIR --extensions-dir=$VSCODEEXTDIR --disable-telemetry --disable-crash-reporter --disable-updates --nogpu
+
 rm -r $VSCODEUSERDATADIR
 rm -r $VSCODEEXTDIR