[Hosted] AAD implementation for item operations (#643)

src/Common/CosmosClient.ts

@@ -10,6 +10,13 @@ const _global = typeof self === "undefined" ? window : self;
 
 export const tokenProvider = async (requestInfo: RequestInfo) => {
   const { verb, resourceId, resourceType, headers } = requestInfo;
+
+  if (userContext.features.enableAadDataPlane && userContext.aadToken) {
+    const AUTH_PREFIX = `type=aad&ver=1.0&sig=`;
+    const authorizationToken = `${AUTH_PREFIX}${userContext.aadToken}`;
+    return authorizationToken;
+  }
+
   if (configContext.platform === Platform.Emulator) {
     // TODO This SDK method mutates the headers object. Find a better one or fix the SDK.
     await setAuthorizationTokenHeaderUsingMasterKey(verb, resourceId, resourceType, headers, EmulatorMasterKey);
@@ -76,7 +83,7 @@ export function client(): Cosmos.CosmosClient {
   if (_client) return _client;
   const options: Cosmos.CosmosClientOptions = {
     endpoint: endpoint() || "https://cosmos.azure.com", // CosmosClient gets upset if we pass a bad URL. This should never actually get called
-    key: userContext.masterKey,
+    ...(!userContext.features.enableAadDataPlane && { key: userContext.masterKey }),
     tokenProvider,
     connectionPolicy: {
       enableEndpointDiscovery: false,

src/Common/dataAccess/createDocument.ts

@@ -1,8 +1,8 @@
 import { CollectionBase } from "../../Contracts/ViewModels";
+import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import { client } from "../CosmosClient";
 import { getEntityName } from "../DocumentUtility";
 import { handleError } from "../ErrorHandlingUtils";
-import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 
 export const createDocument = async (collection: CollectionBase, newDocument: unknown): Promise<unknown> => {
   const entityName = getEntityName();

src/Common/dataAccess/queryDocuments.ts

@@ -1,6 +1,6 @@
-import { Queries } from "../Constants";
 import { FeedOptions, ItemDefinition, QueryIterator, Resource } from "@azure/cosmos";
 import { LocalStorageUtility, StorageKey } from "../../Shared/StorageUtility";
+import { Queries } from "../Constants";
 import { client } from "../CosmosClient";
 
 export const queryDocuments = (

src/Common/dataAccess/queryDocumentsPage.ts

@@ -1,8 +1,8 @@
 import { QueryResults } from "../../Contracts/ViewModels";
 import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
-import { MinimalQueryIterator, nextPage } from "../IteratorUtilities";
-import { handleError } from "../ErrorHandlingUtils";
 import { getEntityName } from "../DocumentUtility";
+import { handleError } from "../ErrorHandlingUtils";
+import { MinimalQueryIterator, nextPage } from "../IteratorUtilities";
 
 export const queryDocumentsPage = async (
   resourceName: string,

src/Contracts/ExplorerContracts.ts

@@ -1,6 +1,6 @@
-import * as Versions from "./Versions";
 import * as ActionContracts from "./ActionContracts";
 import * as Diagnostics from "./Diagnostics";
+import * as Versions from "./Versions";
 
 /**
  * Messaging types used with Data Explorer <-> Portal communication

src/HostedExplorer.tsx

@@ -5,20 +5,20 @@ import { render } from "react-dom";
 import ChevronRight from "../images/chevron-right.svg";
 import "../less/hostedexplorer.less";
 import { AuthType } from "./AuthType";
-import { ConnectExplorer } from "./Platform/Hosted/Components/ConnectExplorer";
 import { DatabaseAccount } from "./Contracts/DataModels";
-import { DirectoryPickerPanel } from "./Platform/Hosted/Components/DirectoryPickerPanel";
-import { AccountSwitcher } from "./Platform/Hosted/Components/AccountSwitcher";
 import "./Explorer/Menus/NavBar/MeControlComponent.less";
-import { useTokenMetadata } from "./hooks/usePortalAccessToken";
-import { MeControl } from "./Platform/Hosted/Components/MeControl";
-import "./Platform/Hosted/ConnectScreen.less";
-import "./Shared/appInsights";
-import { SignInButton } from "./Platform/Hosted/Components/SignInButton";
 import { useAADAuth } from "./hooks/useAADAuth";
-import { FeedbackCommandButton } from "./Platform/Hosted/Components/FeedbackCommandButton";
+import { useTokenMetadata } from "./hooks/usePortalAccessToken";
 import { HostedExplorerChildFrame } from "./HostedExplorerChildFrame";
+import { AccountSwitcher } from "./Platform/Hosted/Components/AccountSwitcher";
+import { ConnectExplorer } from "./Platform/Hosted/Components/ConnectExplorer";
+import { DirectoryPickerPanel } from "./Platform/Hosted/Components/DirectoryPickerPanel";
+import { FeedbackCommandButton } from "./Platform/Hosted/Components/FeedbackCommandButton";
+import { MeControl } from "./Platform/Hosted/Components/MeControl";
+import { SignInButton } from "./Platform/Hosted/Components/SignInButton";
+import "./Platform/Hosted/ConnectScreen.less";
 import { extractMasterKeyfromConnectionString } from "./Platform/Hosted/HostedUtils";
+import "./Shared/appInsights";
 
 initializeIcons();
 
@@ -31,7 +31,7 @@ const App: React.FunctionComponent = () => {
   // For showing/hiding panel
   const [isOpen, { setTrue: openPanel, setFalse: dismissPanel }] = useBoolean(false);
 
-  const { isLoggedIn, armToken, graphToken, account, tenantId, logout, login, switchTenant } = useAADAuth();
+  const { isLoggedIn, armToken, graphToken, aadToken, 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,6 +50,7 @@ const App: React.FunctionComponent = () => {
           authType: AuthType.AAD,
           databaseAccount,
           authorizationToken: armToken,
+          aadToken,
         };
       } else if (authType === AuthType.EncryptedToken) {
         frameWindow.hostedConfig = {

src/HostedExplorerChildFrame.ts

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

src/Platform/Hosted/extractFeatures.ts

@@ -13,6 +13,7 @@ export type Features = {
   readonly enableSpark: boolean;
   readonly enableTtl: boolean;
   readonly executeSproc: boolean;
+  readonly enableAadDataPlane: boolean;
   readonly hostedDataExplorer: boolean;
   readonly junoEndpoint?: string;
   readonly livyEndpoint?: string;
@@ -43,6 +44,7 @@ export function extractFeatures(given = new URLSearchParams(window.location.sear
   return {
     canExceedMaximumValue: "true" === get("canexceedmaximumvalue"),
     cosmosdb: "true" === get("cosmosdb"),
+    enableAadDataPlane: "true" === get("enableaaddataplane"),
     enableChangeFeedPolicy: "true" === get("enablechangefeedpolicy"),
     enableFixedCollectionWithSharedThroughput: "true" === get("enablefixedcollectionwithsharedthroughput"),
     enableKOPanel: "true" === get("enablekopanel"),

src/Terminal/index.ts

@@ -1,11 +1,11 @@
-import "@jupyterlab/terminal/style/index.css";
-import "./index.css";
 import { ServerConnection } from "@jupyterlab/services";
-import { JupyterLabAppFactory } from "./JupyterLabAppFactory";
+import "@jupyterlab/terminal/style/index.css";
+import { HttpHeaders, TerminalQueryParams } from "../Common/Constants";
 import { Action } from "../Shared/Telemetry/TelemetryConstants";
 import * as TelemetryProcessor from "../Shared/Telemetry/TelemetryProcessor";
 import { updateUserContext } from "../UserContext";
-import { HttpHeaders, TerminalQueryParams } from "../Common/Constants";
+import "./index.css";
+import { JupyterLabAppFactory } from "./JupyterLabAppFactory";
 
 const getUrlVars = (): { [key: string]: string } => {
   const vars: { [key: string]: string } = {};
@@ -50,7 +50,7 @@ const createServerSettings = (urlVars: { [key: string]: string }): ServerConnect
 const main = async (): Promise<void> => {
   const urlVars = getUrlVars();
 
-  // Initialize userContext. Currently only subscrptionId is required by TelemetryProcessor
+  // Initialize userContext. Currently only subscriptionId is required by TelemetryProcessor
   updateUserContext({
     subscriptionId: urlVars[TerminalQueryParams.SubscriptionId],
   });

src/UserContext.ts

@@ -11,6 +11,7 @@ interface UserContext {
   readonly resourceGroup?: string;
   readonly databaseAccount?: DatabaseAccount;
   readonly endpoint?: string;
+  readonly aadToken?: string;
   readonly accessToken?: string;
   readonly authorizationToken?: string;
   readonly resourceToken?: string;

src/hooks/useAADAuth.ts

@@ -25,6 +25,7 @@ interface ReturnType {
   isLoggedIn: boolean;
   graphToken: string;
   armToken: string;
+  aadToken: string;
   login: () => void;
   logout: () => void;
   tenantId: string;
@@ -40,6 +41,7 @@ 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 () => {
@@ -79,9 +81,13 @@ export function useAADAuth(): ReturnType {
           authority: `https://login.microsoftonline.com/${tenantId}`,
           scopes: ["https://management.azure.com//.default"],
         }),
-      ]).then(([graphTokenResponse, armTokenResponse]) => {
+        msalInstance.acquireTokenSilent({
+          scopes: ["https://cosmos.azure.com/.default"],
+        }),
+      ]).then(([graphTokenResponse, armTokenResponse, aadTokenResponse]) => {
         setGraphToken(graphTokenResponse.accessToken);
         setArmToken(armTokenResponse.accessToken);
+        setAadToken(aadTokenResponse.accessToken);
       });
     }
   }, [account, tenantId]);
@@ -92,6 +98,7 @@ export function useAADAuth(): ReturnType {
     isLoggedIn,
     graphToken,
     armToken,
+    aadToken,
     login,
     logout,
     switchTenant,

src/hooks/useKnockoutExplorer.ts

@@ -83,6 +83,7 @@ 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;