diff --git a/src/Common/Constants.ts b/src/Common/Constants.ts index 44d2b8b90..b3c15e423 100644 --- a/src/Common/Constants.ts +++ b/src/Common/Constants.ts @@ -138,6 +138,14 @@ export enum MongoBackendEndpointType { remote, } +export class AadScopeEndpoints { + public static readonly Development: string = "https://cosmos.azure.com"; + public static readonly MPAC: string = "https://cosmos.azure.com"; + public static readonly Prod: string = "https://cosmos.azure.com"; + public static readonly Fairfax: string = "https://cosmos.azure.us"; + public static readonly Mooncake: string = "https://cosmos.azure.cn"; +} + export class PortalBackendEndpoints { public static readonly Development: string = "https://localhost:7235"; public static readonly Mpac: string = "https://cdb-ms-mpac-pbe.cosmos.azure.com"; @@ -255,6 +263,7 @@ export class HttpHeaders { public static activityId: string = "x-ms-activity-id"; public static apiType: string = "x-ms-cosmos-apitype"; public static authorization: string = "authorization"; + public static entraIdToken: string = "x-ms-entraid-token"; public static collectionIndexTransformationProgress: string = "x-ms-documentdb-collection-index-transformation-progress"; public static continuation: string = "x-ms-continuation"; diff --git a/src/Common/EnvironmentUtility.test.ts b/src/Common/EnvironmentUtility.test.ts index 1c1d9d8da..79a963b1b 100644 --- a/src/Common/EnvironmentUtility.test.ts +++ b/src/Common/EnvironmentUtility.test.ts @@ -28,3 +28,39 @@ describe("Environment Utility Test", () => { expect(EnvironmentUtility.getEnvironment()).toBe(EnvironmentUtility.Environment.Development); }); }); +describe("normalizeArmEndpoint", () => { + it("should append '/' if not present", () => { + expect(EnvironmentUtility.normalizeArmEndpoint("https://example.com")).toBe("https://example.com/"); + }); + + it("should return the same uri if '/' is present at the end", () => { + expect(EnvironmentUtility.normalizeArmEndpoint("https://example.com/")).toBe("https://example.com/"); + }); + + it("should handle empty string", () => { + expect(EnvironmentUtility.normalizeArmEndpoint("")).toBe(""); + }); +}); + +describe("getEnvironment", () => { + it("should return Prod environment", () => { + updateConfigContext({ + PORTAL_BACKEND_ENDPOINT: PortalBackendEndpoints.Prod, + }); + expect(EnvironmentUtility.getEnvironment()).toBe(EnvironmentUtility.Environment.Prod); + }); + + it("should return Fairfax environment", () => { + updateConfigContext({ + PORTAL_BACKEND_ENDPOINT: PortalBackendEndpoints.Fairfax, + }); + expect(EnvironmentUtility.getEnvironment()).toBe(EnvironmentUtility.Environment.Fairfax); + }); + + it("should return Mooncake environment", () => { + updateConfigContext({ + PORTAL_BACKEND_ENDPOINT: PortalBackendEndpoints.Mooncake, + }); + expect(EnvironmentUtility.getEnvironment()).toBe(EnvironmentUtility.Environment.Mooncake); + }); +}); diff --git a/src/Common/EnvironmentUtility.ts b/src/Common/EnvironmentUtility.ts index cafcbfb3e..6331380a9 100644 --- a/src/Common/EnvironmentUtility.ts +++ b/src/Common/EnvironmentUtility.ts @@ -1,4 +1,5 @@ -import { PortalBackendEndpoints } from "Common/Constants"; +import { AadScopeEndpoints, PortalBackendEndpoints } from "Common/Constants"; +import * as Logger from "Common/Logger"; import { configContext } from "ConfigContext"; export function normalizeArmEndpoint(uri: string): string { @@ -27,3 +28,17 @@ export const getEnvironment = (): Environment => { return environmentMap[configContext.PORTAL_BACKEND_ENDPOINT]; }; + +export const getEnvironmentScopeEndpoint = (): string => { + const environment = getEnvironment(); + const endpoint = AadScopeEndpoints[environment]; + if (!endpoint) { + throw new Error("Cannot determine AAD scope endpoint"); + } + const hrefEndpoint = new URL(endpoint).href.replace(/\/+$/, "/.default"); + Logger.logInfo( + `Using AAD scope endpoint: ${hrefEndpoint}, Environment: ${environment}`, + "EnvironmentUtility/getEnvironmentScopeEndpoint", + ); + return hrefEndpoint; +}; diff --git a/src/Common/MongoProxyClient.ts b/src/Common/MongoProxyClient.ts index a07572bd3..298b9268d 100644 --- a/src/Common/MongoProxyClient.ts +++ b/src/Common/MongoProxyClient.ts @@ -7,6 +7,7 @@ import { MessageTypes } from "../Contracts/ExplorerContracts"; import { Collection } from "../Contracts/ViewModels"; import DocumentId from "../Explorer/Tree/DocumentId"; import { userContext } from "../UserContext"; +import { isDataplaneRbacEnabledForProxyApi } from "../Utils/AuthorizationUtils"; import { logConsoleError } from "../Utils/NotificationConsoleUtils"; import { ApiType, ContentType, HttpHeaders, HttpStatusCodes } from "./Constants"; import { MinimalQueryIterator } from "./IteratorUtilities"; @@ -22,7 +23,13 @@ function authHeaders() { if (userContext.authType === AuthType.EncryptedToken) { return { [HttpHeaders.guestAccessToken]: userContext.accessToken }; } else { - return { [HttpHeaders.authorization]: userContext.authorizationToken }; + const headers: { [key: string]: string } = { + [HttpHeaders.authorization]: userContext.authorizationToken, + }; + if (isDataplaneRbacEnabledForProxyApi(userContext)) { + headers[HttpHeaders.entraIdToken] = userContext.aadToken; + } + return headers; } } diff --git a/src/Explorer/Menus/CommandBar/CommandBarComponentAdapter.tsx b/src/Explorer/Menus/CommandBar/CommandBarComponentAdapter.tsx index eb596c0f7..b65310691 100644 --- a/src/Explorer/Menus/CommandBar/CommandBarComponentAdapter.tsx +++ b/src/Explorer/Menus/CommandBar/CommandBarComponentAdapter.tsx @@ -5,6 +5,7 @@ */ import { CommandBar as FluentCommandBar, ICommandBarItemProps } from "@fluentui/react"; import { useNotebook } from "Explorer/Notebook/useNotebook"; +import { useDataPlaneRbac } from "Explorer/Panes/SettingsPane/SettingsPane"; import { KeyboardActionGroup, useKeyboardActionGroup } from "KeyboardShortcuts"; import { isFabric } from "Platform/Fabric/FabricUtil"; import { userContext } from "UserContext"; @@ -30,7 +31,7 @@ export interface CommandBarStore { } export const useCommandBar: UseStore = create((set) => ({ - contextButtons: [], + contextButtons: [] as CommandButtonComponentProps[], setContextButtons: (contextButtons: CommandButtonComponentProps[]) => set((state) => ({ ...state, contextButtons })), isHidden: false, setIsHidden: (isHidden: boolean) => set((state) => ({ ...state, isHidden })), @@ -43,6 +44,15 @@ export const CommandBar: React.FC = ({ container }: Props) => { const backgroundColor = StyleConstants.BaseLight; const setKeyboardHandlers = useKeyboardActionGroup(KeyboardActionGroup.COMMAND_BAR); + // Subscribe to the store changes that affect button creation + const dataPlaneRbacEnabled = useDataPlaneRbac((state) => state.dataPlaneRbacEnabled); + const aadTokenUpdated = useDataPlaneRbac((state) => state.aadTokenUpdated); + + // Memoize the expensive button creation + const staticButtons = React.useMemo(() => { + return CommandBarComponentButtonFactory.createStaticCommandBarButtons(container, selectedNodeState); + }, [container, selectedNodeState, dataPlaneRbacEnabled, aadTokenUpdated]); + if (userContext.apiType === "Postgres" || userContext.apiType === "VCoreMongo") { const buttons = userContext.apiType === "Postgres" @@ -62,7 +72,6 @@ export const CommandBar: React.FC = ({ container }: Props) => { ); } - const staticButtons = CommandBarComponentButtonFactory.createStaticCommandBarButtons(container, selectedNodeState); const contextButtons = (buttons || []).concat( CommandBarComponentButtonFactory.createContextCommandBarButtons(container, selectedNodeState), ); diff --git a/src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx b/src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx index fbf1b4434..6df096b33 100644 --- a/src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx +++ b/src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx @@ -1,7 +1,6 @@ import { KeyboardAction } from "KeyboardShortcuts"; import { isDataplaneRbacSupported } from "Utils/APITypeUtils"; import * as React from "react"; -import { useEffect, useState } from "react"; import AddSqlQueryIcon from "../../../../images/AddSqlQuery_16x16.svg"; import AddStoredProcedureIcon from "../../../../images/AddStoredProcedure.svg"; import AddTriggerIcon from "../../../../images/AddTrigger.svg"; @@ -68,15 +67,7 @@ export function createStaticCommandBarButtons( } if (isDataplaneRbacSupported(userContext.apiType)) { - const [loginButtonProps, setLoginButtonProps] = useState(undefined); - const dataPlaneRbacEnabled = useDataPlaneRbac((state) => state.dataPlaneRbacEnabled); - const aadTokenUpdated = useDataPlaneRbac((state) => state.aadTokenUpdated); - - useEffect(() => { - const buttonProps = createLoginForEntraIDButton(container); - setLoginButtonProps(buttonProps); - }, [dataPlaneRbacEnabled, aadTokenUpdated, container]); - + const loginButtonProps = createLoginForEntraIDButton(container); if (loginButtonProps) { addDivider(); buttons.push(loginButtonProps); diff --git a/src/Explorer/Tables/TableDataClient.ts b/src/Explorer/Tables/TableDataClient.ts index a842e9986..ff2d9a5ea 100644 --- a/src/Explorer/Tables/TableDataClient.ts +++ b/src/Explorer/Tables/TableDataClient.ts @@ -13,7 +13,7 @@ import { updateDocument } from "../../Common/dataAccess/updateDocument"; import { configContext } from "../../ConfigContext"; import * as ViewModels from "../../Contracts/ViewModels"; import { userContext } from "../../UserContext"; -import { getAuthorizationHeader } from "../../Utils/AuthorizationUtils"; +import { getAuthorizationHeader, isDataplaneRbacEnabledForProxyApi } from "../../Utils/AuthorizationUtils"; import * as NotificationConsoleUtils from "../../Utils/NotificationConsoleUtils"; import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils"; import Explorer from "../Explorer"; @@ -551,6 +551,10 @@ export class CassandraAPIDataClient extends TableDataClient { const authorizationHeaderMetadata: ViewModels.AuthorizationTokenHeaderMetadata = getAuthorizationHeader(); xhr.setRequestHeader(authorizationHeaderMetadata.header, authorizationHeaderMetadata.token); + if (isDataplaneRbacEnabledForProxyApi(userContext)) { + xhr.setRequestHeader(Constants.HttpHeaders.entraIdToken, userContext.aadToken); + } + return true; }; diff --git a/src/Explorer/Tabs/CloudShellTab/ShellTypes/AbstractShellHandler.tsx b/src/Explorer/Tabs/CloudShellTab/ShellTypes/AbstractShellHandler.tsx index b28e0bd12..783dcd99b 100644 --- a/src/Explorer/Tabs/CloudShellTab/ShellTypes/AbstractShellHandler.tsx +++ b/src/Explorer/Tabs/CloudShellTab/ShellTypes/AbstractShellHandler.tsx @@ -24,7 +24,7 @@ export const EXIT_COMMAND_MONGO = ` printf "\\033[1;31mSession ended. Please clo * This command runs mongosh in no-database and quiet mode, * and evaluates the `disableTelemetry()` function to turn off telemetry collection. */ -export const DISABLE_TELEMETRY_COMMAND = `mongosh --nodb --quiet --eval "disableTelemetry()"`; +export const DISABLE_TELEMETRY_COMMAND = `mongosh --nodb --quiet --eval 'disableTelemetry()'`; /** * Abstract class that defines the interface for shell-specific handlers @@ -97,7 +97,7 @@ export abstract class AbstractShellHandler { * is not already present in the environment. */ protected mongoShellSetupCommands(): string[] { - const PACKAGE_VERSION: string = "2.5.5"; + const PACKAGE_VERSION: string = "2.5.6"; return [ "if ! command -v mongosh &> /dev/null; then echo '⚠️ mongosh not found. Installing...'; fi", `if ! command -v mongosh &> /dev/null; then curl -LO https://downloads.mongodb.com/compass/mongosh-${PACKAGE_VERSION}-linux-x64.tgz; fi`, diff --git a/src/Explorer/Tabs/CloudShellTab/ShellTypes/MongoShellHandler.test.tsx b/src/Explorer/Tabs/CloudShellTab/ShellTypes/MongoShellHandler.test.tsx index 231933397..28a79d404 100644 --- a/src/Explorer/Tabs/CloudShellTab/ShellTypes/MongoShellHandler.test.tsx +++ b/src/Explorer/Tabs/CloudShellTab/ShellTypes/MongoShellHandler.test.tsx @@ -18,6 +18,12 @@ interface DatabaseAccount { interface UserContextType { databaseAccount: DatabaseAccount; + features: { + enableAadDataPlane: boolean; + }; + apiType: string; + dataPlaneRbacEnabled: boolean; + aadToken?: string; } // Mock dependencies @@ -29,6 +35,8 @@ jest.mock("../../../../UserContext", () => ({ mongoEndpoint: "https://test-mongo.documents.azure.com:443/", }, }, + features: { enableAadDataPlane: false }, + apiType: "Mongo", }, })); @@ -70,7 +78,7 @@ describe("MongoShellHandler", () => { expect(Array.isArray(commands)).toBe(true); expect(commands.length).toBe(7); - expect(commands[1]).toContain("mongosh-2.5.5-linux-x64.tgz"); + expect(commands[1]).toContain("mongosh-2.5.6-linux-x64.tgz"); }); }); @@ -88,11 +96,12 @@ describe("MongoShellHandler", () => { kind: "test-kind", properties: { mongoEndpoint: "https://test-mongo.documents.azure.com:443/" }, }; + (userContext as UserContextType).dataPlaneRbacEnabled = false; const command = mongoShellHandler.getConnectionCommand(); expect(command).toBe( - 'mongosh --nodb --quiet --eval "disableTelemetry()" && mongosh mongodb://test-mongo.documents.azure.com:10255?appName=CosmosExplorerTerminal --username test-account --password test-key --tls --tlsAllowInvalidCertificates', + "mongosh --nodb --quiet --eval 'disableTelemetry()'; mongosh mongodb://test-mongo.documents.azure.com:10255?appName=CosmosExplorerTerminal --username test-account --password test-key --tls --tlsAllowInvalidCertificates", ); expect(CommonUtils.getHostFromUrl).toHaveBeenCalledWith("https://test-mongo.documents.azure.com:443/"); @@ -115,12 +124,47 @@ describe("MongoShellHandler", () => { }; const command = mongoShellHandler.getConnectionCommand(); - expect(command).toBe("echo 'Database name not found.'"); // Restore original (userContext as UserContextType).databaseAccount = originalDatabaseAccount; }); + + it("should return echo if endpoint is missing", () => { + const testKey = "test-key"; + (userContext as UserContextType).databaseAccount = { + id: "test-id", + name: "", // Empty name to simulate missing name + location: "test-location", + type: "test-type", + kind: "test-kind", + properties: { mongoEndpoint: "" }, + }; + const mongoShellHandler = new MongoShellHandler(testKey); + const command = mongoShellHandler.getConnectionCommand(); + expect(command).toBe("echo 'MongoDB endpoint not found.'"); + }); + + it("should use _getAadConnectionCommand when _isEntraIdEnabled is true", () => { + const testKey = "aad-key"; + (userContext as UserContextType).databaseAccount = { + id: "test-id", + name: "test-account", + location: "test-location", + type: "test-type", + kind: "test-kind", + properties: { mongoEndpoint: "https://test-mongo.documents.azure.com:443/" }, + }; + (userContext as UserContextType).dataPlaneRbacEnabled = true; + + const mongoShellHandler = new MongoShellHandler(testKey); + + const command = mongoShellHandler.getConnectionCommand(); + expect(command).toContain( + "mongosh 'mongodb://test-account:aad-key@test-account.mongo.cosmos.azure.com:10255/?ssl=true&replicaSet=globaldb&authMechanism=PLAIN&retryWrites=false' --tls --tlsAllowInvalidCertificates", + ); + expect(command.startsWith("mongosh --nodb")).toBeTruthy(); + }); }); describe("getTerminalSuppressedData", () => { diff --git a/src/Explorer/Tabs/CloudShellTab/ShellTypes/MongoShellHandler.tsx b/src/Explorer/Tabs/CloudShellTab/ShellTypes/MongoShellHandler.tsx index ab2fb2e7f..b41dbd6db 100644 --- a/src/Explorer/Tabs/CloudShellTab/ShellTypes/MongoShellHandler.tsx +++ b/src/Explorer/Tabs/CloudShellTab/ShellTypes/MongoShellHandler.tsx @@ -1,4 +1,5 @@ import { userContext } from "../../../../UserContext"; +import { isDataplaneRbacEnabledForProxyApi } from "../../../../Utils/AuthorizationUtils"; import { filterAndCleanTerminalOutput, getHostFromUrl, getMongoShellRemoveInfoText } from "../Utils/CommonUtils"; import { AbstractShellHandler, DISABLE_TELEMETRY_COMMAND, EXIT_COMMAND_MONGO } from "./AbstractShellHandler"; @@ -6,12 +7,23 @@ export class MongoShellHandler extends AbstractShellHandler { private _key: string; private _endpoint: string | undefined; private _removeInfoText: string[] = getMongoShellRemoveInfoText(); + private _isEntraIdEnabled: boolean = isDataplaneRbacEnabledForProxyApi(userContext); constructor(private key: string) { super(); this._key = key; this._endpoint = userContext?.databaseAccount?.properties?.mongoEndpoint; } + private _getKeyConnectionCommand(dbName: string): string { + return `mongosh mongodb://${getHostFromUrl(this._endpoint)}:10255?appName=${ + this.APP_NAME + } --username ${dbName} --password ${this._key} --tls --tlsAllowInvalidCertificates`; + } + + private _getAadConnectionCommand(dbName: string): string { + return `mongosh 'mongodb://${dbName}:${this._key}@${dbName}.mongo.cosmos.azure.com:10255/?ssl=true&replicaSet=globaldb&authMechanism=PLAIN&retryWrites=false' --tls --tlsAllowInvalidCertificates`; + } + public getShellName(): string { return "MongoDB"; } @@ -29,19 +41,11 @@ export class MongoShellHandler extends AbstractShellHandler { if (!dbName) { return "echo 'Database name not found.'"; } - return ( - DISABLE_TELEMETRY_COMMAND + - " && " + - "mongosh mongodb://" + - getHostFromUrl(this._endpoint) + - ":10255?appName=" + - this.APP_NAME + - " --username " + - dbName + - " --password " + - this._key + - " --tls --tlsAllowInvalidCertificates" - ); + const connectionCommand = this._isEntraIdEnabled + ? this._getAadConnectionCommand(dbName) + : this._getKeyConnectionCommand(dbName); + const fullCommand = `${DISABLE_TELEMETRY_COMMAND}; ${connectionCommand}`; + return fullCommand; } public getTerminalSuppressedData(): string[] { diff --git a/src/Explorer/Tabs/CloudShellTab/ShellTypes/ShellTypeFactory.test.tsx b/src/Explorer/Tabs/CloudShellTab/ShellTypes/ShellTypeFactory.test.tsx index 327d40b3e..807844514 100644 --- a/src/Explorer/Tabs/CloudShellTab/ShellTypes/ShellTypeFactory.test.tsx +++ b/src/Explorer/Tabs/CloudShellTab/ShellTypes/ShellTypeFactory.test.tsx @@ -7,12 +7,24 @@ import { PostgresShellHandler } from "./PostgresShellHandler"; import { getHandler, getKey } from "./ShellTypeFactory"; import { VCoreMongoShellHandler } from "./VCoreMongoShellHandler"; +interface UserContextType { + databaseAccount: { name: string }; + subscriptionId: string; + resourceGroup: string; + features: { enableAadDataPlane: boolean }; + dataPlaneRbacEnabled: boolean; + aadToken?: string; + apiType?: string; +} + // Mock dependencies jest.mock("../../../../UserContext", () => ({ userContext: { databaseAccount: { name: "testDbName" }, subscriptionId: "testSubId", resourceGroup: "testResourceGroup", + features: { enableAadDataPlane: false }, + dataPlaneRbacEnabled: false, }, })); @@ -109,5 +121,33 @@ describe("ShellTypeHandlerFactory", () => { expect(key).toBe(mockKey); expect(listKeys).toHaveBeenCalledWith("testSubId", "testResourceGroup", "testDbName"); }); + + it("should return MongoShellHandler with primaryMasterKey for TerminalKind.Mongo when RBAC is disabled", async () => { + (listKeys as jest.Mock).mockResolvedValue({ primaryMasterKey: "primaryKey123" }); + (userContext as UserContextType).features.enableAadDataPlane = false; + (userContext as UserContextType).dataPlaneRbacEnabled = false; + const handler = await getHandler(TerminalKind.Mongo); + expect(handler).toBeInstanceOf(MongoShellHandler); + // eslint-disable-next-line @typescript-eslint/ban-ts-comment + // @ts-ignore + expect(handler.key).toBe("primaryKey123"); + }); + + it("should return MongoShellHandler with aadToken for TerminalKind.Mongo when RBAC is enabled", async () => { + (userContext as UserContextType).aadToken = "aadToken123"; + (userContext as UserContextType).features.enableAadDataPlane = true; + (userContext as UserContextType).dataPlaneRbacEnabled = true; + (userContext as UserContextType).apiType = "Mongo"; + const handler = await getHandler(TerminalKind.Mongo); + expect(handler).toBeInstanceOf(MongoShellHandler); + // eslint-disable-next-line @typescript-eslint/ban-ts-comment + // @ts-ignore + expect(handler.key).toBe("aadToken123"); + }); + it("should throw error for unsupported shell type", async () => { + await expect(getHandler("UnknownShell" as unknown as TerminalKind)).rejects.toThrow( + "Unsupported shell type: UnknownShell", + ); + }); }); }); diff --git a/src/Explorer/Tabs/CloudShellTab/ShellTypes/ShellTypeFactory.tsx b/src/Explorer/Tabs/CloudShellTab/ShellTypes/ShellTypeFactory.tsx index 30ecdcaf3..81e068562 100644 --- a/src/Explorer/Tabs/CloudShellTab/ShellTypes/ShellTypeFactory.tsx +++ b/src/Explorer/Tabs/CloudShellTab/ShellTypes/ShellTypeFactory.tsx @@ -1,6 +1,7 @@ import { TerminalKind } from "../../../../Contracts/ViewModels"; import { userContext } from "../../../../UserContext"; import { listKeys } from "../../../../Utils/arm/generatedClients/cosmos/databaseAccounts"; +import { isDataplaneRbacEnabledForProxyApi } from "../../../../Utils/AuthorizationUtils"; import { AbstractShellHandler } from "./AbstractShellHandler"; import { CassandraShellHandler } from "./CassandraShellHandler"; import { MongoShellHandler } from "./MongoShellHandler"; @@ -30,6 +31,9 @@ export async function getKey(): Promise { if (!dbName) { return ""; } + if (isDataplaneRbacEnabledForProxyApi(userContext)) { + return userContext.aadToken || ""; + } const keys = await listKeys(userContext.subscriptionId, userContext.resourceGroup, dbName); return keys?.primaryMasterKey || ""; diff --git a/src/Explorer/Tabs/CloudShellTab/ShellTypes/VCoreMongoShellHandler.test.tsx b/src/Explorer/Tabs/CloudShellTab/ShellTypes/VCoreMongoShellHandler.test.tsx index 5c34a4978..0cf35f1c4 100644 --- a/src/Explorer/Tabs/CloudShellTab/ShellTypes/VCoreMongoShellHandler.test.tsx +++ b/src/Explorer/Tabs/CloudShellTab/ShellTypes/VCoreMongoShellHandler.test.tsx @@ -45,7 +45,7 @@ describe("VCoreMongoShellHandler", () => { expect(Array.isArray(commands)).toBe(true); expect(commands.length).toBe(7); - expect(commands[1]).toContain("mongosh-2.5.5-linux-x64.tgz"); + expect(commands[1]).toContain("mongosh-2.5.6-linux-x64.tgz"); expect(commands[0]).toContain("mongosh not found"); }); diff --git a/src/Explorer/Tabs/CloudShellTab/Utils/AttachAddOn.tsx b/src/Explorer/Tabs/CloudShellTab/Utils/AttachAddOn.tsx index 0a388b28a..f2a2fbe78 100644 --- a/src/Explorer/Tabs/CloudShellTab/Utils/AttachAddOn.tsx +++ b/src/Explorer/Tabs/CloudShellTab/Utils/AttachAddOn.tsx @@ -92,6 +92,18 @@ export class AttachAddon implements ITerminalAddon { * @param {Terminal} terminal - The XTerm terminal instance */ public addMessageListener(terminal: Terminal): void { + let messageBuffer = ""; + let bufferTimeout: NodeJS.Timeout | null = null; + const BUFFER_TIMEOUT = 50; // ms - short timeout for prompt detection + + const processBuffer = () => { + if (messageBuffer.length > 0) { + this.handleCompleteTerminalData(terminal, messageBuffer); + messageBuffer = ""; + } + bufferTimeout = null; + }; + this._disposables.push( addSocketListener(this._socket, "message", (ev) => { let data: ArrayBuffer | string = ev.data; @@ -103,57 +115,136 @@ export class AttachAddon implements ITerminalAddon { data = enc.decode(ev.data as ArrayBuffer); } - // for example of json object look in TerminalHelper in the socket.onMessage - if (data.includes(startStatusJson) && data.includes(endStatusJson)) { - // process as one line - const statusData = data.split(startStatusJson)[1].split(endStatusJson)[0]; - data = data.replace(statusData, ""); - data = data.replace(startStatusJson, ""); - data = data.replace(endStatusJson, ""); - } else if (data.includes(startStatusJson)) { - // check for start - const partialStatusData = data.split(startStatusJson)[1]; - this._socketData += partialStatusData; - data = data.replace(partialStatusData, ""); - data = data.replace(startStatusJson, ""); - } else if (data.includes(endStatusJson)) { - // check for end and process the command - const partialStatusData = data.split(endStatusJson)[0]; - this._socketData += partialStatusData; - data = data.replace(partialStatusData, ""); - data = data.replace(endStatusJson, ""); - this._socketData = ""; - } else if (this._socketData.length > 0) { - // check if the line is all data then just concatenate - this._socketData += data; - data = ""; - } + // Handle status messages + let processedStatusData = data; - if (this._allowTerminalWrite && data.includes(this._startMarker)) { - this._allowTerminalWrite = false; - terminal.write(`Preparing ${this._shellHandler.getShellName()} environment...\r\n`); - } - - if (this._allowTerminalWrite) { - const updatedData = - typeof this._shellHandler?.updateTerminalData === "function" - ? this._shellHandler.updateTerminalData(data) - : data; - - const suppressedData = this._shellHandler?.getTerminalSuppressedData(); - - const shouldNotWrite = suppressedData.filter(Boolean).some((item) => updatedData.includes(item)); - - if (!shouldNotWrite) { - terminal.write(updatedData); + // Process status messages with delimiters + // eslint-disable-next-line no-constant-condition + while (true) { + const startIndex = processedStatusData.indexOf(startStatusJson); + if (startIndex === -1) { + break; } + + const afterStart = processedStatusData.substring(startIndex + startStatusJson.length); + const endIndex = afterStart.indexOf(endStatusJson); + + if (endIndex === -1) { + // Incomplete status message + this._socketData += processedStatusData.substring(startIndex); + processedStatusData = processedStatusData.substring(0, startIndex); + break; + } + + // Remove processed status message + processedStatusData = + processedStatusData.substring(0, startIndex) + afterStart.substring(endIndex + endStatusJson.length); } - if (data.includes(this._shellHandler.getConnectionCommand())) { - this._allowTerminalWrite = true; + // Add to message buffer + messageBuffer += processedStatusData; + + // Clear existing timeout + if (bufferTimeout) { + clearTimeout(bufferTimeout); + bufferTimeout = null; + } + + // Check if this looks like a complete message/command + const isComplete = this.isMessageComplete(messageBuffer, processedStatusData); + + if (isComplete) { + // Message marked as complete, processing immediately + processBuffer(); + } else { + // Set timeout to process buffer after delay + bufferTimeout = setTimeout(processBuffer, BUFFER_TIMEOUT); } }), ); + + // Clean up timeout on dispose + this._disposables.push({ + dispose: () => { + if (bufferTimeout) { + clearTimeout(bufferTimeout); + } + }, + }); + } + + private isMessageComplete(fullBuffer: string, currentChunk: string): boolean { + // Immediate completion indicators + const immediateCompletionPatterns = [ + /\n$/, // Ends with newline + /\r$/, // Ends with carriage return + /\r\n$/, // Ends with CRLF + /; \} \|\| true;$/, // Your command pattern + /disown -a && exit$/, // Exit commands + /printf.*?\\033\[0m\\n"$/, // Your printf pattern + ]; + + // Check current chunk for immediate completion + for (const pattern of immediateCompletionPatterns) { + if (pattern.test(currentChunk)) { + return true; + } + } + + // ANSI sequence detection - these might be complete prompts + const ansiPromptPatterns = [ + /\[\d+G\[0J.*>\s*\[\d+G$/, // Your specific pattern: [1G[0J...> [26G + /\[\d+;\d+H/, // Cursor position sequences + /\]\s*\[\d+G$/, // Ends with cursor positioning + />\s*\[\d+G$/, // Prompt followed by cursor position + ]; + + // Check if buffer ends with what looks like a complete prompt + for (const pattern of ansiPromptPatterns) { + if (pattern.test(fullBuffer)) { + return true; + } + } + + // Check for MongoDB shell prompts specifically + const mongoPromptPatterns = [ + /globaldb \[primary\] \w+>\s*\[\d+G$/, // MongoDB replica set prompt + />\s*\[\d+G$/, // General prompt with cursor positioning + /\w+>\s*$/, // Simple shell prompt + ]; + + for (const pattern of mongoPromptPatterns) { + if (pattern.test(fullBuffer)) { + return true; + } + } + + return false; + } + + private handleCompleteTerminalData(terminal: Terminal, data: string): void { + if (this._allowTerminalWrite && data.includes(this._startMarker)) { + this._allowTerminalWrite = false; + terminal.write(`Preparing ${this._shellHandler.getShellName()} environment...\r\n`); + } + + if (this._allowTerminalWrite) { + const updatedData = + typeof this._shellHandler?.updateTerminalData === "function" + ? this._shellHandler.updateTerminalData(data) + : data; + + const suppressedData = this._shellHandler?.getTerminalSuppressedData(); + const shouldNotWrite = suppressedData.filter(Boolean).some((item) => updatedData.includes(item)); + + if (!shouldNotWrite) { + terminal.write(updatedData); + } + } + + if (data.includes(this._shellHandler.getConnectionCommand())) { + this._allowTerminalWrite = true; + } } public dispose(): void { diff --git a/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2Mongo.test.tsx b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2Mongo.test.tsx index 5628e2c7d..7cee70afe 100644 --- a/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2Mongo.test.tsx +++ b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2Mongo.test.tsx @@ -146,10 +146,16 @@ describe("Documents tab (Mongo API)", () => { updateConfigContext({ platform: Platform.Hosted }); const props: IDocumentsTabComponentProps = createMockProps(); - wrapper = mount(); - wrapper = await waitForComponentToPaint(wrapper); - }); + + // Wait for all pending promises + await act(async () => { + await new Promise((resolve) => setTimeout(resolve, 100)); + }); + + // Wait for any async operations to complete + wrapper = await waitForComponentToPaint(wrapper, 100); + }, 10000); afterEach(() => { wrapper.unmount(); diff --git a/src/Utils/APITypeUtils.ts b/src/Utils/APITypeUtils.ts index dae50eb00..d329ae73c 100644 --- a/src/Utils/APITypeUtils.ts +++ b/src/Utils/APITypeUtils.ts @@ -91,5 +91,11 @@ export const getItemName = (): string => { }; export const isDataplaneRbacSupported = (apiType: string): boolean => { - return apiType === "SQL" || apiType === "Tables" || apiType === "Gremlin"; + return ( + apiType === "SQL" || apiType === "Tables" || apiType === "Gremlin" || apiType === "Mongo" || apiType === "Cassandra" + ); +}; + +export const hasProxyServer = (apiType: string): boolean => { + return apiType === "Mongo" || apiType === "Cassandra"; }; diff --git a/src/Utils/AuthorizationUtils.test.ts b/src/Utils/AuthorizationUtils.test.ts index d9fd43031..bc47c7a5f 100644 --- a/src/Utils/AuthorizationUtils.test.ts +++ b/src/Utils/AuthorizationUtils.test.ts @@ -104,7 +104,7 @@ describe("AuthorizationUtils", () => { it("should return true if dataPlaneRbacEnabled is set to true and API supports RBAC", () => { setAadDataPlane(false); - ["SQL", "Tables", "Gremlin"].forEach((type) => { + ["SQL", "Tables", "Gremlin", "Mongo", "Cassandra"].forEach((type) => { updateUserContext({ dataPlaneRbacEnabled: true, apiType: type as ApiType, @@ -115,7 +115,7 @@ describe("AuthorizationUtils", () => { it("should return false if dataPlaneRbacEnabled is set to true and API does not support RBAC", () => { setAadDataPlane(false); - ["Mongo", "Cassandra", "Postgres", "VCoreMongo"].forEach((type) => { + ["Postgres", "VCoreMongo"].forEach((type) => { updateUserContext({ dataPlaneRbacEnabled: true, apiType: type as ApiType, diff --git a/src/Utils/AuthorizationUtils.ts b/src/Utils/AuthorizationUtils.ts index 359d2ba20..84d4cfcc4 100644 --- a/src/Utils/AuthorizationUtils.ts +++ b/src/Utils/AuthorizationUtils.ts @@ -1,6 +1,7 @@ import * as msal from "@azure/msal-browser"; +import { getEnvironmentScopeEndpoint } from "Common/EnvironmentUtility"; import { Action, ActionModifiers } from "Shared/Telemetry/TelemetryConstants"; -import { isDataplaneRbacSupported } from "Utils/APITypeUtils"; +import { hasProxyServer, isDataplaneRbacSupported } from "Utils/APITypeUtils"; import { AuthType } from "../AuthType"; import * as Constants from "../Common/Constants"; import * as Logger from "../Common/Logger"; @@ -74,10 +75,12 @@ export async function acquireMsalTokenForAccount( if (userContext.databaseAccount.properties?.documentEndpoint === undefined) { throw new Error("Database account has no document endpoint defined"); } - const hrefEndpoint = new URL(userContext.databaseAccount.properties.documentEndpoint).href.replace( - /\/+$/, - "/.default", - ); + let hrefEndpoint = ""; + if (isDataplaneRbacEnabledForProxyApi(userContext)) { + hrefEndpoint = getEnvironmentScopeEndpoint(); + } else { + hrefEndpoint = new URL(userContext.databaseAccount.properties.documentEndpoint).href.replace(/\/+$/, "/.default"); + } const msalInstance = await getMsalInstance(); const knownAccounts = msalInstance.getAllAccounts(); // If user_hint is provided, we will try to use it to find the account. @@ -183,7 +186,11 @@ export async function acquireTokenWithMsal( export function useDataplaneRbacAuthorization(userContext: UserContext): boolean { return ( - userContext.features.enableAadDataPlane || + userContext.features?.enableAadDataPlane || (userContext.dataPlaneRbacEnabled && isDataplaneRbacSupported(userContext.apiType)) ); } + +export function isDataplaneRbacEnabledForProxyApi(userContext: UserContext): boolean { + return useDataplaneRbacAuthorization(userContext) && hasProxyServer(userContext.apiType); +} diff --git a/src/hooks/useKnockoutExplorer.ts b/src/hooks/useKnockoutExplorer.ts index 79d9c635f..46e655a87 100644 --- a/src/hooks/useKnockoutExplorer.ts +++ b/src/hooks/useKnockoutExplorer.ts @@ -1,4 +1,5 @@ import * as Constants from "Common/Constants"; +import { getEnvironmentScopeEndpoint } from "Common/EnvironmentUtility"; import { createUri } from "Common/UrlUtility"; import { DATA_EXPLORER_RPC_VERSION } from "Contracts/DataExplorerMessagesContract"; import { FabricMessageTypes } from "Contracts/FabricMessageTypes"; @@ -62,6 +63,7 @@ import { acquireTokenWithMsal, getAuthorizationHeader, getMsalInstance, + isDataplaneRbacEnabledForProxyApi, } from "../Utils/AuthorizationUtils"; import { isInvalidParentFrameOrigin, shouldProcessMessage } from "../Utils/MessageValidation"; import { get, getReadOnlyKeys, listKeys } from "../Utils/arm/generatedClients/cosmos/databaseAccounts"; @@ -331,7 +333,12 @@ async function configureHostedWithAAD(config: AAD): Promise { const resourceGroup = accountResourceId && accountResourceId.split("resourceGroups/")[1].split("/")[0]; let aadToken; if (account.properties?.documentEndpoint) { - const hrefEndpoint = new URL(account.properties.documentEndpoint).href.replace(/\/$/, "/.default"); + let hrefEndpoint = ""; + if (isDataplaneRbacEnabledForProxyApi(userContext)) { + hrefEndpoint = getEnvironmentScopeEndpoint(); + } else { + hrefEndpoint = new URL(account.properties.documentEndpoint).href.replace(/\/$/, "/.default"); + } const msalInstance = await getMsalInstance(); const cachedAccount = msalInstance.getAllAccounts()?.[0]; msalInstance.setActiveAccount(cachedAccount);