From 8ae59be728b766b5430eb6047c48e3e6bc97bd89 Mon Sep 17 00:00:00 2001
From: Sevo Kukol <sevoku@microsoft.com>
Date: Wed, 20 Nov 2024 16:36:58 +0100
Subject: [PATCH] Add missing multi-tenant support. Select the right Tenant
 when authenticating SDK calls and pass it along to VSCode when retrieving the
 token from the session.

---
 src/docdb/getCosmosClient.ts               |  3 ++-
 src/docdb/tree/DocDBAccountTreeItemBase.ts | 10 ++++++++--
 src/docdb/utils/azureSessionHelper.ts      | 14 ++++++++++----
 src/tree/SubscriptionTreeItem.ts           |  3 ++-
 4 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/src/docdb/getCosmosClient.ts b/src/docdb/getCosmosClient.ts
index 711904fa..4a20baa2 100644
--- a/src/docdb/getCosmosClient.ts
+++ b/src/docdb/getCosmosClient.ts
@@ -20,6 +20,7 @@ export type CosmosDBKeyCredential = {
 
 export type CosmosDBAuthCredential = {
     type: 'auth';
+    tenantId: string;
 };
 
 export type CosmosDBCredential = CosmosDBKeyCredential | CosmosDBAuthCredential;
@@ -106,7 +107,7 @@ export function getCosmosClient(
             ...commonProperties,
             aadCredentials: {
                 getToken: async (scopes, _options) => {
-                    const session = await getSessionFromVSCode(scopes, undefined, { createIfNone: true });
+                    const session = await getSessionFromVSCode(scopes, authCred.tenantId, { createIfNone: true });
                     return {
                         token: session?.accessToken ?? '',
                         expiresOnTimestamp: 0,
diff --git a/src/docdb/tree/DocDBAccountTreeItemBase.ts b/src/docdb/tree/DocDBAccountTreeItemBase.ts
index 3c80426c..ec45d171 100644
--- a/src/docdb/tree/DocDBAccountTreeItemBase.ts
+++ b/src/docdb/tree/DocDBAccountTreeItemBase.ts
@@ -26,7 +26,12 @@ import { deleteCosmosDBAccount } from '../../commands/deleteDatabaseAccount/dele
 import { getThemeAgnosticIconPath, SERVERLESS_CAPABILITY_NAME } from '../../constants';
 import { nonNullProp } from '../../utils/nonNull';
 import { rejectOnTimeout } from '../../utils/timeout';
-import { getCosmosClient, getCosmosKeyCredential, type CosmosDBCredential } from '../getCosmosClient';
+import {
+    getCosmosAuthCredential,
+    getCosmosClient,
+    getCosmosKeyCredential,
+    type CosmosDBCredential,
+} from '../getCosmosClient';
 import { getSignedInPrincipalIdForAccountEndpoint } from '../utils/azureSessionHelper';
 import { ensureRbacPermission, isRbacException, showRbacPermissionError } from '../utils/rbacUtils';
 import { DocDBTreeItemBase } from './DocDBTreeItemBase';
@@ -131,8 +136,9 @@ export abstract class DocDBAccountTreeItemBase extends DocDBTreeItemBase<Databas
                     } catch (e) {
                         if (e instanceof Error && isRbacException(e) && !this.hasShownRbacNotification) {
                             this.hasShownRbacNotification = true;
+                            const tenantId = getCosmosAuthCredential(this.root.credentials)?.tenantId;
                             const principalId =
-                                (await getSignedInPrincipalIdForAccountEndpoint(this.root.endpoint)) ?? '';
+                                (await getSignedInPrincipalIdForAccountEndpoint(this.root.endpoint, tenantId)) ?? '';
                             // 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}]`) &&
diff --git a/src/docdb/utils/azureSessionHelper.ts b/src/docdb/utils/azureSessionHelper.ts
index 20e43789..e7b5626d 100644
--- a/src/docdb/utils/azureSessionHelper.ts
+++ b/src/docdb/utils/azureSessionHelper.ts
@@ -7,14 +7,20 @@
 import { getSessionFromVSCode } from '@microsoft/vscode-azext-azureauth/out/src/getSessionFromVSCode';
 import type * as vscode from 'vscode';
 
-export async function getSignedInPrincipalIdForAccountEndpoint(accountEndpoint: string): Promise<string | undefined> {
-    const session = await getSessionForDatabaseAccount(accountEndpoint);
+export async function getSignedInPrincipalIdForAccountEndpoint(
+    accountEndpoint: string,
+    tenantId: string | undefined,
+): Promise<string | undefined> {
+    const session = await getSessionForDatabaseAccount(accountEndpoint, tenantId);
     const principalId = session?.account.id.split('/')[1] ?? session?.account.id;
     return principalId;
 }
 
-async function getSessionForDatabaseAccount(endpoint: string): Promise<vscode.AuthenticationSession | undefined> {
+async function getSessionForDatabaseAccount(
+    endpoint: string,
+    tenantId: string | undefined,
+): Promise<vscode.AuthenticationSession | undefined> {
     const endpointUrl = new URL(endpoint);
     const scrope = `${endpointUrl.origin}${endpointUrl.pathname}.default`;
-    return await getSessionFromVSCode(scrope, undefined, { createIfNone: false });
+    return await getSessionFromVSCode(scrope, tenantId, { createIfNone: false });
 }
diff --git a/src/tree/SubscriptionTreeItem.ts b/src/tree/SubscriptionTreeItem.ts
index 9d70c58b..ebc17a57 100644
--- a/src/tree/SubscriptionTreeItem.ts
+++ b/src/tree/SubscriptionTreeItem.ts
@@ -252,8 +252,9 @@ export class SubscriptionTreeItem extends SubscriptionTreeItemBase {
                         }
                     }
 
+                    const tenantId = parent.subscription.tenantId ?? databaseAccount.identity?.tenantId;
                     // OAuth is always enabled for Cosmos DB and will be used as a fall back if key auth is unavailable
-                    const authCred = { type: 'auth' };
+                    const authCred = { type: 'auth', tenantId: tenantId };
                     const credentials = [keyCred, authCred].filter(
                         (cred): cred is CosmosDBCredential => cred !== undefined,
                     );