Add RBAC permission request flow for NoSQL accounts

src/docdb/tree/DocDBAccountTreeItemBase.ts

@@ -9,10 +9,12 @@ import { AzExtParentTreeItem, AzExtTreeItem, ICreateChildImplContext } from '@mi
 import * as vscode from 'vscode';
 import { IDeleteWizardContext } from '../../commands/deleteDatabaseAccount/IDeleteWizardContext';
 import { deleteCosmosDBAccount } from '../../commands/deleteDatabaseAccount/deleteCosmosDBAccount';
-import { SERVERLESS_CAPABILITY_NAME, getThemeAgnosticIconPath } from '../../constants';
+import { getThemeAgnosticIconPath, SERVERLESS_CAPABILITY_NAME } from '../../constants';
 import { nonNullProp } from '../../utils/nonNull';
 import { rejectOnTimeout } from '../../utils/timeout';
 import { CosmosDBCredential, getCosmosClient, getCosmosKeyCredential } from '../getCosmosClient';
+import { getSignedInPrincipalIdForAccountEndpoint } from '../utils/azureSessionHelper';
+import { ensureRbacPermission, isRbacException, showRbacPermissionError } from '../utils/rbacUtils';
 import { DocDBTreeItemBase } from './DocDBTreeItemBase';
 
 /**
@@ -22,7 +24,7 @@ import { DocDBTreeItemBase } from './DocDBTreeItemBase';
 export abstract class DocDBAccountTreeItemBase extends DocDBTreeItemBase<DatabaseDefinition & Resource> {
     public readonly label: string;
     public readonly childTypeLabel: string = "Database";
-
+    private hasShownRbacNotification: boolean = false;
 
     constructor(
         parent: AzExtParentTreeItem,
@@ -88,7 +90,21 @@ export abstract class DocDBAccountTreeItemBase extends DocDBTreeItemBase<Databas
             const unableToReachEmulatorMessage: string = "Unable to reach emulator. Please ensure it is started and connected to the port specified by the 'cosmosDB.emulator.port' setting, then try again.";
             return await rejectOnTimeout(2000, () => super.loadMoreChildrenImpl(clearCache), unableToReachEmulatorMessage);
         } else {
-            return await super.loadMoreChildrenImpl(clearCache);
+            try {
+                return await super.loadMoreChildrenImpl(clearCache);
+            } catch (e) {
+                if (e instanceof Error && isRbacException(e) && !this.hasShownRbacNotification) {
+                    this.hasShownRbacNotification = true;
+                    const principalId = await getSignedInPrincipalIdForAccountEndpoint(this.root.endpoint) ?? '';
+                    // chedck if the principal ID matches the one that is signed in, otherwise this might be a security problem, hence show the error message
+                    if (e.message.includes(`[${principalId}]`) && await ensureRbacPermission(this, principalId)) {
+                        return await super.loadMoreChildrenImpl(clearCache);
+                    } else {
+                        void showRbacPermissionError(this.fullId, principalId);
+                    }
+                }
+                throw e; // rethrowing tells the resources extension to show the exception message in the tree
+            }
         }
     }
 

src/docdb/utils/azureSessionHelper.ts

@@ -0,0 +1,20 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+// eslint-disable-next-line import/no-internal-modules
+import { getSessionFromVSCode } from '@microsoft/vscode-azext-azureauth/out/src/getSessionFromVSCode';
+import * as vscode from "vscode";
+
+export async function getSignedInPrincipalIdForAccountEndpoint(accountEndpoint: string): Promise<string | undefined> {
+    const session = await getSessionForDatabaseAccount(accountEndpoint);
+    const principalId = session?.account.id.split('/')[1] ?? session?.account.id;
+    return principalId;
+}
+
+async function getSessionForDatabaseAccount(endpoint: string): Promise<vscode.AuthenticationSession | undefined> {
+    const endpointUrl = new URL(endpoint);
+    const scrope = `${endpointUrl.origin}${endpointUrl.pathname}.default`;
+    return await getSessionFromVSCode(scrope, undefined, { createIfNone: false });
+}

src/docdb/utils/rbacUtils.ts

@@ -0,0 +1,89 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { SqlRoleAssignmentCreateUpdateParameters } from '@azure/arm-cosmosdb';
+import { getResourceGroupFromId } from '@microsoft/vscode-azext-azureutils';
+import { callWithTelemetryAndErrorHandling, IActionContext, IAzureMessageOptions, ISubscriptionContext } from '@microsoft/vscode-azext-utils';
+import { randomUUID } from 'crypto';
+import * as vscode from 'vscode';
+import { createCosmosDBClient } from '../../utils/azureClients';
+import { getDatabaseAccountNameFromId } from '../../utils/azureUtils';
+import { localize } from '../../utils/localize';
+import { DocDBAccountTreeItemBase } from '../tree/DocDBAccountTreeItemBase';
+
+export async function ensureRbacPermission(docDbItem: DocDBAccountTreeItemBase, principalId: string): Promise<boolean> {
+    return await callWithTelemetryAndErrorHandling('cosmosDB.addMissingRbacRole', async (context: IActionContext) => {
+
+        context.errorHandling.suppressDisplay = false;
+        context.errorHandling.rethrow = false;
+
+        const accountName: string = getDatabaseAccountNameFromId(docDbItem.fullId);
+        if (await askForRbacPermissions(accountName, docDbItem.subscription.subscriptionDisplayName, context)) {
+            context.telemetry.properties.lastStep = "addRbacContributorPermission";
+            const resourceGroup: string = getResourceGroupFromId(docDbItem.fullId);
+            const start: number = Date.now();
+            await addRbacContributorPermission(accountName, principalId, resourceGroup, context, docDbItem.subscription);
+            //send duration of the previous call (in seconds) in addition to the duration of the whole event including user prompt
+            context.telemetry.measurements["createRoleAssignment"] = (Date.now() - start) / 1000;
+
+            return true;
+        }
+        return false;
+    }) ?? false;
+}
+
+export function isRbacException(error: Error): boolean {
+    return (error instanceof Error && error.message.includes("does not have required RBAC permissions to perform action"));
+}
+
+export async function showRbacPermissionError(accountName: string, principalId: string): Promise<void> {
+    const message = localize("rbacPermissionErrorMsg", "You do not have the required permissions to access [{0}] with your principal Id [{1}].\nPlease contact the account owner to get the required permissions.", accountName, principalId);
+    const readMoreItem = localize("learnMore", "Learn More");
+    await vscode.window.showErrorMessage(message, { modal: false }, ...[readMoreItem]).then((item) => {
+        if (item === readMoreItem) {
+            void vscode.env.openExternal(vscode.Uri.parse("https://aka.ms/cosmos-native-rbac"));
+        }
+    });
+}
+
+async function askForRbacPermissions(databaseAccount: string, subscription: string, context: IActionContext): Promise<boolean> {
+    const message =
+        [localize("rbacMissingErrorMsg", "You need the 'Data Contributor' RBAC role permission to enable all Azure Databases Extension features for the selected account.\n\n"),
+        localize("rbacMissingErrorAccountName", "Account Name: {0}\n", databaseAccount),
+        localize("rbacMissingErrorSubscriptionName", "Subscription: {0}\n", subscription)
+        ].join("");
+    const options: IAzureMessageOptions = { modal: true, detail: message, learnMoreLink: "https://aka.ms/cosmos-native-rbac", stepName: "askSetRbac" };
+    const setPermissionItem: vscode.MessageItem = { title: localize("rbacExtendPermissionBtn", "Extend RBAC permissions") };
+
+    const result = await context.ui.showWarningMessage(localize("rbacMissingErrorTitle", "No required RBAC permissions"), options, ...[setPermissionItem]);
+    return result === setPermissionItem;
+}
+
+async function addRbacContributorPermission(databaseAccount: string, principalId: string, resourceGroup: string, context: IActionContext, subscription: ISubscriptionContext): Promise<string | undefined> {
+    const defaultRoleId = "00000000-0000-0000-0000-000000000002"; // this is a predefined role with read and write access to data plane resources
+    const fullAccountId = `/subscriptions/${subscription.subscriptionId}/resourceGroups/${resourceGroup}/providers/Microsoft.DocumentDB/databaseAccounts/${databaseAccount}`;
+
+    const createUpdateSqlRoleAssignmentParameters: SqlRoleAssignmentCreateUpdateParameters =
+    {
+        principalId: principalId,
+        roleDefinitionId: fullAccountId + "/sqlRoleDefinitions/" + defaultRoleId,
+        scope: fullAccountId,
+    };
+
+    /*
+    // TODO: find a better way to check if a role assignment for the current user already exists,
+    // iterating over all role assignments and definitions is not efficient.
+    const rbac = client.sqlResources.listSqlRoleAssignments(resourceGroup, databaseAccount)
+    for await (const role of rbac) {
+        console.log(role);
+    }*/
+
+    const roleAssignmentId = randomUUID();
+    const client = await createCosmosDBClient([context, subscription]);
+    const create = await client.sqlResources.beginCreateUpdateSqlRoleAssignmentAndWait(roleAssignmentId, resourceGroup, databaseAccount, createUpdateSqlRoleAssignmentParameters);
+
+    return create.id;
+}
+