From f607b5abd35c1f3e35be168bb1ab52d2c9547bbf Mon Sep 17 00:00:00 2001 From: Laurent Nguyen Date: Tue, 4 Mar 2025 14:00:18 +0100 Subject: [PATCH] Add support for refreshing access tokens --- src/Contracts/FabricMessageTypes.ts | 1 + src/Contracts/FabricMessagesContract.ts | 8 +++ src/Explorer/Explorer.tsx | 4 +- src/Platform/Fabric/FabricUtil.ts | 81 +++++++++++++++++++------ src/hooks/useKnockoutExplorer.ts | 12 ++-- 5 files changed, 82 insertions(+), 24 deletions(-) diff --git a/src/Contracts/FabricMessageTypes.ts b/src/Contracts/FabricMessageTypes.ts index aa374472d..1d4576391 100644 --- a/src/Contracts/FabricMessageTypes.ts +++ b/src/Contracts/FabricMessageTypes.ts @@ -4,6 +4,7 @@ export enum FabricMessageTypes { GetAuthorizationToken = "GetAuthorizationToken", GetAllResourceTokens = "GetAllResourceTokens", + GetAccessToken = "GetAccessToken", Ready = "Ready", } diff --git a/src/Contracts/FabricMessagesContract.ts b/src/Contracts/FabricMessagesContract.ts index 8155b80f8..2cc99c578 100644 --- a/src/Contracts/FabricMessagesContract.ts +++ b/src/Contracts/FabricMessagesContract.ts @@ -73,6 +73,14 @@ export type FabricMessageV3 = message: { visible: boolean; }; + } + | { + type: "accessToken"; + message: { + id: string; + error: string | undefined; + data: { accessToken: string }; + }; }; export enum CosmosDbArtifactType { diff --git a/src/Explorer/Explorer.tsx b/src/Explorer/Explorer.tsx index 2632e6822..35e80d772 100644 --- a/src/Explorer/Explorer.tsx +++ b/src/Explorer/Explorer.tsx @@ -11,7 +11,7 @@ import { IGalleryItem } from "Juno/JunoClient"; import { isFabricMirrored, isFabricMirroredKey, - scheduleRefreshDatabaseResourceToken, + scheduleRefreshFabricToken, } from "Platform/Fabric/FabricUtil"; import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility"; import { acquireMsalTokenForAccount } from "Utils/AuthorizationUtils"; @@ -356,7 +356,7 @@ export default class Explorer { public onRefreshResourcesClick = async (): Promise => { if (isFabricMirroredKey()) { - scheduleRefreshDatabaseResourceToken(true).then(() => this.refreshAllDatabases()); + scheduleRefreshFabricToken(true).then(() => this.refreshAllDatabases()); return; } diff --git a/src/Platform/Fabric/FabricUtil.ts b/src/Platform/Fabric/FabricUtil.ts index 0372a0da7..dc8f28f44 100644 --- a/src/Platform/Fabric/FabricUtil.ts +++ b/src/Platform/Fabric/FabricUtil.ts @@ -12,27 +12,66 @@ let timeoutId: NodeJS.Timeout | undefined; // Prevents multiple parallel requests during DEBOUNCE_DELAY_MS let lastRequestTimestamp: number | undefined = undefined; -const requestDatabaseResourceTokens = async (): Promise => { +/** + * Request fabric token: + * - Mirrored key and AAD: Database Resource Tokens + * - Native: AAD token + * @returns + */ +const requestFabricToken = async (): Promise => { if (lastRequestTimestamp !== undefined && lastRequestTimestamp + DEBOUNCE_DELAY_MS > Date.now()) { return; } if (!userContext.fabricContext || !userContext.databaseAccount) { + // This should not happen + logConsoleError("Fabric context or database account is missing: cannot request tokens"); return; } lastRequestTimestamp = Date.now(); try { - const resourceTokenInfo = await sendCachedDataMessage( - FabricMessageTypes.GetAllResourceTokens, - [], - userContext.fabricContext.artifactInfo?.connectionId, - ); - - if (!userContext.databaseAccount.properties.documentEndpoint) { - userContext.databaseAccount.properties.documentEndpoint = resourceTokenInfo.endpoint; + if (isFabricMirrored()) { + await requestAndStoreDatabaseResourceTokens(); + } else if (isFabricNative()) { + await requestAndStoreAccessToken(); } + scheduleRefreshFabricToken(); + } catch (error) { + logConsoleError(error as string); + throw error; + } finally { + lastRequestTimestamp = undefined; + } +}; + +const requestAndStoreDatabaseResourceTokens = async (): Promise => { + const resourceTokenInfo = await sendCachedDataMessage( + FabricMessageTypes.GetAllResourceTokens, + [], + userContext.fabricContext.artifactInfo?.connectionId, + ); + + if (!userContext.databaseAccount.properties.documentEndpoint) { + userContext.databaseAccount.properties.documentEndpoint = resourceTokenInfo.endpoint; + } + + if (resourceTokenInfo.credentialType === "OAuth2") { + // Mirrored AAD + updateUserContext({ + fabricContext: { + ...userContext.fabricContext, + databaseName: resourceTokenInfo.databaseId, + artifactInfo: undefined, + isReadOnly: resourceTokenInfo.isReadOnly ?? userContext.fabricContext.isReadOnly, + }, + databaseAccount: { ...userContext.databaseAccount }, + aadToken: resourceTokenInfo.accessToken, + }); + } else { + // TODO: In Fabric contract V2, credentialType is undefined. For V3, it is "Key". Check for "Key" when V3 is supported for Fabric Mirroring Key + // Mirrored key updateUserContext({ fabricContext: { ...userContext.fabricContext, @@ -45,21 +84,27 @@ const requestDatabaseResourceTokens = async (): Promise => { }, databaseAccount: { ...userContext.databaseAccount }, }); - scheduleRefreshDatabaseResourceToken(); - } catch (error) { - logConsoleError(error as string); - throw error; - } finally { - lastRequestTimestamp = undefined; } }; +const requestAndStoreAccessToken = async (): Promise => { + const accessTokenInfo = await sendCachedDataMessage<{ accessToken: string }>( + FabricMessageTypes.GetAccessToken, + [], + userContext.fabricContext.artifactInfo?.connectionId, + ); + + updateUserContext({ + aadToken: accessTokenInfo.accessToken, + }); +}; + /** * Check token validity and schedule a refresh if necessary * @param tokenTimestamp * @returns */ -export const scheduleRefreshDatabaseResourceToken = (refreshNow?: boolean): Promise => { +export const scheduleRefreshFabricToken = (refreshNow?: boolean): Promise => { return new Promise((resolve) => { if (timeoutId !== undefined) { clearTimeout(timeoutId); @@ -68,7 +113,7 @@ export const scheduleRefreshDatabaseResourceToken = (refreshNow?: boolean): Prom timeoutId = setTimeout( () => { - requestDatabaseResourceTokens().then(resolve); + requestFabricToken().then(resolve); }, refreshNow ? 0 : TOKEN_VALIDITY_MS, ); @@ -77,7 +122,7 @@ export const scheduleRefreshDatabaseResourceToken = (refreshNow?: boolean): Prom export const checkDatabaseResourceTokensValidity = (tokenTimestamp: number): void => { if (tokenTimestamp + TOKEN_VALIDITY_MS < Date.now()) { - scheduleRefreshDatabaseResourceToken(true); + scheduleRefreshFabricToken(true); } }; diff --git a/src/hooks/useKnockoutExplorer.ts b/src/hooks/useKnockoutExplorer.ts index 9a1fc7b43..52e3894a8 100644 --- a/src/hooks/useKnockoutExplorer.ts +++ b/src/hooks/useKnockoutExplorer.ts @@ -14,7 +14,7 @@ import { useDialog } from "Explorer/Controls/Dialog"; import Explorer from "Explorer/Explorer"; import { useDataPlaneRbac } from "Explorer/Panes/SettingsPane/SettingsPane"; import { useSelectedNode } from "Explorer/useSelectedNode"; -import { scheduleRefreshDatabaseResourceToken } from "Platform/Fabric/FabricUtil"; +import { isFabricMirroredKey, scheduleRefreshFabricToken } from "Platform/Fabric/FabricUtil"; import { AppStateComponentNames, OPEN_TABS_SUBCOMPONENT_NAME, @@ -154,7 +154,7 @@ async function configureFabric(): Promise { }; explorer = createExplorerFabricLegacy(initializationMessage, data.version); - await scheduleRefreshDatabaseResourceToken(true); + await scheduleRefreshFabricToken(true); resolve(explorer); await explorer.refreshAllDatabases(); if (userContext.fabricContext.isVisible) { @@ -169,9 +169,12 @@ async function configureFabric(): Promise { if (initializationMessage.artifactType === CosmosDbArtifactType.MIRRORED_KEY) { // Do not show Home tab for Mirrored useTabs.getState().closeReactTab(ReactTabKind.Home); - await scheduleRefreshDatabaseResourceToken(true); } + // All tokens used in fabric expire + // For Mirrored key, we need the token right away to get the database and containers list. + await scheduleRefreshFabricToken(isFabricMirroredKey()); + resolve(explorer); await explorer.refreshAllDatabases(); @@ -188,7 +191,8 @@ async function configureFabric(): Promise { explorer.onNewCollectionClicked(); break; case "authorizationToken": - case "allResourceTokens_v2": { + case "allResourceTokens_v2": + case "accessToken": { handleCachedDataMessage(data); break; }