Adding centralized cache for Entra tokens #21434
+409
−70
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -43,6 +43,7 @@ import { | |
| TelemetryActions, | ||
| TelemetryViews, | ||
| } from "../sharedInterfaces/telemetry"; | ||
| import { ApiStatus, isStatus, Status } from "../sharedInterfaces/webview"; | ||
| import { ObjectExplorerUtils } from "../objectExplorer/objectExplorerUtils"; | ||
| import { changeLanguageServiceForFile } from "../languageservice/utils"; | ||
| import { AddFirewallRuleWebviewController } from "./addFirewallRuleWebviewController"; | ||
|
|
@@ -120,6 +121,11 @@ export default class ConnectionManager { | |
| Deferred<ConnectionContracts.ConnectionCompleteParams> | ||
| >; | ||
| private _keyVaultTokenCache: Map<string, IToken> = new Map<string, IToken>(); | ||
| private _entraSqlTokenCache: Map<string, IToken> = new Map<string, IToken>(); | ||
| private _entraSqlTokenRefreshInFlight: Map<string, Promise<IToken>> = new Map< | ||
| string, | ||
| Promise<IToken> | ||
| >(); | ||
| private _accountService: AccountService; | ||
| private _firewallService: FirewallService; | ||
| public azureController: AzureController; | ||
|
|
@@ -1028,19 +1034,23 @@ export default class ConnectionManager { | |
| * Does nothing if connection is not using Entra auth. | ||
| * throws if token refresh fails or if account/profile cannot be found. | ||
| */ | ||
| public async refreshEntraTokenIfNeeded(connectionInfo: IConnectionInfo) { | ||
| // 1. Validate that the connection is using Entra auth | ||
| if (connectionInfo.authenticationType !== Constants.azureMfa) { | ||
| return; | ||
| } | ||
|
|
||
| // 2. Validate that the token needs refreshing (isn't expired) | ||
| if ( | ||
| AzureController.isTokenValid(connectionInfo.azureAccountToken, connectionInfo.expiresOn) | ||
| ) { | ||
| this._logger?.verbose( | ||
| `Entra token for account ${connectionInfo.user} (${connectionInfo.email}) is still valid until ${connectionInfo.expiresOn}. No refresh needed.`, | ||
| ); | ||
| return; | ||
| } | ||
|
|
||
| // 3. Collect Entra account information | ||
| let account: IAccount; | ||
| let profile: ConnectionProfile; | ||
|
|
||
|
|
@@ -1052,22 +1062,53 @@ export default class ConnectionManager { | |
| sendErrorEvent( | ||
| TelemetryViews.ConnectionManager, | ||
| TelemetryActions.Connect, | ||
| new Error("Azure MFA connection missing accountId in refreshEntraTokenIfNeeded"), | ||
| true, // includeErrorMessage | ||
| ); | ||
| throw new Error(LocalizedConstants.cannotConnect); | ||
| } | ||
|
|
||
| if (!account) { | ||
| this._logger?.verbose( | ||
| `No account found in account store for accountId ${connectionInfo.accountId}. Cannot refresh Entra token.`, | ||
| ); | ||
| throw new Error(LocalizedConstants.msgAccountNotFound); | ||
| } | ||
|
|
||
| connectionInfo.user = account.displayInfo.displayName; | ||
| connectionInfo.email = account.displayInfo.email; | ||
| profile.user = account.displayInfo.displayName; | ||
| profile.email = account.displayInfo.email; | ||
|
|
||
| // 4. Use cached token if present and valid/unexpired | ||
| const cacheKey = this.getEntraSqlTokenCacheKey( | ||
| connectionInfo, | ||
| account.properties?.owningTenant?.id, | ||
| ); | ||
| const cachedToken = this._entraSqlTokenCache.get(cacheKey); | ||
|
|
||
| this._logger?.verbose( | ||
| `Cached token ${cachedToken ? "found" : "not found"} for cache key ${cacheKey}.`, | ||
| ); | ||
|
|
||
| if (cachedToken) { | ||
| // If there's a cached token, use it if still valid, or remove it from cache if expired | ||
| if (AzureController.isTokenValid(cachedToken.token, cachedToken.expiresOn)) { | ||
| this.applyEntraToken(connectionInfo, cachedToken); | ||
| this._logger?.verbose( | ||
| `Using cached Entra token for account ${account.displayInfo.displayName} (${account.displayInfo.email}) and tenant ${profile.tenantId}. Cached token expires on ${cachedToken.expiresOn}. (currently ${Date.now() / 1000})`, | ||
| ); | ||
|
|
||
| return; | ||
Benjin marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| } else { | ||
| this._logger?.verbose( | ||
| `Cached token for cache key ${cacheKey} is expired. Removing from cache. (currently ${Date.now() / 1000})`, | ||
| ); | ||
| this._entraSqlTokenCache.delete(cacheKey); | ||
| } | ||
| } | ||
|
|
||
| // 5. Lastly, refresh the token, cache the new token, and update the connection info with it | ||
| const refreshTask = async () => { | ||
| return await this.azureController.refreshAccessToken( | ||
| account, | ||
|
|
@@ -1077,63 +1118,109 @@ export default class ConnectionManager { | |
| ); | ||
| }; | ||
|
|
||
| // Dedupe concurrent token refresh requests for the same account into a single request, and share the result | ||
| let refreshPromise = this._entraSqlTokenRefreshInFlight.get(cacheKey); | ||
| if (!refreshPromise) { | ||
| // Token refresh code cannot figure out if the user closed the browser window, | ||
| // so we wrap it in a cancellable progress dialog to allow the user to cancel | ||
| // the operation. | ||
| refreshPromise = new Promise<IToken>((resolve, reject) => { | ||
| void vscode.window.withProgress( | ||
| { | ||
| location: vscode.ProgressLocation.Notification, | ||
| title: LocalizedConstants.ObjectExplorer.AzureSignInMessage( | ||
| account.displayInfo.displayName || account.displayInfo.email, | ||
| ), | ||
| cancellable: true, | ||
| }, | ||
| async (_progress, token) => { | ||
| token.onCancellationRequested(() => { | ||
| reject({ | ||
| status: ApiStatus.Cancelled, | ||
| message: "Azure sign in cancelled by user.", | ||
| } as Status); | ||
| }); | ||
| try { | ||
| const refreshedToken = await refreshTask(); | ||
| if (!refreshedToken) { | ||
| reject({ | ||
| status: ApiStatus.Error, | ||
| message: LocalizedConstants.msgAccountRefreshFailed(), | ||
| } as Status); | ||
| return; | ||
| } | ||
| this._logger?.verbose( | ||
| `Successfully refreshed Entra token for account ${account.displayInfo.displayName} (${account.displayInfo.email}) and tenant ${profile.tenantId}; now expires on ${refreshedToken.expiresOn} (currently ${Date.now() / 1000}).`, | ||
| ); | ||
| resolve(refreshedToken); | ||
Benjin marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| } catch (error) { | ||
| const refreshErrorStatus: Status = { | ||
| status: ApiStatus.Error, | ||
| message: getErrorMessage(error), | ||
| }; | ||
| this._logger?.error( | ||
| `Error refreshing Entra token for account ${account.displayInfo.displayName} (${account.displayInfo.email}) and tenant ${profile.tenantId}: ${refreshErrorStatus.message}`, | ||
| ); | ||
| reject(refreshErrorStatus); | ||
| } | ||
|
Comment on lines
+1137
to
+1165
There was a problem hiding this comment. After cancellation triggers
Contributor
Author
There was a problem hiding this comment. @copilot open a new pull request to apply changes based on this feedback |
||
| }, | ||
| ); | ||
| }).finally(() => { | ||
| this._entraSqlTokenRefreshInFlight.delete(cacheKey); | ||
| }); | ||
| this._entraSqlTokenRefreshInFlight.set(cacheKey, refreshPromise); | ||
| } | ||
|
|
||
| try { | ||
| const azureAccountToken = await refreshPromise; | ||
| this.applyEntraToken(connectionInfo, azureAccountToken); | ||
| // Save refreshed token so other connections for the same account+tenant can reuse it. | ||
| this._entraSqlTokenCache.set(cacheKey, azureAccountToken); | ||
| this._logger?.verbose( | ||
| `Successfully refreshed Entra token for account ${account.displayInfo.displayName} (${account.displayInfo.email}) and tenant ${profile.tenantId}. Cached token for future use with cache key ${cacheKey}.`, | ||
| ); | ||
| } catch (error) { | ||
| this._logger?.verbose( | ||
| `Failed to refresh Entra token for account ${account.displayInfo.displayName} (${account.displayInfo.email}) and tenant ${profile.tenantId}. Error: ${getErrorMessage(error)}`, | ||
| ); | ||
| if (isStatus(error)) { | ||
| if (error.status === ApiStatus.Cancelled) { | ||
| this._logger.verbose("Refresh cancelled: " + error.message); | ||
| throw new Error(LocalizedConstants.cannotConnect); | ||
| } | ||
|
|
||
| if (error.status === ApiStatus.Error) { | ||
| const message = LocalizedConstants.msgAccountRefreshFailed(error.message); | ||
| this._logger.error("Error refreshing account: " + message); | ||
| await this.vscodeWrapper.showErrorMessage(message); | ||
| throw new Error(message); | ||
Benjin marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| } | ||
| } | ||
|
|
||
| throw error; | ||
| } | ||
| } | ||
|
|
||
| private getEntraSqlTokenCacheKey( | ||
| connectionInfo: IConnectionInfo, | ||
| defaultTenantId?: string, | ||
| ): string { | ||
| return `${connectionInfo.accountId ?? ""}|${connectionInfo.tenantId ?? defaultTenantId ?? ""}`; | ||
| } | ||
|
|
||
| private applyEntraToken(connectionInfo: IConnectionInfo, token: IToken): void { | ||
| connectionInfo.azureAccountToken = token.token; | ||
| connectionInfo.expiresOn = token.expiresOn; | ||
| } | ||
|
|
||
| /** | ||
| * Clears both token entries and any in-flight refresh promises. | ||
| */ | ||
| private clearEntraSqlTokenCache(): void { | ||
| this._entraSqlTokenCache.clear(); | ||
| this._entraSqlTokenRefreshInFlight.clear(); | ||
| } | ||
|
|
||
| /** | ||
| * Handles password-based credential authentication by prompting for password if needed. | ||
| * This method checks if a password is required and prompts the user if it's not saved or available. | ||
|
|
@@ -1422,7 +1509,7 @@ export default class ConnectionManager { | |
| // Handle Entra token validity | ||
| if (credentials.authenticationType === Constants.azureMfa) { | ||
| try { | ||
| await this.refreshEntraTokenIfNeeded(credentials); | ||
| } catch (error) { | ||
| telemetryActivity?.endFailed( | ||
| error, | ||
|
|
@@ -1653,14 +1740,17 @@ export default class ConnectionManager { | |
| }; | ||
| } else if (errorType === SqlConnectionErrorType.EntraTokenExpired) { | ||
| try { | ||
| await this.refreshEntraTokenIfNeeded(credentials); | ||
| return { | ||
| isHandled: true, | ||
| updatedCredentials: credentials, | ||
| errorHandled: SqlConnectionErrorType.EntraTokenExpired, | ||
| }; | ||
| } catch (error) { | ||
| const errorMessage = getErrorMessage(error); | ||
| if (errorMessage !== LocalizedConstants.cannotConnect) { | ||
| Utils.showErrorMsg(errorMessage); | ||
| } | ||
Benjin marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| return { | ||
| isHandled: false, | ||
| updatedCredentials: credentials, | ||
|
|
@@ -1853,7 +1943,7 @@ export default class ConnectionManager { | |
| // No connection for this URI, nothing to do | ||
| return; | ||
| } | ||
| await this.refreshEntraTokenIfNeeded(connectionInfo.credentials); | ||
| } | ||
|
|
||
| public async addAccount(): Promise<IAccount> { | ||
|
|
@@ -1907,6 +1997,7 @@ export default class ConnectionManager { | |
|
|
||
| public onClearAzureTokenCache(): void { | ||
| this.azureController.clearTokenCache(); | ||
| this.clearEntraSqlTokenCache(); | ||
| this.vscodeWrapper.showInformationMessage( | ||
| LocalizedConstants.Accounts.clearedEntraTokenCache, | ||
| ); | ||
|
|
||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.