ckaczor/azuredatastudio
1
0
Fork 0
mirror of https://github.com/ckaczor/azuredatastudio.git synced 2026-02-16 18:46:40 -05:00
Code Issues Packages Projects Releases Wiki Activity

Fix access token refresh design (ADS only) (#21206)

Browse Source
This commit is contained in:
Cheena Malhotra
2022-11-21 19:35:00 -08:00
committed by GitHub
parent 8db40ab55f
commit 3b6ce47acc
5 changed files with 67 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/sql/platform/connection/common/connectionManagement.ts
View File

@@ -176,7 +176,7 @@ export interface IConnectionManagementService {
isConnected(fileUri: string): boolean; isConnected(fileUri: string): boolean;
refreshAzureAccountTokenIfNecessary(uri: string): Promise<boolean>; refreshAzureAccountTokenIfNecessary(uriOrConnectionProfile: string | ConnectionProfile): Promise<boolean>;
/** /**
* Returns true if the connection profile is connected * Returns true if the connection profile is connected
*/ */

2
src/sql/platform/connection/test/common/testConnectionManagementService.ts
View File

@@ -319,7 +319,7 @@ export class TestConnectionManagementService implements IConnectionManagementSer
return undefined!; return undefined!;
} }
refreshAzureAccountTokenIfNecessary(uri: string): Promise<boolean> { refreshAzureAccountTokenIfNecessary(uriOrConnectionProfile: string | ConnectionProfile): Promise<boolean> {
return undefined; return undefined;
} }

2
src/sql/workbench/contrib/objectExplorer/test/browser/connectionTreeActions.test.ts
View File

