From 86e0c6963ff728e0f4afa7f12fe2dda39906f8fb Mon Sep 17 00:00:00 2001 From: Christopher Suh Date: Fri, 13 May 2022 07:33:14 -0700 Subject: [PATCH] Refreshes token for intellisense (#19214) * wip * wip * wip * working refresh token for intellisense * pr review comments * pr review comments * add link * pr comments * change authority -> tenantId * refactor tenant * fix build * add js doc comments, other pr changes * fix error messaging * add log * added logs * remove expiresOn * fix error messaging * pr comments * remove localized strings from logs --- extensions/mssql/src/contracts.ts | 63 ++++++++++++++++++++++ extensions/mssql/src/features.ts | 45 +++++++++++++++- extensions/mssql/src/localizedConstants.ts | 7 +++ 3 files changed, 114 insertions(+), 1 deletion(-) diff --git a/extensions/mssql/src/contracts.ts b/extensions/mssql/src/contracts.ts index 87861f1e1f..34054babe9 100644 --- a/extensions/mssql/src/contracts.ts +++ b/extensions/mssql/src/contracts.ts @@ -49,6 +49,69 @@ export namespace SecurityTokenRequest { } // ------------------------------- ------------------------------------------ +// ------------------------------- < Refresh Token Notification > --------------------------------- + +/** + * Parameters for a refresh token notification sent from STS to ADS + */ +export interface RefreshTokenParams { + /** + * The tenant ID + */ + tenantId: string; + /** + * The provider that indicates the type of linked account to query + */ + provider: string; + /** + * The identifier of the target resource of the requested token + */ + resource: string; + /** + * The account ID + */ + accountId: string; + /** + * The URI for the editor that needs a token refresh + */ + uri: string; +} + +export namespace RefreshTokenNotification { + export const type = new NotificationType('account/refreshToken'); +} + + + +// ------------------------------- ------------------------------- + +// ------------------------------- < Token Refreshed Notification > --------------------------------- + +/** + * Parameters for a new refresh token sent from ADS to STS + */ +export interface TokenRefreshedParams { + /** + * The refresh token + */ + token: string; + /** + * The token expiration, a Unix epoch + */ + expiresOn: Number; + /** + * The URI for the editor that needs a token refresh + */ + uri: string; +} + +export namespace TokenRefreshedNotification { + export const type = new NotificationType('account/tokenRefreshed'); +} + +// ------------------------------- ------------------------------- + + // ------------------------------- < Agent Management > ------------------------------------ // Job management parameters export interface AgentJobsParams { diff --git a/extensions/mssql/src/features.ts b/extensions/mssql/src/features.ts index fa4feb46ed..5d1c83d10d 100644 --- a/extensions/mssql/src/features.ts +++ b/extensions/mssql/src/features.ts @@ -13,6 +13,7 @@ import * as Utils from './utils'; import * as UUID from 'vscode-languageclient/lib/utils/uuid'; import { DataItemCache } from './util/dataCache'; import * as azurecore from 'azurecore'; +import * as localizedConstants from './localizedConstants'; const localize = nls.loadMessageBundle(); @@ -43,7 +44,17 @@ export class AccountFeature implements StaticFeature { let timeToLiveInSeconds = 10; this.tokenCache = new DataItemCache(this.getToken, timeToLiveInSeconds); this._client.onRequest(contracts.SecurityTokenRequest.type, async (request): Promise => { - return this.tokenCache.getData(request); + return await this.tokenCache.getData(request); + }); + this._client.onNotification(contracts.RefreshTokenNotification.type, async (request) => { + // Refresh token, then inform client the token has been updated. This is done as separate notification messages due to the synchronous processing nature of STS currently https://github.com/microsoft/azuredatastudio/issues/17179 + let result = await this.refreshToken(request); + if (!result) { + void window.showErrorMessage(localizedConstants.tokenRefreshFailed('autocompletion')); + console.log(`Token Refresh Failed ${request.toString()}`); + throw Error(localizedConstants.tokenRefreshFailed('autocompletion')); + } + this._client.sendNotification(contracts.TokenRefreshedNotification.type, result); }); } @@ -92,6 +103,38 @@ export class AccountFeature implements StaticFeature { return params; } + protected async refreshToken(request: contracts.RefreshTokenParams): Promise { + + // find account + const accountList = await azdata.accounts.getAllAccounts(); + const account = accountList.find(a => a.key.accountId === request.accountId); + if (account) { + console.log(`Failed to find azure account ${request.accountId} when executing token refresh`); + throw Error(localizedConstants.failedToFindAccount(request.accountId)); + } + + // find tenant + const tenant = account.properties.tenants.find(tenant => tenant.id === request.tenantId); + if (!tenant) { + console.log(`Failed to find tenant ${request.tenantId} in account ${account.displayInfo.displayName} when refreshing security token`); + throw Error(localizedConstants.failedToFindTenants(request.tenantId, account.displayInfo.displayName)); + } + + // Get the updated token, which will handle refreshing it if necessary + const securityToken = await azdata.accounts.getAccountSecurityToken(account, tenant.id, azdata.AzureResource.ResourceManagement); + if (!securityToken) { + console.log('Editor token refresh failed, autocompletion will be disabled until the editor is disconnected and reconnected'); + throw Error(localizedConstants.tokenRefreshFailedNoSecurityToken); + } + let params: contracts.TokenRefreshedParams = { + token: securityToken.token, + expiresOn: securityToken.expiresOn, + uri: request.uri + }; + + return params; + } + static AccountQuickPickItem = class implements QuickPickItem { account: azdata.Account; label: string; diff --git a/extensions/mssql/src/localizedConstants.ts b/extensions/mssql/src/localizedConstants.ts index 372be4472a..f7bdd9b1fd 100644 --- a/extensions/mssql/src/localizedConstants.ts +++ b/extensions/mssql/src/localizedConstants.ts @@ -55,3 +55,10 @@ export function sparkJobSubmissionGetApplicationIdFailed(err: string): string { export function sparkJobSubmissionLocalFileNotExisted(path: string): string { return localize('sparkJobSubmission.LocalFileNotExisted', "Local file {0} does not existed. ", path); } export const sparkJobSubmissionNoSqlBigDataClusterFound = localize('sparkJobSubmission.NoSqlBigDataClusterFound', "No SQL Server Big Data Cluster found."); export function sparkConnectionRequired(name: string): string { return localize('sparkConnectionRequired', "Please connect to the Spark cluster before View {0} History.", name); } + + +export function failedToFindTenants(tenantId: string, accountName: string): string { return localize('mssql.failedToFindTenants', "Failed to find tenant '{0}' in account '{1}' when refreshing security token", tenantId, accountName); } +export function tokenRefreshFailed(name: string): string { return localize('mssql.tokenRefreshFailed', "{0} AAD token refresh failed, please reconnect to enable {0}", name); } +export const tokenRefreshFailedNoSecurityToken = localize('mssql.tokenRefreshFailedNoSecurityToken', "Editor token refresh failed, autocompletion will be disabled until the editor is disconnected and reconnected"); +export function failedToFindAccount(accountName: string) { return localize('mssql.failedToFindAccount', "Failed to find azure account {0} when executing token refresh", accountName); } +