From caf4e9f51f5daa68d68fd2c8298ea043acd06993 Mon Sep 17 00:00:00 2001 From: Senthamil Sindhu Date: Wed, 3 Jul 2024 16:01:23 -0700 Subject: [PATCH] Add fetchAndUpdate Keys --- src/Common/CosmosClient.ts | 35 ++------------ .../Panes/SettingsPane/SettingsPane.tsx | 9 ++++ src/hooks/useKnockoutExplorer.ts | 46 +++++++++++++------ 3 files changed, 45 insertions(+), 45 deletions(-) diff --git a/src/Common/CosmosClient.ts b/src/Common/CosmosClient.ts index 682ab4e39..b5e50d5cb 100644 --- a/src/Common/CosmosClient.ts +++ b/src/Common/CosmosClient.ts @@ -3,12 +3,10 @@ import { getAuthorizationTokenUsingResourceTokens } from "Common/getAuthorizatio import { AuthorizationToken } from "Contracts/FabricMessageTypes"; import { checkDatabaseResourceTokensValidity } from "Platform/Fabric/FabricUtil"; import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility"; -import { listKeys } from "Utils/arm/generatedClients/cosmos/databaseAccounts"; -import { DatabaseAccountListKeysResult } from "Utils/arm/generatedClients/cosmos/types"; import { AuthType } from "../AuthType"; import { PriorityLevel } from "../Common/Constants"; import { Platform, configContext } from "../ConfigContext"; -import { updateUserContext, userContext } from "../UserContext"; +import { userContext } from "../UserContext"; import { logConsoleError } from "../Utils/NotificationConsoleUtils"; import * as PriorityBasedExecutionUtils from "../Utils/PriorityBasedExecutionUtils"; import { EmulatorMasterKey, HttpHeaders } from "./Constants"; @@ -81,12 +79,6 @@ export const tokenProvider = async (requestInfo: Cosmos.RequestInfo) => { } } - let retryAttempts: number = 50; - while (retryAttempts > 0 && userContext.listKeysFetchInProgress) { - retryAttempts--; - await sleep(100); - } - if (userContext.masterKey) { // TODO This SDK method mutates the headers object. Find a better one or fix the SDK. await Cosmos.setAuthorizationTokenHeaderUsingMasterKey( @@ -97,23 +89,7 @@ export const tokenProvider = async (requestInfo: Cosmos.RequestInfo) => { userContext.masterKey, ); return decodeURIComponent(headers.authorization); - } else if (userContext.dataPlaneRbacEnabled == false) { - const { databaseAccount: account, subscriptionId, resourceGroup } = userContext; - const keys: DatabaseAccountListKeysResult = await listKeys(subscriptionId, resourceGroup, account.name); - - if (keys.primaryMasterKey) { - updateUserContext({ masterKey: keys.primaryMasterKey }); - // TODO This SDK method mutates the headers object. Find a better one or fix the SDK. - await Cosmos.setAuthorizationTokenHeaderUsingMasterKey( - verb, - resourceId, - resourceType, - headers, - keys.primaryMasterKey, - ); - return decodeURIComponent(headers.authorization); - } - } + } if (userContext.resourceToken) { return userContext.resourceToken; @@ -124,10 +100,6 @@ export const tokenProvider = async (requestInfo: Cosmos.RequestInfo) => { return decodeURIComponent(result.PrimaryReadWriteToken); }; -function sleep(ms: number) { - return new Promise((resolve) => setTimeout(resolve, ms)); -} - export const requestPlugin: Cosmos.Plugin = async (requestContext, diagnosticNode, next) => { requestContext.endpoint = new URL(configContext.PROXY_PATH, window.location.href).href; requestContext.headers["x-ms-proxy-target"] = endpoint(); @@ -197,9 +169,10 @@ export function client(): Cosmos.CosmosClient { // The header will be ignored if priority based execution is disabled on the account. _defaultHeaders["x-ms-cosmos-priority-level"] = PriorityLevel.Default; } - + const options: Cosmos.CosmosClientOptions = { endpoint: endpoint() || "https://cosmos.azure.com", // CosmosClient gets upset if we pass a bad URL. This should never actually get called + key: userContext.dataPlaneRbacEnabled ? "" : userContext.masterKey, tokenProvider, userAgentSuffix: "Azure Portal", defaultHeaders: _defaultHeaders, diff --git a/src/Explorer/Panes/SettingsPane/SettingsPane.tsx b/src/Explorer/Panes/SettingsPane/SettingsPane.tsx index 09b4d457f..9e6d2afd6 100644 --- a/src/Explorer/Panes/SettingsPane/SettingsPane.tsx +++ b/src/Explorer/Panes/SettingsPane/SettingsPane.tsx @@ -36,6 +36,8 @@ import Explorer from "../../Explorer"; import { RightPaneForm, RightPaneFormProps } from "../RightPaneForm/RightPaneForm"; import { AuthType } from "AuthType"; import create, { UseStore } from "zustand"; +import { DatabaseAccountListKeysResult } from "@azure/arm-cosmosdb/esm/models"; +import { listKeys } from "Utils/arm/generatedClients/cosmos/databaseAccounts"; export interface DataPlaneRbacState { dataPlaneRbacEnabled: boolean; @@ -166,8 +168,15 @@ export const SettingsPane: FunctionComponent<{ explorer: Explorer }> = ({ updateUserContext({ dataPlaneRbacEnabled: false, }); + const { databaseAccount: account, subscriptionId, resourceGroup } = userContext; + const keys: DatabaseAccountListKeysResult = await listKeys(subscriptionId, resourceGroup, account.name); + + if (keys.primaryMasterKey) { + updateUserContext({ masterKey: keys.primaryMasterKey }); + useDataPlaneRbac.setState({ dataPlaneRbacEnabled: false }); } + } LocalStorageUtility.setEntryBoolean(StorageKey.RUThresholdEnabled, ruThresholdEnabled); LocalStorageUtility.setEntryBoolean(StorageKey.QueryTimeoutEnabled, queryTimeoutEnabled); diff --git a/src/hooks/useKnockoutExplorer.ts b/src/hooks/useKnockoutExplorer.ts index b98341b63..cea0f1428 100644 --- a/src/hooks/useKnockoutExplorer.ts +++ b/src/hooks/useKnockoutExplorer.ts @@ -417,11 +417,24 @@ function configureEmulator(): Explorer { return explorer; } + async function fetchAndUpdateKeys(subscriptionId: string, resourceGroup: string, account: string) { + try { + updateUserContext({ listKeysFetchInProgress: true }); + const keys = await listKeys(subscriptionId, resourceGroup, account); + + updateUserContext({ masterKey: keys.primaryMasterKey, listKeysFetchInProgress: false }); + } catch (error) { + updateUserContext({ listKeysFetchInProgress: false }); + console.error("Error during fetching keys or updating user context:", error); + } +} + async function configurePortal(): Promise { updateUserContext({ authType: AuthType.AAD, }); let explorer: Explorer; + return new Promise((resolve) => { // In development mode, try to load the iframe message from session storage. // This allows webpack hot reload to function properly in the portal @@ -477,7 +490,7 @@ async function configurePortal(): Promise { setTimeout(() => explorer.openNPSSurveyDialog(), 3000); } - const { databaseAccount: account, subscriptionId, resourceGroup } = userContext; + const { databaseAccount: account , subscriptionId, resourceGroup} = userContext; if (userContext.apiType === "SQL") { if (LocalStorageUtility.hasItem(StorageKey.DataPlaneRbacEnabled)) { @@ -489,16 +502,28 @@ async function configurePortal(): Promise { } else { dataPlaneRbacEnabled = isDataPlaneRbacSetting === Constants.RBACOptions.setTrueRBACOption; } + if(!dataPlaneRbacEnabled) { + (async () => { + await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name); + })(); + } updateUserContext({ dataPlaneRbacEnabled }); useDataPlaneRbac.setState({ dataPlaneRbacEnabled: dataPlaneRbacEnabled }); } else { const dataPlaneRbacEnabled = account.properties.disableLocalAuth; + if(!dataPlaneRbacEnabled) { + (async () => { + await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name); + })(); + } + updateUserContext({ dataPlaneRbacEnabled }); useDataPlaneRbac.setState({ dataPlaneRbacEnabled: dataPlaneRbacEnabled }); } - } else { + } + else { (async () => { await fetchAndUpdateKeys(subscriptionId, resourceGroup, account.name); })(); @@ -529,18 +554,6 @@ async function configurePortal(): Promise { }); } -async function fetchAndUpdateKeys(subscriptionId: string, resourceGroup: string, account: string) { - try { - updateUserContext({ listKeysFetchInProgress: true }); - const keys = await listKeys(subscriptionId, resourceGroup, account); - - updateUserContext({ masterKey: keys.primaryMasterKey, listKeysFetchInProgress: false }); - } catch (error) { - updateUserContext({ listKeysFetchInProgress: false }); - console.error("Error during fetching keys or updating user context:", error); - } -} - function shouldForwardMessage(message: PortalMessage, messageOrigin: string) { // Only allow forwarding messages from the same origin return messageOrigin === window.document.location.origin && message.type === MessageTypes.TelemetryInfo; @@ -558,6 +571,11 @@ function updateContextsFromPortalMessage(inputs: DataExplorerInputsFrame) { const authorizationToken = inputs.authorizationToken || ""; const databaseAccount = inputs.databaseAccount; + if(userContext.apiType !== "SQL") { + const masterKey = inputs.masterKey || ""; + updateUserContext({masterKey}); + } + updateConfigContext({ BACKEND_ENDPOINT: inputs.extensionEndpoint || configContext.BACKEND_ENDPOINT, ARM_ENDPOINT: normalizeArmEndpoint(inputs.csmEndpoint || configContext.ARM_ENDPOINT),