Add akv token request logic (#9556)

Add support for running queries that require a decryption key from Azure Key Vault when using Always Encrypted.

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

@@ -159,7 +159,8 @@ export class AzureAccountProvider implements azdata.AccountProvider {
 		const resourceIdMap = new Map<azdata.AzureResource, string>([
 			[azdata.AzureResource.ResourceManagement, self._metadata.settings.armResource.id],
 			[azdata.AzureResource.Sql, self._metadata.settings.sqlResource.id],
-			[azdata.AzureResource.OssRdbms, self._metadata.settings.ossRdbmsResource.id]
+			[azdata.AzureResource.OssRdbms, self._metadata.settings.ossRdbmsResource.id],
+			[azdata.AzureResource.AzureKeyVault, self._metadata.settings.azureKeyVaultResource.id]
 		]);
 
 		let accessTokenPromises: Thenable<void>[] = [];

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

@@ -96,7 +96,8 @@ export class AzureAccountProvider implements azdata.AccountProvider {
 		const resourceIdMap = new Map<azdata.AzureResource, string>([
 			[azdata.AzureResource.ResourceManagement, this.metadata.settings.armResource.id],
 			[azdata.AzureResource.Sql, this.metadata.settings.sqlResource.id],
-			[azdata.AzureResource.OssRdbms, this.metadata.settings.ossRdbmsResource.id]
+			[azdata.AzureResource.OssRdbms, this.metadata.settings.ossRdbmsResource.id],
+			[azdata.AzureResource.AzureKeyVault, this.metadata.settings.azureKeyVaultResource.id]
 		]);
 		const tenantRefreshPromises: Promise<{ tenantId: any, securityToken: AzureAccountSecurityToken }>[] = [];
 		const tokenCollection: AzureAccountSecurityTokenCollection = {};

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

@@ -79,6 +79,11 @@ interface Settings {
 	 */
 	ossRdbmsResource?: Resource;
 
+	/**
+	 * Information that describes the Azure Key Vault resource
+	 */
+	azureKeyVaultResource?: Resource;
+
 	/**
 	 * A list of tenant IDs to authenticate against. If defined, then these IDs will be used
 	 * instead of querying the tenants endpoint of the armResource

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

@@ -33,6 +33,10 @@ const publicAzureSettings: ProviderSettings = {
 				id: 'https://ossrdbms-aad.database.windows.net',
 				endpoint: 'https://ossrdbms-aad.database.windows.net'
 			},
+			azureKeyVaultResource: {
+				id: 'https://vault.azure.net',
+				endpoint: 'https://vault.azure.net'
+			},
 			redirectUri: 'http://localhost/redirect'
 		}
 	}
@@ -56,6 +60,10 @@ const usGovAzureSettings: ProviderSettings = {
 				id: 'https://management.core.usgovcloudapi.net/',
 				endpoint: 'https://management.usgovcloudapi.net'
 			},
+			azureKeyVaultResource: {
+				id: 'https://vault.usgovcloudapi.net',
+				endpoint: 'https://vault.usgovcloudapi.net'
+			},
 			redirectUri: 'http://localhost/redirect'
 		}
 	}
@@ -79,6 +87,10 @@ const germanyAzureSettings: ProviderSettings = {
 				id: 'https://management.core.cloudapi.de/',
 				endpoint: 'https://management.microsoftazure.de'
 			},
+			azureKeyVaultResource: {
+				id: 'https://vault.microsoftazure.de',
+				endpoint: 'https://vault.microsoftazure.de'
+			},
 			redirectUri: 'http://localhost/redirect'
 		}
 	}
@@ -101,6 +113,10 @@ const chinaAzureSettings: ProviderSettings = {
 				id: 'https://management.core.chinacloudapi.cn/',
 				endpoint: 'https://managemement.chinacloudapi.net'
 			},
+			azureKeyVaultResource: {
+				id: 'https://vault.azure.cn',
+				endpoint: 'https://vault.azure.cn'
+			},
 			redirectUri: 'http://localhost/redirect'
 		}
 	}

extensions/mssql/src/contracts.ts

@@ -31,6 +31,24 @@ export class TelemetryParams {
 
 // ------------------------------- </ Telemetry Sent Event > ----------------------------------
 
+// ------------------------------- < Security Token Request > ------------------------------------------
+export interface RequestSecurityTokenParams {
+	authority: string;
+	provider: string;
+	resource: string;
+	scope: string;
+}
+
+export interface RequestSecurityTokenResponse {
+	accountKey: string;
+	token: string;
+}
+
+export namespace SecurityTokenRequest {
+	export const type = new RequestType<RequestSecurityTokenParams, RequestSecurityTokenResponse, void, void>('account/securityTokenRequest');
+}
+// ------------------------------- </ Security Token Request > ------------------------------------------
+
 // ------------------------------- < Agent Management > ------------------------------------
 // Job management parameters
 export interface AgentJobsParams {

extensions/mssql/src/features.ts

@@ -2,16 +2,18 @@
  *  Copyright (c) Microsoft Corporation. All rights reserved.
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
-
+import * as nls from 'vscode-nls';
 import { SqlOpsDataClient, SqlOpsFeature } from 'dataprotocol-client';
 import { ClientCapabilities, StaticFeature, RPCMessageType, ServerCapabilities } from 'vscode-languageclient';
-import { Disposable } from 'vscode';
+import { Disposable, window } 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';
 
+const localize = nls.loadMessageBundle();
+
 export class TelemetryFeature implements StaticFeature {
 
 	constructor(private _client: SqlOpsDataClient) { }
@@ -27,6 +29,50 @@ export class TelemetryFeature implements StaticFeature {
 	}
 }
 
+export class AccountFeature implements StaticFeature {
+
+	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;
+		});
+	}
+}
+
 export class AgentServicesFeature extends SqlOpsFeature<undefined> {
 	private static readonly messagesTypes: RPCMessageType[] = [
 		contracts.AgentJobsRequest.type,

extensions/mssql/src/sqlToolsServer.ts

@@ -11,7 +11,7 @@ import * as path from 'path';
 import { getCommonLaunchArgsAndCleanupOldLogFiles } from './utils';
 import { Telemetry, LanguageClientErrorHandler } from './telemetry';
 import { SqlOpsDataClient, ClientOptions } from 'dataprotocol-client';
-import { TelemetryFeature, AgentServicesFeature, SerializationFeature } from './features';
+import { TelemetryFeature, AgentServicesFeature, SerializationFeature, AccountFeature } from './features';
 import { CredentialStore } from './credentialstore/credentialstore';
 import { AzureResourceProvider } from './resourceProvider/resourceProvider';
 import { SchemaCompareService } from './schemaCompare/schemaCompareService';
@@ -150,6 +150,7 @@ function getClientOptions(context: AppContext): ClientOptions {
 			// we only want to add new features
 			...SqlOpsDataClient.defaultFeatures,
 			TelemetryFeature,
+			AccountFeature,
 			AgentServicesFeature,
 			SerializationFeature,
 			SchemaCompareService.asFeature(context),

src/sql/azdata.proposed.d.ts

@@ -129,7 +129,8 @@ declare module 'azdata' {
 	 * Add OssRdbms for sqlops AzureResource.
 	 */
 	export enum AzureResource {
-		OssRdbms = 2
+		OssRdbms = 2,
+		AzureKeyVault = 3
 	}
 
 	export interface ModelBuilder {

src/sql/platform/accounts/common/interfaces.ts

@@ -45,7 +45,8 @@ export interface IAccountManagementService {
 export enum AzureResource {
 	ResourceManagement = 0,
 	Sql = 1,
-	OssRdbms = 2
+	OssRdbms = 2,
+	AzureKeyVault = 3
 }
 
 export interface IAccountStore {

src/sql/workbench/api/common/sqlExtHostTypes.ts

@@ -394,7 +394,8 @@ export class TreeComponentItem extends vsExtTypes.TreeItem {
 export enum AzureResource {
 	ResourceManagement = 0,
 	Sql = 1,
-	OssRdbms = 2
+	OssRdbms = 2,
+	AzureKeyVault = 3
 }
 
 export class TreeItem extends vsExtTypes.TreeItem {