From c79cfd709ad45f87257f8324d4abfbbb462b030c Mon Sep 17 00:00:00 2001 From: Justin M <63619224+JustinMDotNet@users.noreply.github.com> Date: Mon, 28 Sep 2020 11:59:16 -0700 Subject: [PATCH] 3644 Kusto Token Refresh (#12576) * 3644 Added RequestSecurityTokenParams, RequestSecurityTokenResponse, and SecurityTokenRequest to Kusto/contracts.ts. Added AccountFeature to features.ts. Registered feature in kustoServer.ts * 3644 Removed TryCatch in kusto features > getToken * 3644 Added AccountId to Kusto > RequestSecurityTokenParams. Refactored kusto features getToken to use the accountId for the query window. * 3644 Removed unused AccountQuickPickItem --- extensions/kusto/src/contracts.ts | 18 ++++++++++ extensions/kusto/src/features.ts | 54 ++++++++++++++++++++++++++++- extensions/kusto/src/kustoServer.ts | 3 +- 3 files changed, 73 insertions(+), 2 deletions(-) diff --git a/extensions/kusto/src/contracts.ts b/extensions/kusto/src/contracts.ts index 491c6cfb02..b16943abfe 100644 --- a/extensions/kusto/src/contracts.ts +++ b/extensions/kusto/src/contracts.ts @@ -29,6 +29,24 @@ export class TelemetryParams { // ------------------------------- ---------------------------------- +// ------------------------------- < Security Token Request > ------------------------------------------ +export interface RequestSecurityTokenParams { + authority: string; + provider: string; + resource: string; + accountId: string; +} + +export interface RequestSecurityTokenResponse { + accountKey: string; + token: string; +} + +export namespace SecurityTokenRequest { + export const type = new RequestType('account/securityTokenRequest'); +} +// ------------------------------- ------------------------------------------ + // ------------------------------- ----------------------------- export namespace SerializeDataStartRequest { export const type = new RequestType('serialize/start'); diff --git a/extensions/kusto/src/features.ts b/extensions/kusto/src/features.ts index 3404224873..d100dcaa81 100644 --- a/extensions/kusto/src/features.ts +++ b/extensions/kusto/src/features.ts @@ -5,12 +5,13 @@ 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'; +import { localize } from './localize'; export class TelemetryFeature implements StaticFeature { @@ -27,6 +28,57 @@ 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 (request): Promise => { + return this.getToken(request); + }); + } + + protected async getToken(request: contracts.RequestSecurityTokenParams): Promise { + const accountList = await azdata.accounts.getAllAccounts(); + let account: azdata.Account | undefined; + + if (accountList.length < 1) { + // TODO: Prompt user to add account + window.showErrorMessage(localize('kusto.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 { + account = accountList.find(a => a.key.accountId === request.accountId); + } + + if (!account) { + window.showErrorMessage(localize('kusto.accountDoesNotExist', "Account does not exist.")); + return undefined; + } + + const tenant = account.properties.tenants.find((t: { [key: string]: string }) => request.authority.includes(t.id)); + const unauthorizedMessage = localize('kusto.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; + } + const securityToken = await azdata.accounts.getAccountSecurityToken(account, tenant.id, azdata.AzureResource.Sql); + + if (!securityToken?.token) { + window.showErrorMessage(unauthorizedMessage); + return undefined; + } + + let params: contracts.RequestSecurityTokenResponse = { + accountKey: JSON.stringify(account.key), + token: securityToken.token + }; + + return params; + } +} + export class SerializationFeature extends SqlOpsFeature { private static readonly messageTypes: RPCMessageType[] = [ contracts.SerializeDataStartRequest.type, diff --git a/extensions/kusto/src/kustoServer.ts b/extensions/kusto/src/kustoServer.ts index e9baf275b0..f1c2ff9c05 100644 --- a/extensions/kusto/src/kustoServer.ts +++ b/extensions/kusto/src/kustoServer.ts @@ -12,7 +12,7 @@ import { getCommonLaunchArgsAndCleanupOldLogFiles } from './utils'; import { localize } from './localize'; import { Telemetry, LanguageClientErrorHandler } from './telemetry'; import { SqlOpsDataClient, ClientOptions } from 'dataprotocol-client'; -import { TelemetryFeature, SerializationFeature } from './features'; +import { TelemetryFeature, SerializationFeature, AccountFeature } from './features'; import { AppContext } from './appContext'; import { CompletionExtensionParams, CompletionExtLoadRequest } from './contracts'; import { promises as fs } from 'fs'; @@ -135,6 +135,7 @@ function getClientOptions(context: AppContext): ClientOptions { // we only want to add new features ...SqlOpsDataClient.defaultFeatures, TelemetryFeature, + AccountFeature, SerializationFeature ], outputChannel: new CustomOutputChannel()