diff --git a/src/Contracts/ViewModels.ts b/src/Contracts/ViewModels.ts index 2df718293..dcb19f1a8 100644 --- a/src/Contracts/ViewModels.ts +++ b/src/Contracts/ViewModels.ts @@ -406,6 +406,7 @@ export enum TerminalKind { Cassandra = 2, Postgres = 3, VCoreMongo = 4, + CosmosDB = 5, } export interface DataExplorerInputsFrame { diff --git a/src/Explorer/Explorer.tsx b/src/Explorer/Explorer.tsx index bcacb8d65..bf1c5a33e 100644 --- a/src/Explorer/Explorer.tsx +++ b/src/Explorer/Explorer.tsx @@ -666,6 +666,10 @@ export default class Explorer { title = "Mongo Shell"; break; + case ViewModels.TerminalKind.CosmosDB: + title = "Cosmos DB Shell"; + break; + default: throw new Error("Terminal kind: ${kind} not supported"); } diff --git a/src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx b/src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx index 25ed4ea4f..911191e0d 100644 --- a/src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx +++ b/src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx @@ -17,11 +17,11 @@ import SynapseIcon from "../../../../images/synapse-link.svg"; import VSCodeIcon from "../../../../images/vscode.svg"; import { AuthType } from "../../../AuthType"; import * as Constants from "../../../Common/Constants"; -import { Platform, configContext } from "../../../ConfigContext"; +import { configContext, Platform } from "../../../ConfigContext"; import * as ViewModels from "../../../Contracts/ViewModels"; import { - userContext, isVCoreMongoNativeAuthDisabled, + userContext, VCoreMongoNativeAuthDisabledMessage, VCoreMongoNativeAuthLearnMoreUrl, } from "../../../UserContext"; @@ -81,6 +81,15 @@ export function createStaticCommandBarButtons( } } + if ( + userContext.apiType === "SQL" && + userContext.features.enableCloudShell && + userContext.features.enableCosmosDBShell + ) { + addDivider(); + buttons.push(createOpenTerminalButtonByKind(container, ViewModels.TerminalKind.CosmosDB)); + } + if (!selectedNodeState.isDatabaseNodeOrNoneSelected()) { const isQuerySupported = userContext.apiType === "SQL" || userContext.apiType === "Gremlin"; @@ -498,6 +507,8 @@ function createOpenTerminalButtonByKind( return "PSQL"; case ViewModels.TerminalKind.VCoreMongo: return "MongoDB (DocumentDB)"; + case ViewModels.TerminalKind.CosmosDB: + return "Cosmos DB"; default: return ""; } @@ -507,7 +518,9 @@ function createOpenTerminalButtonByKind( "This feature is not yet available in your account's region. View supported regions here: https://aka.ms/cosmos-enable-notebooks."; const isNativeAuthDisabled = terminalKind === ViewModels.TerminalKind.VCoreMongo && isVCoreMongoNativeAuthDisabled(); const disableButton = - (!useNotebook.getState().isNotebooksEnabledForAccount && !useNotebook.getState().isNotebookEnabled) || + (!useNotebook.getState().isNotebooksEnabledForAccount && + !useNotebook.getState().isNotebookEnabled && + !userContext.features.enableCloudShell) || isNativeAuthDisabled; return { iconSrc: HostedTerminalIcon, diff --git a/src/Explorer/Tabs/CloudShellTab/Data/CloudShellClient.test.tsx b/src/Explorer/Tabs/CloudShellTab/Data/CloudShellClient.test.tsx index 61daecd15..cc89cb34f 100644 --- a/src/Explorer/Tabs/CloudShellTab/Data/CloudShellClient.test.tsx +++ b/src/Explorer/Tabs/CloudShellTab/Data/CloudShellClient.test.tsx @@ -127,6 +127,7 @@ describe("CloudShellClient", () => { vnetSettings: {}, }, }, + timeoutMs: 30000, }); expect(result).toEqual(mockResponse); }); @@ -154,6 +155,7 @@ describe("CloudShellClient", () => { vnetSettings: mockVNetSettings, }, }, + timeoutMs: 30000, }); }); @@ -212,6 +214,7 @@ describe("CloudShellClient", () => { path: "/subscriptions/sub-id/providers/Microsoft.CloudShell/register", method: "POST", apiVersion: "2022-12-01", + timeoutMs: 30000, }); expect(result).toEqual(mockResponse); }); @@ -227,6 +230,7 @@ describe("CloudShellClient", () => { path: "/subscriptions/sub-id/providers/Microsoft.CloudShell/register", method: "POST", apiVersion: "2022-12-01", + timeoutMs: 30000, }); }); }); @@ -251,6 +255,7 @@ describe("CloudShellClient", () => { osType: OsType.Linux, }, }, + timeoutMs: 30000, }); expect(result).toEqual(mockResponse); }); @@ -274,6 +279,7 @@ describe("CloudShellClient", () => { osType: OsType.Linux, }, }, + timeoutMs: 30000, }); }); }); diff --git a/src/Explorer/Tabs/CloudShellTab/Data/CloudShellClient.tsx b/src/Explorer/Tabs/CloudShellTab/Data/CloudShellClient.tsx index ee4bd01e0..2f7926cae 100644 --- a/src/Explorer/Tabs/CloudShellTab/Data/CloudShellClient.tsx +++ b/src/Explorer/Tabs/CloudShellTab/Data/CloudShellClient.tsx @@ -14,6 +14,13 @@ import { } from "../Models/DataModels"; import { getLocale } from "../Utils/CommonUtils"; +// armRequest defaults to a 5s timeout with no retry for PUT/POST. Cloud Shell provisioning +// (registering the provider, applying settings, and provisioning the console itself) can +// legitimately take much longer than that on first use or when the console region differs +// from the user's usual region, so a much more generous timeout is used for these calls to +// avoid a spurious AbortError ("User aborted query.") aborting the whole session start. +const CLOUDSHELL_ARM_TIMEOUT_MS = 30000; + export const getUserSettings = async (): Promise => { return await armRequest({ host: configContext.ARM_ENDPOINT, @@ -51,6 +58,7 @@ export const putEphemeralUserSettings = async ( method: "PUT", apiVersion: "2023-02-01-preview", body: ephemeralSettings, + timeoutMs: CLOUDSHELL_ARM_TIMEOUT_MS, }); }; @@ -69,6 +77,7 @@ export const registerCloudShellProvider = async (subscriptionId: string) => { path: `/subscriptions/${subscriptionId}/providers/Microsoft.CloudShell/register`, method: "POST", apiVersion: "2022-12-01", + timeoutMs: CLOUDSHELL_ARM_TIMEOUT_MS, }); }; @@ -88,6 +97,7 @@ export const provisionConsole = async (consoleLocation: string): Promise ({ + userContext: { + databaseAccount: { + properties: { + documentEndpoint: "https://test-account.documents.azure.com:443/", + }, + }, + }, +})); + +describe("CosmosDBShellHandler", () => { + const mockKey = "testKey"; + const endpoint = "https://test-account.documents.azure.com:443/"; + const tokenConnectionCommand = `export COSMOSDB_SHELL_TOKEN='aadToken123'; cosmosdbshell --connect '${endpoint}' --connect-mode gateway --verbose`; + const keyConnectionCommand = `export COSMOSDB_SHELL_ACCOUNT_KEY='${mockKey}'; cosmosdbshell --connect '${endpoint}' --connect-mode gateway --verbose`; + let cosmosDBShellHandler: CosmosDBShellHandler; + + beforeEach(() => { + cosmosDBShellHandler = new CosmosDBShellHandler({ kind: "key", value: mockKey }); + jest.clearAllMocks(); + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + afterAll(() => { + jest.resetAllMocks(); + jest.restoreAllMocks(); + jest.resetModules(); + }); + + describe("Positive Tests", () => { + it("should return correct shell name", () => { + expect(cosmosDBShellHandler.getShellName()).toBe("Cosmos DB"); + }); + + it("should install the CosmosDBShell tool in setup commands", () => { + const commands = cosmosDBShellHandler.getSetUpCommands(); + + expect(Array.isArray(commands)).toBe(true); + expect(commands.some((c) => c.includes("dotnet tool install --global CosmosDBShell --prerelease"))).toBe(true); + expect(commands.some((c) => c.includes("dotnet tool update --global CosmosDBShell --prerelease"))).toBe(true); + expect(commands.some((c) => c.includes("$HOME/.dotnet/tools"))).toBe(true); + }); + + it("should bootstrap a .NET SDK 10 when the tool and SDK are missing", () => { + const commands = cosmosDBShellHandler.getSetUpCommands(); + + expect(commands.some((c) => c.includes("dotnet-install.sh") && c.includes("--channel 10.0"))).toBe(true); + expect(commands.some((c) => c === "export DOTNET_ROOT=$HOME/.dotnet")).toBe(true); + }); + + it("should not export any credential env var in setup commands for a key credential", () => { + const commands = cosmosDBShellHandler.getSetUpCommands(); + + expect(commands.some((c) => c.includes("COSMOSDB_SHELL_"))).toBe(false); + }); + + it("should not export any credential env var in setup commands for a token credential", () => { + const handler = new CosmosDBShellHandler({ kind: "token", value: "aadToken123" }); + const commands = handler.getSetUpCommands(); + + expect(commands.some((c) => c.includes("COSMOSDB_SHELL_"))).toBe(false); + }); + + it("should export the account key and connect to the bare endpoint on the same command line", () => { + expect(cosmosDBShellHandler.getConnectionCommand()).toBe(keyConnectionCommand); + }); + + it("should export the token and connect to the bare endpoint on the same command line", () => { + const handler = new CosmosDBShellHandler({ kind: "token", value: "aadToken123" }); + + expect(handler.getConnectionCommand()).toBe(tokenConnectionCommand); + }); + + it("should never pass an interactive or ambient credential flag", () => { + const handler = new CosmosDBShellHandler({ kind: "token", value: "aadToken123" }); + const connectionCommand = handler.getConnectionCommand(); + + expect(connectionCommand).not.toContain("--connect-tenant"); + expect(connectionCommand).not.toContain("--connect-hint"); + expect(connectionCommand).not.toContain("--connect-authority-host"); + expect(connectionCommand).not.toContain("--connect-azure-cli"); + expect(connectionCommand).not.toContain("--connect-managed-identity"); + expect(connectionCommand).not.toContain("--connect-vscode-credential"); + }); + + it("should return empty array for terminal suppressed data", () => { + expect(cosmosDBShellHandler.getTerminalSuppressedData()).toEqual([]); + }); + }); + + describe("Negative Tests", () => { + it("should not export a credential env var when no credential was resolved", () => { + const handler = new CosmosDBShellHandler(undefined); + const commands = handler.getSetUpCommands(); + + expect(commands.some((c) => c.includes("COSMOSDB_SHELL_"))).toBe(false); + }); + + it("should not launch the shell when no credential was resolved", () => { + const handler = new CosmosDBShellHandler(undefined); + const connectionCommand = handler.getConnectionCommand(); + + expect(connectionCommand).not.toContain("cosmosdbshell --connect"); + expect(connectionCommand).toContain("Unable to acquire a Cosmos DB credential"); + expect(connectionCommand).toContain("Login for Entra ID"); + }); + + it("should echo the specific reason when one is provided", () => { + const handler = new CosmosDBShellHandler(undefined, "listing the account keys failed (Forbidden)"); + const connectionCommand = handler.getConnectionCommand(); + + expect(connectionCommand).toContain("Unable to acquire a Cosmos DB credential"); + expect(connectionCommand).toContain("listing the account keys failed (Forbidden)"); + }); + }); +}); diff --git a/src/Explorer/Tabs/CloudShellTab/ShellTypes/CosmosDBShellHandler.tsx b/src/Explorer/Tabs/CloudShellTab/ShellTypes/CosmosDBShellHandler.tsx new file mode 100644 index 000000000..ed1b9658d --- /dev/null +++ b/src/Explorer/Tabs/CloudShellTab/ShellTypes/CosmosDBShellHandler.tsx @@ -0,0 +1,129 @@ +import { userContext } from "../../../../UserContext"; +import { AbstractShellHandler } from "./AbstractShellHandler"; + +/** + * A credential resolved by Data Explorer and handed to the Cosmos DB Shell. + * + * Both kinds are delivered the same way the Mongo shell delivers its credential: exported + * as an environment variable on the same command line as the `cosmosdbshell` invocation + * (`export VAR='...'; cosmosdbshell --connect `), immediately before the tool + * reads it, rather than as a separate setup step or embedded in a connection string. + * + * - `token` is an Entra ID (AAD) bearer token scoped to the account's data plane, exported + * as COSMOSDB_SHELL_TOKEN. + * - `key` is an account master (or read-only) key, exported as COSMOSDB_SHELL_ACCOUNT_KEY. + * + * The kind is tracked explicitly so the correct environment variable is always used. + */ +export interface CosmosDBShellCredential { + kind: "token" | "key"; + value: string; +} + +/** + * Shell handler for the Azure Cosmos DB Shell (https://github.com/Azure/CosmosDBShell), + * a .NET global tool that targets the Cosmos DB NoSQL (SQL Core) API. + */ +export class CosmosDBShellHandler extends AbstractShellHandler { + private _endpoint: string | undefined; + + constructor( + private credential: CosmosDBShellCredential | undefined, + private unavailableReason?: string, + ) { + super(); + this._endpoint = userContext?.databaseAccount?.properties?.documentEndpoint; + } + + public getShellName(): string { + return "Cosmos DB"; + } + + /** + * Setup commands for the Cosmos DB Shell: + * + * 1. Put the private .NET install dir and global tools dir on PATH (and set + * DOTNET_ROOT) so both the `dotnet` host and installed tool resolve. + * 2. Bootstrap a private .NET SDK 10 into $HOME/.dotnet when neither the + * cosmosdbshell tool nor a suitable SDK is already available. The + * CosmosDBShell global tool targets net10.0, so `dotnet tool install` + * requires the .NET SDK 10.0+, which the Azure Cloud Shell host does not + * ship by default. + * 3. Install the CosmosDBShell global tool if it is not already available, or + * update it to the latest release when it is (so an older cached install in + * the persistent Cloud Shell $HOME picks up newer connect options). + * 4. Persist the PATH/DOTNET_ROOT changes for future sessions. + * + * The credential itself is not exported here: it travels on the same command line as + * the `cosmosdbshell` invocation (see getConnectionCommand), mirroring how the Mongo + * shell handler builds its connection command. + * + * Installation steps run conditionally only if cosmosdbshell is not already + * present in the environment. + */ + public getSetUpCommands(): string[] { + return [ + "export DOTNET_ROOT=$HOME/.dotnet", + "export PATH=$HOME/.dotnet:$HOME/.dotnet/tools:$PATH", + "if ! command -v cosmosdbshell &> /dev/null; then echo '⚠️ cosmosdbshell not found. Installing .NET SDK 10 and CosmosDBShell...'; fi", + "if ! command -v cosmosdbshell &> /dev/null && ! dotnet --list-sdks 2>/dev/null | grep -q '^10\\.'; then curl -sSL https://dot.net/v1/dotnet-install.sh | bash -s -- --channel 10.0 --install-dir $HOME/.dotnet; fi", + "if ! command -v cosmosdbshell &> /dev/null; then dotnet tool install --global CosmosDBShell --prerelease; else dotnet tool update --global CosmosDBShell --prerelease; fi", + "grep -qxF 'export DOTNET_ROOT=$HOME/.dotnet' ~/.bashrc || echo 'export DOTNET_ROOT=$HOME/.dotnet' >> ~/.bashrc", + "grep -qxF 'export PATH=$HOME/.dotnet:$HOME/.dotnet/tools:$PATH' ~/.bashrc || echo 'export PATH=$HOME/.dotnet:$HOME/.dotnet/tools:$PATH' >> ~/.bashrc", + ]; + } + + private _getKeyConnectionCommand(key: string): string { + // Export the key immediately before invoking the tool, on the same command line, so it + // never appears as a --connect flag value or lands in a separate setup step. The tool + // reads the account key from COSMOSDB_SHELL_ACCOUNT_KEY and connects to the bare endpoint. + return `export COSMOSDB_SHELL_ACCOUNT_KEY='${key}'; cosmosdbshell --connect '${this._endpoint}' --connect-mode gateway --verbose`; + } + + private _getTokenConnectionCommand(token: string): string { + // Same pattern for the Entra ID token: exported right before the invocation so it never + // appears in argv/ps, then the tool reads it from COSMOSDB_SHELL_TOKEN. + return `export COSMOSDB_SHELL_TOKEN='${token}'; cosmosdbshell --connect '${this._endpoint}' --connect-mode gateway --verbose`; + } + + public getConnectionCommand(): string { + if (!this._endpoint) { + return `echo '${this.getShellName()} endpoint not found.'`; + } + + if (!this.credential) { + // Never let the tool continue without a credential. With no account key and no + // COSMOSDB_SHELL_TOKEN exported, its credential chain would fall through to + // DefaultAzureCredential, which in Azure Cloud Shell tries the managed identity + // first and fails with "AudienceNotSupported" (the Cloud Shell MSI cannot mint a + // token for the *.documents.azure.com audience). The interactive browser and + // device-code flows are not usable from this embedded terminal either, so fail + // fast with actionable guidance instead. + // + // There is no way to inspect the browser dev tools console from inside this remote + // terminal, so the specific reason (when known) is echoed here too. + const reasonSuffix = this.unavailableReason ? ` — ${this.unavailableReason}` : ""; + return `echo 'Unable to acquire a ${this.getShellName()} credential${reasonSuffix}. Use "Login for Entra ID" in the Data Explorer toolbar, verify you have Cosmos DB data-plane RBAC access to this account, then reopen the shell.'`; + } + + // Force gateway (HTTPS/443) connection mode. The shell otherwise defaults to + // direct (TCP) mode for real accounts, and Azure Cloud Shell blocks the + // direct-mode TCP ports, causing the connection to fail. `--verbose` surfaces + // the full exception details when a connection attempt fails. + // + // Surface which credential kind Explorer resolved in the browser console. A wrong + // guess here (e.g. a token when the account actually needs a key) otherwise fails + // silently inside the remote shell with no client-side signal to debug against. + console.warn(`CloudShell: connecting to Cosmos DB shell with a "${this.credential.kind}" credential`, { + endpoint: this._endpoint, + }); + + return this.credential.kind === "key" + ? this._getKeyConnectionCommand(this.credential.value) + : this._getTokenConnectionCommand(this.credential.value); + } + + public getTerminalSuppressedData(): string[] { + return []; + } +} diff --git a/src/Explorer/Tabs/CloudShellTab/ShellTypes/ShellTypeFactory.test.tsx b/src/Explorer/Tabs/CloudShellTab/ShellTypes/ShellTypeFactory.test.tsx index 807844514..2096ef1b1 100644 --- a/src/Explorer/Tabs/CloudShellTab/ShellTypes/ShellTypeFactory.test.tsx +++ b/src/Explorer/Tabs/CloudShellTab/ShellTypes/ShellTypeFactory.test.tsx @@ -1,19 +1,23 @@ import { TerminalKind } from "../../../../Contracts/ViewModels"; import { userContext } from "../../../../UserContext"; -import { listKeys } from "../../../../Utils/arm/generatedClients/cosmos/databaseAccounts"; +import { getReadOnlyKeys, listKeys } from "../../../../Utils/arm/generatedClients/cosmos/databaseAccounts"; +import { acquireMsalTokenForAccount, getMsalInstance } from "../../../../Utils/AuthorizationUtils"; import { CassandraShellHandler } from "./CassandraShellHandler"; +import { CosmosDBShellHandler } from "./CosmosDBShellHandler"; import { MongoShellHandler } from "./MongoShellHandler"; import { PostgresShellHandler } from "./PostgresShellHandler"; -import { getHandler, getKey } from "./ShellTypeFactory"; +import type { CosmosDBShellCredentialDiagnostics } from "./ShellTypeFactory"; +import { getCosmosDBShellCredential, getHandler, getKey } from "./ShellTypeFactory"; import { VCoreMongoShellHandler } from "./VCoreMongoShellHandler"; interface UserContextType { - databaseAccount: { name: string }; + databaseAccount: { name: string; properties?: { disableLocalAuth?: boolean } }; subscriptionId: string; resourceGroup: string; features: { enableAadDataPlane: boolean }; dataPlaneRbacEnabled: boolean; aadToken?: string; + masterKey?: string; apiType?: string; } @@ -30,13 +34,29 @@ jest.mock("../../../../UserContext", () => ({ jest.mock("../../../../Utils/arm/generatedClients/cosmos/databaseAccounts", () => ({ listKeys: jest.fn(), + getReadOnlyKeys: jest.fn(), })); +jest.mock("../../../../Utils/AuthorizationUtils", () => { + const actual = jest.requireActual("../../../../Utils/AuthorizationUtils"); + return { + ...actual, + getMsalInstance: jest.fn(), + acquireMsalTokenForAccount: jest.fn(), + }; +}); + describe("ShellTypeHandlerFactory", () => { const mockKey = "testKey"; + const mockMsalAccounts = (accounts: { username: string }[]): void => { + (getMsalInstance as jest.Mock).mockResolvedValue({ getAllAccounts: (): { username: string }[] => accounts }); + }; + beforeEach(() => { (listKeys as jest.Mock).mockResolvedValue({ primaryMasterKey: mockKey }); + mockMsalAccounts([]); + (acquireMsalTokenForAccount as jest.Mock).mockResolvedValue(""); }); afterEach(() => { @@ -69,7 +89,7 @@ describe("ShellTypeHandlerFactory", () => { type DatabaseAccountType = { name: string }; (userContext.databaseAccount as DatabaseAccountType).name = ""; - const key = await getKey(); + const key = await getKey(false); expect(key).toBe(""); expect(listKeys).not.toHaveBeenCalled(); @@ -80,7 +100,7 @@ describe("ShellTypeHandlerFactory", () => { it("should return empty string when listKeys returns null", async () => { (listKeys as jest.Mock).mockResolvedValue(null); - const key = await getKey(); + const key = await getKey(false); expect(key).toBe(""); }); @@ -89,7 +109,7 @@ describe("ShellTypeHandlerFactory", () => { /* no primaryMasterKey */ }); - const key = await getKey(); + const key = await getKey(false); expect(key).toBe(""); }); }); @@ -116,12 +136,57 @@ describe("ShellTypeHandlerFactory", () => { expect(handler).toBeInstanceOf(CassandraShellHandler); }); + it("should return CosmosDBShellHandler with key for CosmosDB terminal kind", async () => { + const handler = await getHandler(TerminalKind.CosmosDB); + expect(handler).toBeInstanceOf(CosmosDBShellHandler); + }); + it("should get key successfully when database name exists", async () => { - const key = await getKey(); + const key = await getKey(false); expect(key).toBe(mockKey); expect(listKeys).toHaveBeenCalledWith("testSubId", "testResourceGroup", "testDbName"); }); + it("should return the aadToken without listing keys when Entra ID auth is requested", async () => { + (userContext as UserContextType).aadToken = "aadToken123"; + + const key = await getKey(true); + expect(key).toBe("aadToken123"); + expect(listKeys).not.toHaveBeenCalled(); + }); + + it("should return an empty string when Entra ID auth is requested but no cached token or account exists", async () => { + (userContext as UserContextType).aadToken = undefined; + mockMsalAccounts([]); + + const key = await getKey(true); + expect(key).toBe(""); + expect(acquireMsalTokenForAccount).not.toHaveBeenCalled(); + expect(listKeys).not.toHaveBeenCalled(); + }); + + it("should silently mint a Cosmos token when Entra ID auth is requested and an MSAL account is cached", async () => { + (userContext as UserContextType).aadToken = undefined; + mockMsalAccounts([{ username: "user@contoso.com" }]); + (acquireMsalTokenForAccount as jest.Mock).mockResolvedValue("mintedToken123"); + + const key = await getKey(true); + expect(key).toBe("mintedToken123"); + expect(acquireMsalTokenForAccount).toHaveBeenCalled(); + expect(listKeys).not.toHaveBeenCalled(); + }); + + it("should return an empty string when the silent token acquisition fails", async () => { + (userContext as UserContextType).aadToken = undefined; + mockMsalAccounts([{ username: "user@contoso.com" }]); + (acquireMsalTokenForAccount as jest.Mock).mockRejectedValue(new Error("interaction_required")); + jest.spyOn(console, "error").mockImplementation(() => undefined); + + const key = await getKey(true); + expect(key).toBe(""); + expect(listKeys).not.toHaveBeenCalled(); + }); + 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; @@ -150,4 +215,148 @@ describe("ShellTypeHandlerFactory", () => { ); }); }); + + describe("getCosmosDBShellCredential", () => { + beforeEach(() => { + (userContext as UserContextType).aadToken = undefined; + (userContext as UserContextType).masterKey = undefined; + (userContext as UserContextType).apiType = "SQL"; + (userContext as UserContextType).features.enableAadDataPlane = false; + (userContext as UserContextType).dataPlaneRbacEnabled = false; + (userContext as UserContextType).databaseAccount.properties = { disableLocalAuth: false }; + }); + + it("should reuse the master key Data Explorer already resolved without calling ARM", async () => { + (userContext as UserContextType).masterKey = "cachedMasterKey"; + + const credential = await getCosmosDBShellCredential(); + + expect(credential).toEqual({ kind: "key", value: "cachedMasterKey" }); + expect(listKeys).not.toHaveBeenCalled(); + expect(getReadOnlyKeys).not.toHaveBeenCalled(); + }); + + it("should fall back to the read-only keys when the caller lacks list-keys permission", async () => { + const authorizationFailed = Object.assign(new Error("AuthorizationFailed"), { code: "AuthorizationFailed" }); + (listKeys as jest.Mock).mockRejectedValue(authorizationFailed); + (getReadOnlyKeys as jest.Mock).mockResolvedValue({ primaryReadonlyMasterKey: "readOnlyKey" }); + + const credential = await getCosmosDBShellCredential(); + + expect(credential).toEqual({ kind: "key", value: "readOnlyKey" }); + expect(getReadOnlyKeys).toHaveBeenCalledWith("testSubId", "testResourceGroup", "testDbName"); + }); + + it("should return undefined when both the read-write and read-only key fetches fail", async () => { + const authorizationFailed = Object.assign(new Error("AuthorizationFailed"), { code: "AuthorizationFailed" }); + (listKeys as jest.Mock).mockRejectedValue(authorizationFailed); + (getReadOnlyKeys as jest.Mock).mockRejectedValue(new Error("Forbidden")); + jest.spyOn(console, "error").mockImplementation(() => undefined); + + const credential = await getCosmosDBShellCredential(); + + expect(credential).toBeUndefined(); + }); + + it("should return the account key when Entra ID auth is not enabled", async () => { + const credential = await getCosmosDBShellCredential(); + + expect(credential).toEqual({ kind: "key", value: mockKey }); + expect(listKeys).toHaveBeenCalledWith("testSubId", "testResourceGroup", "testDbName"); + }); + + it("should return the cached aadToken when Entra ID auth is enabled", async () => { + (userContext as UserContextType).dataPlaneRbacEnabled = true; + (userContext as UserContextType).aadToken = "aadToken123"; + + const credential = await getCosmosDBShellCredential(); + + expect(credential).toEqual({ kind: "token", value: "aadToken123" }); + expect(listKeys).not.toHaveBeenCalled(); + }); + + it("should return a silently minted token when an MSAL account is cached", async () => { + (userContext as UserContextType).dataPlaneRbacEnabled = true; + mockMsalAccounts([{ username: "user@contoso.com" }]); + (acquireMsalTokenForAccount as jest.Mock).mockResolvedValue("mintedToken123"); + + const credential = await getCosmosDBShellCredential(); + + expect(credential).toEqual({ kind: "token", value: "mintedToken123" }); + expect(listKeys).not.toHaveBeenCalled(); + }); + + it("should fall back to the account key when no token could be resolved", async () => { + (userContext as UserContextType).dataPlaneRbacEnabled = true; + + const credential = await getCosmosDBShellCredential(); + + expect(acquireMsalTokenForAccount).not.toHaveBeenCalled(); + expect(credential).toEqual({ kind: "key", value: mockKey }); + }); + + it("should fall back to the account key when the silent token acquisition throws", async () => { + (userContext as UserContextType).dataPlaneRbacEnabled = true; + mockMsalAccounts([{ username: "user@contoso.com" }]); + (acquireMsalTokenForAccount as jest.Mock).mockRejectedValue(new Error("interaction_required")); + jest.spyOn(console, "error").mockImplementation(() => undefined); + + const credential = await getCosmosDBShellCredential(); + + expect(credential).toEqual({ kind: "key", value: mockKey }); + }); + + it("should not fall back to the account key when local auth is disabled", async () => { + (userContext as UserContextType).databaseAccount.properties = { disableLocalAuth: true }; + jest.spyOn(console, "warn").mockImplementation(() => undefined); + + const credential = await getCosmosDBShellCredential(); + + expect(credential).toBeUndefined(); + expect(listKeys).not.toHaveBeenCalled(); + }); + + it("should return undefined when listing the account keys fails", async () => { + (listKeys as jest.Mock).mockRejectedValue(new Error("Forbidden")); + jest.spyOn(console, "error").mockImplementation(() => undefined); + + const credential = await getCosmosDBShellCredential(); + + expect(credential).toBeUndefined(); + }); + + it("should return undefined when the database name is missing", async () => { + (userContext as UserContextType).databaseAccount.name = ""; + + const credential = await getCosmosDBShellCredential(); + + expect(credential).toBeUndefined(); + expect(listKeys).not.toHaveBeenCalled(); + + (userContext as UserContextType).databaseAccount.name = "testDbName"; + }); + + it("should populate a specific reason when both key fetches fail, for the shell to echo", async () => { + (listKeys as jest.Mock).mockRejectedValue(new Error("Forbidden")); + (getReadOnlyKeys as jest.Mock).mockRejectedValue(new Error("Forbidden")); + jest.spyOn(console, "error").mockImplementation(() => undefined); + + const diagnostics: CosmosDBShellCredentialDiagnostics = {}; + const credential = await getCosmosDBShellCredential(diagnostics); + + expect(credential).toBeUndefined(); + expect(diagnostics.reason).toContain("Forbidden"); + }); + + it("should populate a reason when local auth is disabled and no token was resolved", async () => { + (userContext as UserContextType).databaseAccount.properties = { disableLocalAuth: true }; + jest.spyOn(console, "warn").mockImplementation(() => undefined); + + const diagnostics: CosmosDBShellCredentialDiagnostics = {}; + const credential = await getCosmosDBShellCredential(diagnostics); + + expect(credential).toBeUndefined(); + expect(diagnostics.reason).toContain("local (key) auth is disabled"); + }); + }); }); diff --git a/src/Explorer/Tabs/CloudShellTab/ShellTypes/ShellTypeFactory.tsx b/src/Explorer/Tabs/CloudShellTab/ShellTypes/ShellTypeFactory.tsx index 81e068562..0c3be0be6 100644 --- a/src/Explorer/Tabs/CloudShellTab/ShellTypes/ShellTypeFactory.tsx +++ b/src/Explorer/Tabs/CloudShellTab/ShellTypes/ShellTypeFactory.tsx @@ -1,9 +1,15 @@ import { TerminalKind } from "../../../../Contracts/ViewModels"; import { userContext } from "../../../../UserContext"; -import { listKeys } from "../../../../Utils/arm/generatedClients/cosmos/databaseAccounts"; -import { isDataplaneRbacEnabledForProxyApi } from "../../../../Utils/AuthorizationUtils"; +import { getReadOnlyKeys, listKeys } from "../../../../Utils/arm/generatedClients/cosmos/databaseAccounts"; +import { + acquireMsalTokenForAccount, + getMsalInstance, + isCloudShellEntraAuthEnabled, + isDataplaneRbacEnabledForProxyApi, +} from "../../../../Utils/AuthorizationUtils"; import { AbstractShellHandler } from "./AbstractShellHandler"; import { CassandraShellHandler } from "./CassandraShellHandler"; +import { CosmosDBShellCredential, CosmosDBShellHandler } from "./CosmosDBShellHandler"; import { MongoShellHandler } from "./MongoShellHandler"; import { PostgresShellHandler } from "./PostgresShellHandler"; import { VCoreMongoShellHandler } from "./VCoreMongoShellHandler"; @@ -16,23 +22,244 @@ export async function getHandler(shellType: TerminalKind): Promise { +/** + * Populated by {@link getCosmosDBShellCredential} (and its helpers) with a human-readable + * explanation when no credential could be resolved, so the shell can echo a specific cause + * instead of a generic message. There is no way to inspect the browser dev tools console + * from inside the remote Cloud Shell terminal, so this is surfaced directly in the shell. + */ +export interface CosmosDBShellCredentialDiagnostics { + reason?: string; +} + +/** + * Resolves the credential Data Explorer already holds for the current account and hands + * it to the Cosmos DB Shell. + * + * The Cloud Shell cannot authenticate to Cosmos DB on its own: its managed identity is + * rejected with "AudienceNotSupported" for the `*.documents.azure.com` audience, its `az` + * session is not signed in, and the interactive browser/device-code flows are not usable + * from the embedded terminal. The credential therefore has to come from Data Explorer. + * + * Resolution order: + * 1. Entra ID data-plane token — the cached `userContext.aadToken`, or a silently minted + * one. Silent acquisition is only attempted when an MSAL account is already cached so + * it can never trigger an interactive popup. + * 2. Account master key — the key Data Explorer already resolved (`userContext.masterKey`), + * or, if that is not populated, one fetched via ARM (`listKeys`, falling back to + * `getReadOnlyKeys` on any failure for read-only callers). The key is handed to the + * handler, which exports it as COSMOSDB_SHELL_ACCOUNT_KEY on the same command line as + * the `cosmosdbshell` invocation. Skipped when the account has local auth disabled + * (keys do not exist). + * + * Returns `undefined` when nothing could be resolved, which makes the handler surface + * actionable guidance instead of letting the tool attempt its own sign-in. + */ +export async function getCosmosDBShellCredential( + diagnostics?: CosmosDBShellCredentialDiagnostics, +): Promise { + const dbName = userContext.databaseAccount?.name; + if (!dbName) { + if (diagnostics) { + diagnostics.reason = "no database account name is available"; + } + return undefined; + } + + if (isCloudShellEntraAuthEnabled(userContext)) { + const token = await resolveCosmosDataPlaneToken(); + if (token) { + return { kind: "token", value: token }; + } + } + + if (userContext.databaseAccount?.properties?.disableLocalAuth) { + const reason = + "local (key) auth is disabled on this account and no Entra ID token could be resolved; sign in via " + + '"Login for Entra ID" in the toolbar and reopen the shell'; + console.warn(`CloudShell: ${reason}.`); + if (diagnostics) { + diagnostics.reason = reason; + } + return undefined; + } + + const key = await resolveAccountKey(dbName, diagnostics); + if (!key) { + if (!diagnostics?.reason) { + const reason = + "no account key could be resolved. Data Explorer may be authenticating through the portal proxy " + + "(per-request tokens), which cannot be reused by the shell. Ensure you have either data-plane RBAC " + + 'access (then use "Login for Entra ID") or permission to list the account keys'; + console.warn(`CloudShell: ${reason}.`); + if (diagnostics) { + diagnostics.reason = reason; + } + } + return undefined; + } + return { kind: "key", value: key }; +} + +/** + * Returns a Cosmos data-plane token without ever prompting. Silent acquisition is guarded + * on an existing cached MSAL account because `acquireMsalTokenForAccount` falls back to an + * interactive popup when none exists, which cannot complete in the hosted Cloud Shell. + */ +async function resolveCosmosDataPlaneToken(): Promise { + if (userContext.aadToken) { + return userContext.aadToken; + } + + try { + const msalInstance = await getMsalInstance(); + if (msalInstance.getAllAccounts().length === 0) { + return ""; + } + return (await acquireMsalTokenForAccount(userContext.databaseAccount, true)) || ""; + } catch (error) { + console.error("Failed to silently acquire a Cosmos data-plane token for the Cloud Shell", error); + return ""; + } +} + +/** + * Resolves the subscription id and resource group for the current account. These are + * normally populated on `userContext`, but in some hosting contexts they can be missing, + * so they are parsed from the account's ARM resource id as a fallback. The id has the form + * `/subscriptions//resourceGroups//providers/Microsoft.DocumentDB/databaseAccounts/`. + */ +function resolveAccountArmScope(): { subscriptionId: string; resourceGroup: string } { + let subscriptionId = userContext.subscriptionId; + let resourceGroup = userContext.resourceGroup; + + const accountId = userContext.databaseAccount?.id; + if ((!subscriptionId || !resourceGroup) && accountId) { + const match = accountId.match(/\/subscriptions\/([^/]+)\/resourceGroups\/([^/]+)\//i); + if (match) { + subscriptionId = subscriptionId || match[1]; + resourceGroup = resourceGroup || match[2]; + } + } + + return { subscriptionId, resourceGroup }; +} + +async function resolveAccountKey(dbName: string, diagnostics?: CosmosDBShellCredentialDiagnostics): Promise { + // Prefer the key Data Explorer already resolved for this account so the shell reuses the + // exact credential DE is connected with and avoids a redundant ARM round-trip. + if (userContext.masterKey) { + return userContext.masterKey; + } + + // Otherwise fetch it via ARM, mirroring DE's own `fetchAndUpdateKeys`: try the read-write + // keys first, then fall back to the read-only keys. The fallback is attempted on any + // failure (not just "AuthorizationFailed") because a read-only caller can surface the + // missing-permission error in different shapes; without it such a user gets no credential + // and the shell cannot connect. + const { subscriptionId, resourceGroup } = resolveAccountArmScope(); + if (!subscriptionId || !resourceGroup) { + const reason = "the subscription id or resource group for this account could not be determined"; + console.error(`CloudShell: cannot list Cosmos DB account keys because ${reason}.`); + if (diagnostics) { + diagnostics.reason = reason; + } + return ""; + } + + let listKeysErrorMessage: string | undefined; + try { + const keys = await listKeys(subscriptionId, resourceGroup, dbName); + if (keys?.primaryMasterKey) { + return keys.primaryMasterKey; + } + listKeysErrorMessage = "the response did not include a usable key"; + } catch (error) { + listKeysErrorMessage = error instanceof Error ? error.message : String(error); + console.error("Failed to list read-write account keys for the Cloud Shell; trying read-only keys", error); + } + + try { + const readOnlyKeys = await getReadOnlyKeys(subscriptionId, resourceGroup, dbName); + if (readOnlyKeys?.primaryReadonlyMasterKey) { + return readOnlyKeys.primaryReadonlyMasterKey; + } + if (diagnostics) { + diagnostics.reason = + `listing account keys failed (${listKeysErrorMessage}) and the read-only keys response did not ` + + "include a usable key either"; + } + return ""; + } catch (readOnlyError) { + console.error("Failed to list read-only account keys for the Cloud Shell", readOnlyError); + if (diagnostics) { + const readOnlyErrorMessage = readOnlyError instanceof Error ? readOnlyError.message : String(readOnlyError); + diagnostics.reason = + `listing the account keys failed (${listKeysErrorMessage}) and listing the read-only keys also failed ` + + `(${readOnlyErrorMessage}) — this is usually a missing listKeys/listReadOnlyKeys RBAC permission`; + } + return ""; + } +} + +/** + * Resolves the credential to inject into the Cloud Shell for a Mongo or Cassandra + * connection. + * + * @param useEntraIdAuth When true, returns an Entra ID (AAD) bearer token scoped to + * the account's data plane; otherwise returns the account master key. The caller + * decides which credential applies so it stays in sync with the connection command + * the matching handler builds. + * + * The Cosmos DB (NoSQL) shell does not use this function — see + * {@link getCosmosDBShellCredential}, which also reports which kind of credential it + * resolved so the correct environment variable is exported. + */ +export async function getKey(useEntraIdAuth: boolean): Promise { const dbName = userContext.databaseAccount.name; if (!dbName) { return ""; } - if (isDataplaneRbacEnabledForProxyApi(userContext)) { - return userContext.aadToken || ""; + if (useEntraIdAuth) { + if (userContext.aadToken) { + return userContext.aadToken; + } + + try { + const msalInstance = await getMsalInstance(); + if (msalInstance.getAllAccounts().length === 0) { + // No cached account to silently acquire against; a token request here would + // require an interactive popup, which cannot complete in the Cloud Shell. + console.warn( + "CloudShell: no cached MSAL account available to mint a Cosmos data-plane token; " + + "the shell will use interactive device-code authentication.", + ); + return ""; + } + const token = (await acquireMsalTokenForAccount(userContext.databaseAccount, true)) || ""; + if (!token) { + console.warn("CloudShell: silent Cosmos data-plane token acquisition returned an empty token."); + } + return token; + } catch (error) { + console.error("Failed to silently acquire a Cosmos data-plane token for the Cloud Shell", error); + return ""; + } } const keys = await listKeys(userContext.subscriptionId, userContext.resourceGroup, dbName); diff --git a/src/Explorer/Tabs/CloudShellTab/Utils/CommonUtils.test.tsx b/src/Explorer/Tabs/CloudShellTab/Utils/CommonUtils.test.tsx new file mode 100644 index 000000000..a6186084c --- /dev/null +++ b/src/Explorer/Tabs/CloudShellTab/Utils/CommonUtils.test.tsx @@ -0,0 +1,67 @@ +import { Terminal } from "xterm"; +import { askConfirmation } from "./CommonUtils"; + +describe("askConfirmation", () => { + const createTerminalMock = () => { + let keyHandler: ((event: { key: string }) => void) | undefined; + const dispose = jest.fn(); + const terminal = { + writeln: jest.fn(), + focus: jest.fn(), + onKey: jest.fn((handler: (event: { key: string }) => void) => { + keyHandler = handler; + return { dispose }; + }), + } as unknown as Terminal; + + return { + terminal, + dispose, + pressKey: (key: string) => keyHandler?.({ key }), + }; + }; + + it("resolves true when the user presses Y", async () => { + const { terminal, dispose, pressKey } = createTerminalMock(); + const promise = askConfirmation(terminal, "Proceed?"); + + pressKey("Y"); + + await expect(promise).resolves.toBe(true); + expect(dispose).toHaveBeenCalledTimes(1); + }); + + it("resolves true when the user presses lowercase y", async () => { + const { terminal, pressKey } = createTerminalMock(); + const promise = askConfirmation(terminal, "Proceed?"); + + pressKey("y"); + + await expect(promise).resolves.toBe(true); + }); + + it("resolves false when the user presses N", async () => { + const { terminal, dispose, pressKey } = createTerminalMock(); + const promise = askConfirmation(terminal, "Proceed?"); + + pressKey("N"); + + await expect(promise).resolves.toBe(false); + expect(dispose).toHaveBeenCalledTimes(1); + }); + + it("ignores keys other than Y or N and keeps listening until a valid answer", async () => { + const { terminal, dispose, pressKey } = createTerminalMock(); + const promise = askConfirmation(terminal, "Proceed?"); + + pressKey("a"); + pressKey("1"); + pressKey("\r"); + expect(dispose).not.toHaveBeenCalled(); + + pressKey("y"); + + await expect(promise).resolves.toBe(true); + expect(dispose).toHaveBeenCalledTimes(1); + }); +}); diff --git a/src/Explorer/Tabs/CloudShellTab/Utils/CommonUtils.tsx b/src/Explorer/Tabs/CloudShellTab/Utils/CommonUtils.tsx index dcd6e8681..87759d572 100644 --- a/src/Explorer/Tabs/CloudShellTab/Utils/CommonUtils.tsx +++ b/src/Explorer/Tabs/CloudShellTab/Utils/CommonUtils.tsx @@ -24,9 +24,15 @@ export const askConfirmation = async (terminal: Terminal, question: string): Pro terminal.focus(); return new Promise((resolve) => { const keyListener = terminal.onKey(({ key }: { key: string }) => { + const normalizedKey = key.toLowerCase(); + // Only "y" or "n" are accepted. Any other key is ignored so an accidental + // keypress does not abort the flow; keep listening until a valid answer. + if (normalizedKey !== "y" && normalizedKey !== "n") { + return; + } keyListener.dispose(); terminal.writeln(key); - return resolve(key.toLowerCase() === "y"); + return resolve(normalizedKey === "y"); }); }); }; @@ -46,6 +52,8 @@ export const getShellNameForDisplay = (terminalKind: TerminalKind): string => { case TerminalKind.Mongo: case TerminalKind.VCoreMongo: return "MongoDB"; + case TerminalKind.CosmosDB: + return "Cosmos DB"; default: return ""; } diff --git a/src/Platform/Hosted/extractFeatures.ts b/src/Platform/Hosted/extractFeatures.ts index 88aaac1ca..d8ccecae0 100644 --- a/src/Platform/Hosted/extractFeatures.ts +++ b/src/Platform/Hosted/extractFeatures.ts @@ -35,6 +35,7 @@ export type Features = { readonly disableConnectionStringLogin: boolean; readonly enableContainerCopy: boolean; readonly enableCloudShell: boolean; + readonly enableCosmosDBShell: boolean; readonly enableRestoreContainer: boolean; // only for Fabric readonly mongoDisableNativeAuth: boolean; @@ -104,6 +105,7 @@ export function extractFeatures(given = new URLSearchParams(window.location.sear enableContainerCopy: "true" === get("enablecontainercopy"), enableRestoreContainer: "true" === get("enablerestorecontainer"), enableCloudShell: true, + enableCosmosDBShell: "true" === get("enablecosmosdbshell"), mongoDisableNativeAuth: "true" === get("mongodisablenativeauth"), }; } diff --git a/src/Utils/AuthorizationUtils.test.ts b/src/Utils/AuthorizationUtils.test.ts index 68d448ba7..fd361ac22 100644 --- a/src/Utils/AuthorizationUtils.test.ts +++ b/src/Utils/AuthorizationUtils.test.ts @@ -17,6 +17,7 @@ describe("AuthorizationUtils", () => { updateUserContext({ features: { enableContainerCopy: false, + enableCosmosDBShell: false, enableAadDataPlane: enabled, canExceedMaximumValue: false, cosmosdb: false, diff --git a/src/Utils/AuthorizationUtils.ts b/src/Utils/AuthorizationUtils.ts index 92b6b4a53..0938c6f33 100644 --- a/src/Utils/AuthorizationUtils.ts +++ b/src/Utils/AuthorizationUtils.ts @@ -236,3 +236,21 @@ export function useDataplaneRbacAuthorization(userContext: UserContext): boolean export function isDataplaneRbacEnabledForProxyApi(userContext: UserContext): boolean { return useDataplaneRbacAuthorization(userContext) && hasProxyServer(userContext.apiType); } + +/** + * Determines whether Cloud Shell connections should authenticate with an Entra ID + * token (COSMOSDB_SHELL_TOKEN) instead of the account master key. + * + * Returns true when data-plane RBAC authorization is in effect, or when the account + * has local (key) auth disabled — in the latter case a master key would be rejected + * by the service, so we must fall back to Entra ID even if the RBAC toggle is off. + * + * Unlike {@link isDataplaneRbacEnabledForProxyApi}, this does not require a proxy + * server, so it works for the SQL/CosmosDB API which connects directly to the + * data plane. + */ +export function isCloudShellEntraAuthEnabled(userContext: UserContext): boolean { + return ( + useDataplaneRbacAuthorization(userContext) || userContext.databaseAccount?.properties?.disableLocalAuth === true + ); +}