From a6473adf40434cf36dfab4197e645359aa4bdb2a Mon Sep 17 00:00:00 2001
From: Senthamil Sindhu <sindhuba@microsoft.com>
Date: Tue, 19 Nov 2024 20:56:54 -0800
Subject: [PATCH] Add recent changes

---
 src/Common/CosmosClient.ts        | 41 +++++++++++++++++++++++++++++--
 src/hooks/useDatabaseAccounts.tsx |  1 -
 src/hooks/useKnockoutExplorer.ts  |  6 +----
 3 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/src/Common/CosmosClient.ts b/src/Common/CosmosClient.ts
index f7a4fbfc5..13264e31a 100644
--- a/src/Common/CosmosClient.ts
+++ b/src/Common/CosmosClient.ts
@@ -8,11 +8,13 @@ import { AuthType } from "../AuthType";
 import { BackendApi, PriorityLevel } from "../Common/Constants";
 import * as Logger from "../Common/Logger";
 import { Platform, configContext } from "../ConfigContext";
-import { userContext } from "../UserContext";
+import { updateUserContext, userContext } from "../UserContext";
 import { logConsoleError } from "../Utils/NotificationConsoleUtils";
 import * as PriorityBasedExecutionUtils from "../Utils/PriorityBasedExecutionUtils";
 import { EmulatorMasterKey, HttpHeaders } from "./Constants";
 import { getErrorMessage } from "./ErrorHandlingUtils";
+import { runCommand } from "hooks/useDatabaseAccounts";
+import { acquireTokenWithMsal, getMsalInstance } from "Utils/AuthorizationUtils";
 
 const _global = typeof self === "undefined" ? window : self;
 
@@ -32,7 +34,42 @@ export const tokenProvider = async (requestInfo: Cosmos.RequestInfo) => {
       return null;
     }
     const AUTH_PREFIX = `type=aad&ver=1.0&sig=`;
-    const authorizationToken = `${AUTH_PREFIX}${userContext.aadToken}`;
+    let authorizationToken;
+
+    try {
+       authorizationToken = `${AUTH_PREFIX}${userContext.aadToken}`;
+    } catch (error) {
+      if (error.code === "ExpiredAuthenticationToken") {
+        // Renew the AAD token using runCommand
+        const newToken = await runCommand(async () => {
+          // Implement the logic to acquire a new AAD token
+          const msalInstance = await getMsalInstance();
+          const cachedAccount = msalInstance.getAllAccounts()?.[0];
+          const cachedTenantId = localStorage.getItem("cachedTenantId");
+  
+          msalInstance.setActiveAccount(cachedAccount);
+          const newAccessToken = await acquireTokenWithMsal(msalInstance, {
+            authority: `${configContext.AAD_ENDPOINT}${cachedTenantId}`,
+            scopes: [`${configContext.ARM_ENDPOINT}/.default`],
+          });
+  
+          // Update user context with the new token
+          updateUserContext({ aadToken: newAccessToken });
+          authorizationToken = `${AUTH_PREFIX}${userContext.aadToken}`;
+          return newAccessToken;
+        });
+  
+        // Retry getting the token after renewing
+        const retryResult = await getTokenFromAuthService(verb, resourceType, resourceId);
+        headers[HttpHeaders.msDate] = retryResult.XDate;
+        return decodeURIComponent(retryResult.PrimaryReadWriteToken);
+      } else {
+        console.error('An error occurred:', error.message);
+        throw error; 
+      }
+    }
+
+    
     return authorizationToken;
   }
 
diff --git a/src/hooks/useDatabaseAccounts.tsx b/src/hooks/useDatabaseAccounts.tsx
index 3a280728a..728e9cf29 100644
--- a/src/hooks/useDatabaseAccounts.tsx
+++ b/src/hooks/useDatabaseAccounts.tsx
@@ -136,7 +136,6 @@ interface QueryResponse {
   $skipToken?: string;
 }
 
-// Define a generic runCommand function
 export async function runCommand<T>(
   fn: (...args: any[]) => Promise<T>,
   ...args: any[]
diff --git a/src/hooks/useKnockoutExplorer.ts b/src/hooks/useKnockoutExplorer.ts
index 6de4ada5e..2cd550c41 100644
--- a/src/hooks/useKnockoutExplorer.ts
+++ b/src/hooks/useKnockoutExplorer.ts
@@ -51,7 +51,6 @@ import {
 } from "../Utils/AuthorizationUtils";
 import { isInvalidParentFrameOrigin, shouldProcessMessage } from "../Utils/MessageValidation";
 import { getReadOnlyKeys, listKeys } from "../Utils/arm/generatedClients/cosmos/databaseAccounts";
-import { getReadOnlyKeys, listKeys } from "../Utils/arm/generatedClients/cosmos/databaseAccounts";
 import { applyExplorerBindings } from "../applyExplorerBindings";
 
 // This hook will create a new instance of Explorer.ts and bind it to the DOM
@@ -463,9 +462,6 @@ function configureEmulator(): Explorer {
   return explorer;
 }
 
-export async function fetchAndUpdateKeys(subscriptionId: string, resourceGroup: string, account: string) {
-  Logger.logInfo(`Fetching keys for ${userContext.apiType} account ${account}`, "Explorer/fetchAndUpdateKeys");
-  let keys;
 export async function fetchAndUpdateKeys(subscriptionId: string, resourceGroup: string, account: string) {
   Logger.logInfo(`Fetching keys for ${userContext.apiType} account ${account}`, "Explorer/fetchAndUpdateKeys");
   let keys;
@@ -838,4 +834,4 @@ async function updateContextForSampleData(explorer: Explorer): Promise<void> {
 
 interface SampledataconnectionResponse {
   connectionString: string;
-}
+}
\ No newline at end of file