From 2fda88177003e000f79f0b0215ec3e6b1d7bf23e Mon Sep 17 00:00:00 2001
From: Zachary Foster <zfoster@users.noreply.github.com>
Date: Mon, 24 May 2021 14:03:51 -0400
Subject: [PATCH] Use customer endpoint for RBAC AAD auth (#818)

---
 src/HostedExplorer.tsx  |  7 ++++---
 src/hooks/useAADAuth.ts | 20 ++++++++++++++++++++
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/src/HostedExplorer.tsx b/src/HostedExplorer.tsx
index 4751ea5fa..9ae06f844 100644
--- a/src/HostedExplorer.tsx
+++ b/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";
@@ -7,7 +7,7 @@ import "../less/hostedexplorer.less";
 import { AuthType } from "./AuthType";
 import { DatabaseAccount } from "./Contracts/DataModels";
 import "./Explorer/Menus/NavBar/MeControlComponent.less";
-import { useAADAuth } from "./hooks/useAADAuth";
+import { useAADAuth, useAADDataPlane } from "./hooks/useAADAuth";
 import { useTokenMetadata } from "./hooks/usePortalAccessToken";
 import { HostedExplorerChildFrame } from "./HostedExplorerChildFrame";
 import { AccountSwitcher } from "./Platform/Hosted/Components/AccountSwitcher";
@@ -31,8 +31,9 @@ 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 { aadToken } = useAADDataPlane(databaseAccount);
   const [authType, setAuthType] = React.useState<AuthType>(encryptedToken ? AuthType.EncryptedToken : undefined);
   const [connectionString, setConnectionString] = React.useState<string>();
 
diff --git a/src/hooks/useAADAuth.ts b/src/hooks/useAADAuth.ts
index 1c6e946b2..9defd7d96 100644
--- a/src/hooks/useAADAuth.ts
+++ b/src/hooks/useAADAuth.ts
@@ -1,6 +1,7 @@
 import * as msal from "@azure/msal-browser";
 import { useBoolean } from "@fluentui/react-hooks";
 import * as React from "react";
+import { DatabaseAccount } from "../Contracts/DataModels";
 
 const config: msal.Configuration = {
   cache: {
@@ -104,3 +105,22 @@ export function useAADAuth(): ReturnType {
     switchTenant,
   };
 }
+
+export function useAADDataPlane(databaseAccount: DatabaseAccount): { aadToken: string } {
+  const [aadToken, setAadToken] = React.useState<string>();
+
+  React.useEffect(() => {
+    if (databaseAccount?.properties?.documentEndpoint) {
+      const hrefEndpoint = new URL(databaseAccount.properties.documentEndpoint).href.replace(/\/$/, "/.default");
+      msalInstance
+        .acquireTokenSilent({
+          scopes: [hrefEndpoint],
+        })
+        .then((aadTokenResponse) => {
+          setAadToken(aadTokenResponse.accessToken);
+        });
+    }
+  }, [databaseAccount]);
+
+  return { aadToken };
+}