From 3b86a9477fe8257e5b9483eff2c848f05e321776 Mon Sep 17 00:00:00 2001
From: Senthamil Sindhu <sindhuba@microsoft.com>
Date: Tue, 23 Jul 2024 11:14:19 -0700
Subject: [PATCH] Add code for arm token refresh

---
 src/hooks/useDatabaseAccounts.tsx | 66 ++++++++++++++++++++++++++++++-
 1 file changed, 64 insertions(+), 2 deletions(-)

diff --git a/src/hooks/useDatabaseAccounts.tsx b/src/hooks/useDatabaseAccounts.tsx
index f517b2e30..4cdb36bca 100644
--- a/src/hooks/useDatabaseAccounts.tsx
+++ b/src/hooks/useDatabaseAccounts.tsx
@@ -3,6 +3,8 @@ import { QueryRequestOptions, QueryResponse } from "Contracts/AzureResourceGraph
 import useSWR from "swr";
 import { configContext } from "../ConfigContext";
 import { DatabaseAccount } from "../Contracts/DataModels";
+import { acquireTokenWithMsal, getMsalInstance } from "Utils/AuthorizationUtils";
+import React from "react";
 /* eslint-disable  @typescript-eslint/no-explicit-any */
 
 interface AccountListResult {
@@ -72,14 +74,24 @@ export async function fetchDatabaseAccountsFromGraph(
     });
 
     if (!response.ok) {
-      throw new Error(await response.text());
-    }
+     // throw new Error(await response.text());
+   // }
 
     const queryResponse: QueryResponse = (await response.json()) as QueryResponse;
     skipToken = queryResponse.$skipToken;
     queryResponse.data?.map((databaseAccount: any) => {
       databaseAccounts.push(databaseAccount as DatabaseAccount);
     });
+  } else {
+    try{
+      console.log("Token expired");
+      await acquireNewTokenAndRetry(body);
+    }
+    catch (error) {
+     throw new Error(error);
+    }
+  
+  }
   } while (skipToken);
 
   return databaseAccounts.sort((a, b) => a.name.localeCompare(b.name));
@@ -92,3 +104,53 @@ export function useDatabaseAccounts(subscriptionId: string, armToken: string): D
   );
   return data;
 }
+
+async function acquireNewTokenAndRetry(body: any) {
+  try {
+    const msalInstance = await getMsalInstance();
+    
+const cachedAccount = msalInstance.getAllAccounts()?.[0];
+const cachedTenantId = localStorage.getItem("cachedTenantId");
+
+   // const [tenantId, setTenantId] = React.useState<string>(cachedTenantId);
+
+  
+    msalInstance.setActiveAccount(cachedAccount);
+
+    const newAccessToken = await acquireTokenWithMsal(msalInstance, {
+      authority: `${configContext.AAD_ENDPOINT}${cachedTenantId}`,
+      scopes: [`${configContext.ARM_ENDPOINT}/.default`],
+    });
+    const newBearer = `Bearer ${newAccessToken}`;
+    const newHeaders = new Headers();
+    newHeaders.append("Authorization", newBearer);
+    newHeaders.append(HttpHeaders.contentType, "application/json");
+    const apiVersion = "2021-03-01";
+    const managementResourceGraphAPIURL = `${configContext.ARM_ENDPOINT}providers/Microsoft.ResourceGraph/resources?api-version=${apiVersion}`;
+    
+  const databaseAccounts: DatabaseAccount[] = [];
+  let skipToken: string;
+  
+
+    // Retry the request with the new token
+    const response = await fetch(managementResourceGraphAPIURL, {
+      method: "POST",
+      headers: newHeaders,
+      body: JSON.stringify(body),
+    });
+
+    if (response.ok) {
+      // Handle successful response with new token
+      const queryResponse: QueryResponse = await response.json();
+      skipToken = queryResponse.$skipToken;
+      queryResponse.data?.forEach((databaseAccount: any) => {
+        databaseAccounts.push(databaseAccount as DatabaseAccount);
+      });
+    } else {
+      throw new Error(`Failed to fetch data after acquiring new token. Status: ${response.status}, ${await response.text()}`);
+    }
+  } catch (error) {
+    console.error("Error acquiring new token and retrying:", error);
+    throw error;
+  }
+}