Redo user endpoint dynamic token (#827)

* Redo user endpoint dynamic token

* Fixes aad endpoint race condition, tenant switching, and account permissions

* Export const msalInstance

* Format

* fix import

* format

* Redo getMsalInstance

* format again

* Check for doc endpoint

src/HostedExplorer.tsx

@@ -1,5 +1,5 @@
-import { useBoolean } from "@fluentui/react-hooks";
 import { initializeIcons } from "@fluentui/react";
+import { useBoolean } from "@fluentui/react-hooks";
 import * as React from "react";
 import { render } from "react-dom";
 import ChevronRight from "../images/chevron-right.svg";
@@ -31,7 +31,7 @@ const App: React.FunctionComponent = () => {
   // For showing/hiding panel
   const [isOpen, { setTrue: openPanel, setFalse: dismissPanel }] = useBoolean(false);
 
-  const { isLoggedIn, armToken, graphToken, aadToken, account, tenantId, logout, login, switchTenant } = useAADAuth();
+  const { isLoggedIn, armToken, graphToken, account, tenantId, logout, login, switchTenant } = useAADAuth();
   const [databaseAccount, setDatabaseAccount] = React.useState<DatabaseAccount>();
   const [authType, setAuthType] = React.useState<AuthType>(encryptedToken ? AuthType.EncryptedToken : undefined);
   const [connectionString, setConnectionString] = React.useState<string>();
@@ -50,7 +50,6 @@ const App: React.FunctionComponent = () => {
           authType: AuthType.AAD,
           databaseAccount,
           authorizationToken: armToken,
-          aadToken,
         };
       } else if (authType === AuthType.EncryptedToken) {
         frameWindow.hostedConfig = {

src/HostedExplorerChildFrame.ts

@@ -7,7 +7,6 @@ export interface HostedExplorerChildFrame extends Window {
 }
 
 export interface AAD {
-  aadToken: string;
   authType: AuthType.AAD;
   databaseAccount: DatabaseAccount;
   authorizationToken: string;

src/Utils/AuthorizationUtils.ts

@@ -1,3 +1,4 @@
+import * as msal from "@azure/msal-browser";
 import { AuthType } from "../AuthType";
 import * as Constants from "../Common/Constants";
 import * as Logger from "../Common/Logger";
@@ -40,3 +41,21 @@ export function decryptJWTToken(token: string) {
 
   return JSON.parse(tokenPayload);
 }
+
+export function getMsalInstance() {
+  const config: msal.Configuration = {
+    cache: {
+      cacheLocation: "localStorage",
+    },
+    auth: {
+      authority: "https://login.microsoftonline.com/common",
+      clientId: "203f1145-856a-4232-83d4-a43568fba23d",
+    },
+  };
+
+  if (process.env.NODE_ENV === "development") {
+    config.auth.redirectUri = "https://dataexplorer-dev.azurewebsites.net";
+  }
+  const msalInstance = new msal.PublicClientApplication(config);
+  return msalInstance;
+}

src/hooks/useAADAuth.ts

@@ -1,22 +1,9 @@
 import * as msal from "@azure/msal-browser";
 import { useBoolean } from "@fluentui/react-hooks";
 import * as React from "react";
+import { getMsalInstance } from "../Utils/AuthorizationUtils";
 
-const config: msal.Configuration = {
-  cache: {
-    cacheLocation: "localStorage",
-  },
-  auth: {
-    authority: "https://login.microsoftonline.com/common",
-    clientId: "203f1145-856a-4232-83d4-a43568fba23d",
-  },
-};
-
-if (process.env.NODE_ENV === "development") {
-  config.auth.redirectUri = "https://dataexplorer-dev.azurewebsites.net";
-}
-
-const msalInstance = new msal.PublicClientApplication(config);
+const msalInstance = getMsalInstance();
 
 const cachedAccount = msalInstance.getAllAccounts()?.[0];
 const cachedTenantId = localStorage.getItem("cachedTenantId");
@@ -25,7 +12,6 @@ interface ReturnType {
   isLoggedIn: boolean;
   graphToken: string;
   armToken: string;
-  aadToken: string;
   login: () => void;
   logout: () => void;
   tenantId: string;
@@ -41,7 +27,6 @@ export function useAADAuth(): ReturnType {
   const [tenantId, setTenantId] = React.useState<string>(cachedTenantId);
   const [graphToken, setGraphToken] = React.useState<string>();
   const [armToken, setArmToken] = React.useState<string>();
-  const [aadToken, setAadToken] = React.useState<string>();
 
   msalInstance.setActiveAccount(account);
   const login = React.useCallback(async () => {
@@ -81,13 +66,9 @@ export function useAADAuth(): ReturnType {
           authority: `https://login.microsoftonline.com/${tenantId}`,
           scopes: ["https://management.azure.com//.default"],
         }),
-        msalInstance.acquireTokenSilent({
-          scopes: ["https://cosmos.azure.com/.default"],
-        }),
-      ]).then(([graphTokenResponse, armTokenResponse, aadTokenResponse]) => {
+      ]).then(([graphTokenResponse, armTokenResponse]) => {
         setGraphToken(graphTokenResponse.accessToken);
         setArmToken(armTokenResponse.accessToken);
-        setAadToken(aadTokenResponse.accessToken);
       });
     }
   }, [account, tenantId]);
@@ -98,7 +79,6 @@ export function useAADAuth(): ReturnType {
     isLoggedIn,
     graphToken,
     armToken,
-    aadToken,
     login,
     logout,
     switchTenant,

src/hooks/useKnockoutExplorer.ts

@@ -28,11 +28,13 @@ import { CollectionCreation } from "../Shared/Constants";
 import { DefaultExperienceUtility } from "../Shared/DefaultExperienceUtility";
 import { PortalEnv, updateUserContext, userContext } from "../UserContext";
 import { listKeys } from "../Utils/arm/generatedClients/cosmos/databaseAccounts";
+import { DatabaseAccountListKeysResult } from "../Utils/arm/generatedClients/cosmos/types";
+import { getMsalInstance } from "../Utils/AuthorizationUtils";
 import { isInvalidParentFrameOrigin } from "../Utils/MessageValidation";
 
 // This hook will create a new instance of Explorer.ts and bind it to the DOM
 // This hook has a LOT of magic, but ideally we can delete it once we have removed KO and switched entirely to React
-// Pleas tread carefully :)
+// Please tread carefully :)
 
 export function useKnockoutExplorer(platform: Platform, explorerParams: ExplorerParams): Explorer {
   const [explorer, setExplorer] = useState<Explorer>();
@@ -83,16 +85,33 @@ async function configureHostedWithAAD(config: AAD, explorerParams: ExplorerParam
   updateUserContext({
     authType: AuthType.AAD,
     authorizationToken: `Bearer ${config.authorizationToken}`,
-    aadToken: config.aadToken,
   });
   const account = config.databaseAccount;
   const accountResourceId = account.id;
   const subscriptionId = accountResourceId && accountResourceId.split("subscriptions/")[1].split("/")[0];
   const resourceGroup = accountResourceId && accountResourceId.split("resourceGroups/")[1].split("/")[0];
-  const keys = await listKeys(subscriptionId, resourceGroup, account.name);
+  let aadToken;
+  let keys: DatabaseAccountListKeysResult = {};
+  if (account.properties?.documentEndpoint) {
+    const hrefEndpoint = new URL(account.properties.documentEndpoint).href.replace(/\/$/, "/.default");
+    const msalInstance = getMsalInstance();
+    const cachedAccount = msalInstance.getAllAccounts()?.[0];
+    msalInstance.setActiveAccount(cachedAccount);
+    const aadTokenResponse = await msalInstance.acquireTokenSilent({
+      forceRefresh: true,
+      scopes: [hrefEndpoint],
+    });
+    aadToken = aadTokenResponse.accessToken;
+  }
+  try {
+    keys = await listKeys(subscriptionId, resourceGroup, account.name);
+  } catch (e) {
+    console.warn(e);
+  }
   updateUserContext({
     subscriptionId,
     resourceGroup,
+    aadToken,
     databaseAccount: config.databaseAccount,
     masterKey: keys.primaryMasterKey,
   });