@@ -403,6 +403,7 @@ suite('SQL Connection Tree Action tests', () => {
resolve(connection); resolve(connection);
})); }));
connectionManagementService.setup(x => x.isConnected(undefined, TypeMoq.It.isAny())).returns(() => isConnectedReturnValue); connectionManagementService.setup(x => x.isConnected(undefined, TypeMoq.It.isAny())).returns(() => isConnectedReturnValue);
connectionManagementService.setup(x => x.refreshAzureAccountTokenIfNecessary(TypeMoq.It.isAny())).returns(async () => true);
let objectExplorerSession = { let objectExplorerSession = {
success: true, success: true,
@@ -449,6 +450,7 @@ suite('SQL Connection Tree Action tests', () => {
return refreshAction.run().then((value) => { return refreshAction.run().then((value) => {
connectionManagementService.verify(x => x.isConnected(undefined, TypeMoq.It.isAny()), TypeMoq.Times.atLeastOnce()); connectionManagementService.verify(x => x.isConnected(undefined, TypeMoq.It.isAny()), TypeMoq.Times.atLeastOnce());
connectionManagementService.verify(x => x.refreshAzureAccountTokenIfNecessary(TypeMoq.It.isAny()), TypeMoq.Times.atLeastOnce());
objectExplorerService.verify(x => x.getObjectExplorerNode(TypeMoq.It.isAny()), TypeMoq.Times.atLeastOnce()); objectExplorerService.verify(x => x.getObjectExplorerNode(TypeMoq.It.isAny()), TypeMoq.Times.atLeastOnce());
objectExplorerService.verify(x => x.refreshTreeNode(TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.exactly(0)); objectExplorerService.verify(x => x.refreshTreeNode(TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.exactly(0));
tree.verify(x => x.refresh(TypeMoq.It.isAny()), TypeMoq.Times.atLeastOnce()); tree.verify(x => x.refresh(TypeMoq.It.isAny()), TypeMoq.Times.atLeastOnce());

97
src/sql/workbench/services/connection/browser/connectionManagementService.ts
View File

@@ -906,14 +906,15 @@ export class ConnectionManagementService extends Disposable implements IConnecti
} }
const tenantId = connection.azureTenantId; const tenantId = connection.azureTenantId;
const token = await this._accountManagementService.getAccountSecurityToken(account, tenantId, azureResource); const token = await this._accountManagementService.getAccountSecurityToken(account, tenantId, azureResource);
this._logService.debug(`Got token for tenant ${token}`);
if (!token) { if (!token) {
this._logService.info(`No security tokens found for account`); this._logService.warn(`No security tokens found for account`);
} else {
this._logService.debug(`Got access token for tenant ${tenantId} that expires in ${(token.expiresOn - new Date().getTime()) / 1000} seconds`);
connection.options['azureAccountToken'] = token.token;
connection.options['expiresOn'] = token.expiresOn;
connection.options['password'] = '';
return true;
} }
connection.options['azureAccountToken'] = token.token;
connection.options['expiresOn'] = token.expiresOn;
connection.options['password'] = '';
return true;
} else { } else {
this._logService.info(`Could not find Azure account with name ${accountId}`); this._logService.info(`Could not find Azure account with name ${accountId}`);
} }
@@ -925,56 +926,78 @@ export class ConnectionManagementService extends Disposable implements IConnecti
/** /**
* Refresh Azure access token if it's expired. * Refresh Azure access token if it's expired.
* @param uri connection uri * @param uriOrConnectionProfile connection uri or connection profile
* @returns true if no need to refresh or successfully refreshed token * @returns true if no need to refresh or successfully refreshed token
*/ */
public async refreshAzureAccountTokenIfNecessary(uri: string): Promise<boolean> { public async refreshAzureAccountTokenIfNecessary(uriOrConnectionProfile: string | ConnectionProfile): Promise<boolean> {
const profile = this._connectionStatusManager.getConnectionProfile(uri); if (!uriOrConnectionProfile) {
if (!profile) { this._logService.warn(`refreshAzureAccountTokenIfNecessary: Neither Connection uri nor connection profile received.`);
this._logService.warn(`Connection not found for uri ${uri}`);
return false; return false;
} }
//wait for the pending reconnction promise if any let uri: string;
let connectionProfile: ConnectionProfile;
if (typeof uriOrConnectionProfile === 'string') {
uri = uriOrConnectionProfile;
connectionProfile = this._connectionStatusManager.getConnectionProfile(uri);
if (!connectionProfile) {
this._logService.warn(`Connection not found for uri ${uri} when refreshing token`);
return false;
}
} else {
connectionProfile = uriOrConnectionProfile;
uri = this.getConnectionUri(connectionProfile);
}
// Wait for the pending reconnction promise if any
// We expect uri to be defined
const previousReconnectPromise = this._uriToReconnectPromiseMap[uri]; const previousReconnectPromise = this._uriToReconnectPromiseMap[uri];
if (previousReconnectPromise) { if (previousReconnectPromise) {
this._logService.info(`Found pending reconnect promise for uri ${uri}, waiting.`); this._logService.debug(`Found pending reconnect promise for uri ${uri}, waiting.`);
try { try {
const previousConnectionResult = await previousReconnectPromise; const previousConnectionResult = await previousReconnectPromise;
if (previousConnectionResult && previousConnectionResult.connected) { if (previousConnectionResult && previousConnectionResult.connected) {
this._logService.info(`Previous pending reconnection for uri ${uri} succeeded.`); this._logService.debug(`Previous pending reconnection for uri ${uri} succeeded.`);
return true; return true;
} }
this._logService.info(`Previous pending reconnection for uri ${uri} failed.`); this._logService.debug(`Previous pending reconnection for uri ${uri} failed.`);
} catch (err) { } catch (err) {
this._logService.info(`Previous pending reconnect promise for uri ${uri} is rejected with error ${err}, will attempt to reconnect if necessary.`); this._logService.debug(`Previous pending reconnect promise for uri ${uri} is rejected with error ${err}, will attempt to reconnect if necessary.`);
} }
} }
const expiry = profile.options.expiresOn; // We expect connectionProfile to be defined
if (typeof expiry === 'number' && !Number.isNaN(expiry)) { if (connectionProfile && connectionProfile.authenticationType === Constants.AuthenticationType.AzureMFA) {
const currentTime = new Date().getTime() / 1000; const expiry = connectionProfile.options.expiresOn;
const maxTolerance = 2 * 60; // two minutes if (typeof expiry === 'number' && !Number.isNaN(expiry)) {
if (expiry - currentTime < maxTolerance) { const currentTime = new Date().getTime() / 1000;
this._logService.info(`Access token expired for connection ${profile.id} with uri ${uri}`); const maxTolerance = 2 * 60; // two minutes
try { if (expiry - currentTime < maxTolerance) {
const connectionResultPromise = this.connect(profile, uri); this._logService.debug(`Access token expired for connection ${connectionProfile.id} with uri ${uri}`);
this._uriToReconnectPromiseMap[uri] = connectionResultPromise; try {
const connectionResult = await connectionResultPromise; const connectionResultPromise = this.connect(connectionProfile, uri);
if (!connectionResult) { this._uriToReconnectPromiseMap[uri] = connectionResultPromise;
this._logService.error(`Failed to refresh connection ${profile.id} with uri ${uri}, invalid connection result.`); const connectionResult = await connectionResultPromise;
throw new Error(nls.localize('connection.invalidConnectionResult', "Connection result is invalid")); if (!connectionResult) {
} else if (!connectionResult.connected) { this._logService.error(`Failed to refresh connection ${connectionProfile.id} with uri ${uri}, invalid connection result.`);
this._logService.error(`Failed to refresh connection ${profile.id} with uri ${uri}, error code: ${connectionResult.errorCode}, error message: ${connectionResult.errorMessage}`); throw new Error(nls.localize('connection.invalidConnectionResult', "Connection result is invalid"));
throw new Error(nls.localize('connection.refreshAzureTokenFailure', "Failed to refresh Azure account token for connection")); } else if (!connectionResult.connected) {
this._logService.error(`Failed to refresh connection ${connectionProfile.id} with uri ${uri}, error code: ${connectionResult.errorCode}, error message: ${connectionResult.errorMessage}`);
throw new Error(nls.localize('connection.refreshAzureTokenFailure', "Failed to refresh Azure account token for connection"));
}
this._logService.debug(`Successfully refreshed token for connection ${connectionProfile.id} with uri ${uri}, result: ${connectionResult.connected} ${connectionResult.connectionProfile}, ${this._connectionStatusManager.getConnectionProfile(uri)}`);
return true;
} finally {
delete this._uriToReconnectPromiseMap[uri];
} }
this._logService.info(`Successfully refreshed token for connection ${profile.id} with uri ${uri}, result: ${connectionResult.connected} ${connectionResult.connectionProfile}, isConnected: ${this.isConnected(uri)}, ${this._connectionStatusManager.getConnectionProfile(uri)}`);
return true;
} finally {
delete this._uriToReconnectPromiseMap[uri];
} }
else {
this._logService.debug(`No need to refresh Azure acccount token for connection ${connectionProfile.id} with uri ${uri}`);
}
} else {
this._logService.warn(`Invalid expiry time ${expiry} for connection ${connectionProfile.id} with uri ${uri}`);
} }
this._logService.info(`No need to refresh Azure acccount token for connection ${profile.id} with uri ${uri}`);
} }
return true; return true;
} }

3
src/sql/workbench/services/objectExplorer/browser/connectionTreeAction.ts
View File

@@ -48,6 +48,7 @@ export class RefreshAction extends Action {
if (this.element instanceof ConnectionProfile) { if (this.element instanceof ConnectionProfile) {
let connection: ConnectionProfile = this.element; let connection: ConnectionProfile = this.element;
if (this._connectionManagementService.isConnected(undefined, connection)) { if (this._connectionManagementService.isConnected(undefined, connection)) {
await this._connectionManagementService.refreshAzureAccountTokenIfNecessary(connection);
treeNode = this._objectExplorerService.getObjectExplorerNode(connection); treeNode = this._objectExplorerService.getObjectExplorerNode(connection);
if (treeNode === undefined) { if (treeNode === undefined) {
await this._objectExplorerService.updateObjectExplorerNodes(connection.toIConnectionProfile()); await this._objectExplorerService.updateObjectExplorerNodes(connection.toIConnectionProfile());
@@ -55,6 +56,8 @@ export class RefreshAction extends Action {
} }
} }
} else if (this.element instanceof TreeNode) { } else if (this.element instanceof TreeNode) {
let connection: ConnectionProfile = this.element.getConnectionProfile();
this._connectionManagementService.refreshAzureAccountTokenIfNecessary(connection);
treeNode = this.element; treeNode = this.element;
} }