Refresh expired azure account tokens before running query (#17137)

* wip

* wip

* fix merge conflict

* wip

* wip

* add expiresOn

* refresh azure account token when running query

* remove package-lock.json

* cleanup

* cleanup

* fix lint

* pr comments & localization

* fix lint

* cleanup

* pr review comments

localization/xliff/enu/constants/localizedConstants.enu.xlf

@@ -329,6 +329,36 @@
         <trans-unit id="msgConnecting">
             <source xml:lang="en">Connecting to server "{0}" on document "{1}".</source>
         </trans-unit>
+        <trans-unit id="msgConnectionNotFound">
+            <source xml:lang="en">Connection not found for uri "{0}".</source>
+        </trans-unit>
+        <trans-unit id="msgFoundPendingReconnect">
+            <source xml:lang="en">Found pending reconnect promise for uri {0}, waiting.</source>
+        </trans-unit>
+        <trans-unit id="msgPendingReconnectSuccess">
+            <source xml:lang="en">Previous pending reconnection for uri {0}, succeeded.</source>
+        </trans-unit>
+        <trans-unit id="msgFoundPendingReconnectFailed">
+            <source xml:lang="en">Found pending reconnect promise for uri {0}, failed.</source>
+        </trans-unit>
+        <trans-unit id="msgFoundPendingReconnectError">
+            <source xml:lang="en">Previous pending reconnect promise for uri {0} is rejected with error {1}, will attempt to reconnect if necessary.</source>
+        </trans-unit>
+        <trans-unit id="msgAcessTokenExpired">
+            <source xml:lang="en">Access token expired for connection {0} with uri {1}</source>
+        </trans-unit>
+        <trans-unit id="msgRefreshTokenError">
+            <source xml:lang="en">Error when refreshing token</source>
+        </trans-unit>
+        <trans-unit id="msgRefreshConnection">
+            <source xml:lang="en">Failed to refresh connection ${0} with uri {1}, invalid connection result.</source>
+        </trans-unit>
+        <trans-unit id="msgRefreshTokenSuccess">
+            <source xml:lang="en">Successfully refreshed token for connection {0} with uri {1}, {2}</source>
+        </trans-unit>
+        <trans-unit id="msgRefreshTokenNotNeeded">
+            <source xml:lang="en">No need to refresh Azure acccount token for connection {0} with uri {1}</source>
+        </trans-unit>
         <trans-unit id="msgConnectedServerInfo">
             <source xml:lang="en">Connected to server "{0}" on document "{1}". Server information: {2}</source>
         </trans-unit>

src/azure/accountService.ts

@@ -10,7 +10,7 @@ import * as Constants from '../constants/constants';
 import { AzureController } from './azureController';
 import { AccountStore } from './accountStore';
 import providerSettings from '../azure/providerSettings';
-import { Tenant } from '@microsoft/ads-adal-library';
+import { Tenant, Token } from '@microsoft/ads-adal-library';
 
 export class AccountService {
 
@@ -70,12 +70,12 @@ export class AccountService {
         // TODO: match type for mapping in mssql and sqltoolsservice
         let mapping = {};
         mapping[this.getHomeTenant(this.account).id] = {
-            token: await this.refreshToken(this.account)
+            token: (await this.refreshToken(this.account)).token
         };
         return mapping;
     }
 
-    public async refreshToken(account): Promise<string> {
+    public async refreshToken(account): Promise<Token> {
         return await this._azureController.refreshToken(account, this._accountStore, providerSettings.resources.azureManagementResource);
     }
 

src/azure/azureController.ts

@@ -108,6 +108,7 @@ export class AzureController {
                 this._vscodeWrapper.showErrorMessage(errorMessage);
             }
             profile.azureAccountToken = token.token;
+            profile.expiresOn = token.expiresOn;
             profile.email = account.displayInfo.email;
             profile.accountId = account.key.id;
         } else if (config === utils.azureAuthTypeToString(AzureAuthType.DeviceCode)) {
@@ -122,6 +123,7 @@ export class AzureController {
                 this._vscodeWrapper.showErrorMessage(errorMessage);
             }
             profile.azureAccountToken = token.token;
+            profile.expiresOn = token.expiresOn;
             profile.email = account.displayInfo.email;
             profile.accountId = account.key.id;
         }
@@ -146,13 +148,14 @@ export class AzureController {
                 }
             });
         }
-        profile.azureAccountToken = azureAccountToken;
+        profile.azureAccountToken = azureAccountToken.token;
+        profile.expiresOn = azureAccountToken.expiresOn;
         profile.email = account.displayInfo.email;
         profile.accountId = account.key.id;
         return profile;
     }
 
