Add RBAC permission request flow for NoSQL accounts (#2289)

* Add RBAC permission request flow for NoSQL accounts

If connecting to a database account using AAD fails because of missing
RBAC permissions, we now notify the user with instructions and
an option to assign a contributor role for them.

If this fails, show an error notification with a link to RBAC instructions.

* Show RBAC notification only once per account item.

Sometimes the resource tree continously refreshes all expanded nodes
and if authentication fails the RBAC notification pops up again and again.

With this change we show the notification only once for each account node.
Refreshing the parent (resource group) node will create a new instance
and show the notification again.

* Minor code style changes

* Use IActionContext to show the RBAC notification

* Add RBAC flow localization support

* Add Telemetry to RBAC permission request

In addition leave error handling to parent resources extension.

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 {
+            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;
+}
+