-    public async refreshToken(account: IAccount, accountStore: AccountStore, settings: AADResource): Promise<string | undefined> {
+    public async refreshToken(account: IAccount, accountStore: AccountStore, settings: AADResource): Promise<Token | undefined> {
         try {
             let token: Token;
             if (account.properties.azureAuthType === 0) {
@@ -175,7 +178,7 @@ export class AzureController {
                 token = await azureDeviceCode.getAccountSecurityToken(
                     account, azureDeviceCode.getHomeTenant(account).id, providerSettings.resources.databaseResource);
             }
-            return token.token;
+            return token;
         } catch (ex) {
             let errorMsg = this.azureErrorLookup.getSimpleError(ex.errorCode);
             this._vscodeWrapper.showErrorMessage(errorMsg);

src/controllers/connectionManager.ts

@@ -354,10 +354,11 @@ export default class ConnectionManager {
         };
     }
 
-    private handleConnectionSuccess(fileUri: string,
-                                    connection: ConnectionInfo,
-                                    newCredentials: IConnectionInfo,
-                                    result: ConnectionContracts.ConnectionCompleteParams): void {
+    private handleConnectionSuccess(
+        fileUri: string,
+        connection: ConnectionInfo,
+        newCredentials: IConnectionInfo,
+        result: ConnectionContracts.ConnectionCompleteParams): void {
         connection.connectionId = result.connectionId;
         connection.serverInfo = result.serverInfo;
         connection.credentials = newCredentials;
@@ -398,7 +399,7 @@ export default class ConnectionManager {
                     await vscode.env.openExternal(vscode.Uri.parse(Constants.integratedAuthHelpLink));
                 }
             } else if (platformInfo.runtimeId === Runtime.OSX_10_11_64 &&
-            result.messages.indexOf('Unable to load DLL \'System.Security.Cryptography.Native\'') !== -1) {
+                result.messages.indexOf('Unable to load DLL \'System.Security.Cryptography.Native\'') !== -1) {
                 const action = await this.vscodeWrapper.showErrorMessage(Utils.formatString(LocalizedConstants.msgConnectionError2,
                     LocalizedConstants.macOpenSslErrorMessage), LocalizedConstants.macOpenSslHelpButton);
                 if (action && action === LocalizedConstants.macOpenSslHelpButton) {
@@ -507,6 +508,7 @@ export default class ConnectionManager {
      * @returns The list of databases retrieved from the connection
      */
     public async listDatabases(connectionUri: string): Promise<string[]> {
+        await this.refreshAzureAccountToken(connectionUri);
         const listParams = new ConnectionContracts.ListDatabasesParams();
         listParams.ownerUri = connectionUri;
         const result = await this.client.sendRequest(ConnectionContracts.ListDatabasesRequest.type, listParams);
@@ -534,11 +536,11 @@ export default class ConnectionManager {
         if (fileUri && this._vscodeWrapper.isEditingSqlFile) {
             if (isSqlCmdMode) {
                 SqlToolsServerClient.instance.sendNotification(LanguageServiceContracts.LanguageFlavorChangedNotification.type,
-                    <LanguageServiceContracts.DidChangeLanguageFlavorParams> {
-                    uri: fileUri,
-                    language: isSqlCmd ? 'sqlcmd' : 'sql',
-                    flavor: 'MSSQL'
-                });
+                    <LanguageServiceContracts.DidChangeLanguageFlavorParams>{
+                        uri: fileUri,
+                        language: isSqlCmd ? 'sqlcmd' : 'sql',
+                        flavor: 'MSSQL'
+                    });
                 return true;
             }
             const flavor = await this._connectionUI.promptLanguageFlavor();
@@ -547,11 +549,11 @@ export default class ConnectionManager {
             }
             this.statusView.languageFlavorChanged(fileUri, flavor);
             SqlToolsServerClient.instance.sendNotification(LanguageServiceContracts.LanguageFlavorChangedNotification.type,
-                <LanguageServiceContracts.DidChangeLanguageFlavorParams> {
-                uri: fileUri,
-                language: 'sql',
-                flavor: flavor
-            });
+                <LanguageServiceContracts.DidChangeLanguageFlavorParams>{
+                    uri: fileUri,
+                    language: 'sql',
+                    flavor: flavor
+                });
             return true;
         } else {
             await this._vscodeWrapper.showWarningMessage(LocalizedConstants.msgOpenSqlFile);
@@ -675,7 +677,9 @@ export default class ConnectionManager {
         const self = this;
         // Check if the azure account token is present before sending connect request
         if (connectionCreds.authenticationType === Constants.azureMfa) {
-            if (!connectionCreds.azureAccountToken) {
+            const currentTime = new Date().getTime() / 1000;
+            const maxTolerance = 2 * 60; // two minutes
+            if (!connectionCreds.azureAccountToken || connectionCreds.expiresOn - currentTime < maxTolerance) {
                 let account = this.accountStore.getAccount(connectionCreds.accountId);
                 let profile = new ConnectionProfile(connectionCreds);
                 let azureAccountToken = await this.azureController.refreshToken(account, this.accountStore, providerSettings.resources.databaseResource);
@@ -690,7 +694,8 @@ export default class ConnectionManager {
                         throw new Error(`${LocalizedConstants.cannotConnect}`);
                     }
                 } else {
-                    connectionCreds.azureAccountToken = azureAccountToken;
+                    connectionCreds.azureAccountToken = azureAccountToken.token;
+                    connectionCreds.expiresOn = azureAccountToken.expiresOn;
                 }
             }
         }
@@ -795,7 +800,7 @@ export default class ConnectionManager {
 
     public onDidOpenTextDocument(doc: vscode.TextDocument): void {
         let uri = doc.uri.toString(true);
-        if (doc.languageId === 'sql' && typeof(this._connections[uri]) === 'undefined') {
+        if (doc.languageId === 'sql' && typeof (this._connections[uri]) === 'undefined') {
             this.statusView.notConnected(uri);
         }
     }
@@ -814,6 +819,53 @@ export default class ConnectionManager {
         }
     }
 
+    public async refreshAzureAccountToken(uri: string): Promise<void> {
+        const profile = this.getConnectionInfo(uri);
+        if (!profile) {
+            this.vscodeWrapper.logToOutputChannel(Utils.formatString(LocalizedConstants.msgConnectionNotFound, uri));
+            return;
+        }
+
+        // Wait for the pending reconnction promise if any
+        const previousReconnectPromise = this._uriToConnectionPromiseMap.get(uri);
+        if (previousReconnectPromise) {
+            this.vscodeWrapper.logToOutputChannel(Utils.formatString(LocalizedConstants.msgFoundPendingReconnect, uri));
+            try {
+                const previousConnectionResult = await previousReconnectPromise;
+                if (previousConnectionResult) {
+                    this.vscodeWrapper.logToOutputChannel(Utils.formatString(LocalizedConstants.msgPendingReconnectSuccess, uri));
+                    return;
+                }
+                this.vscodeWrapper.logToOutputChannel(Utils.formatString(LocalizedConstants.msgFoundPendingReconnectFailed, uri));
+            } catch (err) {
+                this.vscodeWrapper.logToOutputChannel(Utils.formatString(LocalizedConstants.msgFoundPendingReconnectError, uri, err));
+            }
+        }
+
+        const expiry = profile.credentials.expiresOn;
+        if (typeof expiry === 'number' && !Number.isNaN(expiry)) {
+            const currentTime = new Date().getTime() / 1000;
+            const maxTolerance = 2 * 60; // two minutes
+            if (expiry - currentTime < maxTolerance) {
+                this.vscodeWrapper.logToOutputChannel(Utils.formatString(LocalizedConstants.msgAcessTokenExpired, profile.connectionId, uri));
+                try {
+                    let connectionResult = await this.connect(uri, profile.credentials);
+                    if (!connectionResult) {
+                        this.vscodeWrapper.showErrorMessage(Utils.formatString(LocalizedConstants.msgRefreshConnection, profile.connectionId, uri));
+                        throw new Error('Unable to refresh connection');
+                    }
+                    this.vscodeWrapper.logToOutputChannel(Utils.formatString(LocalizedConstants.msgRefreshTokenSuccess,
+                        profile.connectionId, uri, this.getConnectionInfo(uri)));
+                    return;
+                } catch {
+                    this.vscodeWrapper.showInformationMessage(Utils.formatString(LocalizedConstants.msgRefreshTokenError));
+                }
+            }
+            this.vscodeWrapper.logToOutputChannel(Utils.formatString(LocalizedConstants.msgRefreshTokenNotNeeded, profile.connectionId, uri));
+        }
+        return;
+    }
+
     public async removeAccount(prompter: IPrompter): Promise<void> {
         // list options for accounts to remove
         let questions: IQuestion[] = [];

src/controllers/mainController.ts

@@ -734,6 +734,8 @@ export default class MainController implements vscode.Disposable {
             if (!self.connectionManager.isConnected(uri)) {
                 await self.onNewConnection();
             }
+            // check if current connection is still valid / active - if not, refresh azure account token
+            await this._connectionMgr.refreshAzureAccountToken(uri);
 
             let title = path.basename(editor.document.fileName);
             let querySelection: ISelectionData;

src/models/connectionCredentials.ts

@@ -23,6 +23,7 @@ export class ConnectionCredentials implements IConnectionInfo {
     public port: number;
     public authenticationType: string;
     public azureAccountToken: string | undefined;
+    public expiresOn: number | undefined;
     public encrypt: boolean;
     public trustServerCertificate: boolean | undefined;
     public persistSecurityInfo: boolean | undefined;

src/models/connectionProfile.ts

@@ -26,7 +26,8 @@ export class ConnectionProfile extends ConnectionCredentials implements IConnect
     public savePassword: boolean;
     public emptyPasswordInput: boolean;
     public azureAuthType: AzureAuthType;
-    public azureAccountToken: string;
+    public azureAccountToken: string | undefined;
+    public expiresOn: number | undefined;
     public accountStore: AccountStore;
     public accountId: string;
 
@@ -36,6 +37,7 @@ export class ConnectionProfile extends ConnectionCredentials implements IConnect
             this.accountId = connectionCredentials.accountId;
             this.authenticationType = connectionCredentials.authenticationType;
             this.azureAccountToken = connectionCredentials.azureAccountToken;
+            this.expiresOn = connectionCredentials.expiresOn;
             this.database = connectionCredentials.database;
             this.email = connectionCredentials.email;
             this.password = connectionCredentials.password;

src/objectExplorer/objectExplorerService.ts

@@ -453,7 +453,8 @@ export class ObjectExplorerService {
                                 }
                             });
                         } else {
-                            connectionCredentials.azureAccountToken = azureAccountToken;
+                            connectionCredentials.azureAccountToken = azureAccountToken.token;
+                            connectionCredentials.expiresOn = azureAccountToken.expiresOn;
                         }
                     }
                 }

test/connectionProfile.test.ts

@@ -31,6 +31,7 @@ function createTestCredentials(): IConnectionInfo {
         port:                           1234,
         authenticationType:             AuthenticationTypes[AuthenticationTypes.SqlLogin],
         azureAccountToken:              '',
+        expiresOn:                      0,
         encrypt:                        false,
         trustServerCertificate:         false,
         persistSecurityInfo:            false,

test/firewallService.test.ts

@@ -12,7 +12,7 @@ import { HandleFirewallRuleRequest, IHandleFirewallRuleResponse,
 import VscodeWrapper from '../src/controllers/vscodeWrapper';
 import { assert } from 'chai';
 import { IAzureSession, IAzureResourceFilter } from '../src/models/interfaces';
-import { Tenant } from '@microsoft/ads-adal-library';
+import { Tenant, Token } from '@microsoft/ads-adal-library';
 import { IAccount } from '../src/models/contracts/azure/accountInterfaces';
 
 
@@ -89,7 +89,13 @@ suite('Firewall Service Tests', () => {
             displayInfo: undefined,
             isStale: undefined
         };
-        accountService.setup(v => v.refreshToken(mockAccount)).returns(() => Promise.resolve('mockToken'));
+        let mockToken: Token = {
+            key: '',
+            tokenType: '',
+            token: '',
+            expiresOn: 0
+        };
+        accountService.setup(v => v.refreshToken(mockAccount)).returns(() => Promise.resolve(mockToken));
         accountService.object.setAccount(mockAccount);
         let result = await firewallService.object.createFirewallRule(server, startIpAddress, endIpAddress);
         assert.isNotNull(result, 'Create Firewall Rule request is sent successfully');

test/perFileConnection.test.ts

@@ -53,6 +53,7 @@ function createTestCredentials(): IConnectionInfo {
         port:                           1234,
         authenticationType:             AuthenticationTypes[AuthenticationTypes.SqlLogin],
         azureAccountToken:              '',
+        expiresOn:                      0,
         encrypt:                        false,
         trustServerCertificate:         false,
         persistSecurityInfo:            false,

typings/vscode-mssql.d.ts

@@ -127,6 +127,11 @@ declare module 'vscode-mssql' {
          */
         azureAccountToken: string | undefined;
 
+        /**
+         * Access token expiry timestamp
+         */
+        expiresOn: number | undefined;
+
         /**
          * Gets or sets a Boolean value that indicates whether SQL Server uses SSL encryption for all data sent between the client and server if
          * the server has a certificate installed.