Enable azure cloud console in ADS (#8546)

* initial changes

* Enable cloud console

* Delete unnecessary code

* error handling

* error handling

* Deal with promises

* Fix externals

* Fix externals for ws

* Cleanup name of terminal

* Update yarn.lock

* Fix yarn.lock

* Fix externals

* Cleanup parts of the code

* Fix more issues

* Fix cloud terminal

* Go back to our client ID

* Fix message

* Respect preferred location

* Fix govt cloud

* Some more messaging

* Enable items on right click

* Some feedback

* Change to status message

extensions/azurecore/src/account-provider/auths/azureAuth.ts

@@ -124,6 +124,7 @@ export abstract class AzureAuth implements vscode.Disposable {
 			this.metadata.settings.sqlResource,
 			this.metadata.settings.graphResource,
 			this.metadata.settings.ossRdbmsResource,
+			this.metadata.settings.microsoftResource,
 			this.metadata.settings.azureKeyVaultResource
 		];
 
@@ -153,6 +154,7 @@ export abstract class AzureAuth implements vscode.Disposable {
 		try {
 			await this.refreshAccessToken(account.key, refreshToken);
 		} catch (ex) {
+			account.isStale = true;
 			if (ex.message) {
 				await vscode.window.showErrorMessage(ex.message);
 			}
@@ -163,6 +165,10 @@ export abstract class AzureAuth implements vscode.Disposable {
 
 
 	public async getSecurityToken(account: azdata.Account, azureResource: azdata.AzureResource): Promise<TokenResponse | undefined> {
+		if (account.isStale === true) {
+			return undefined;
+		}
+
 		const resource = this.resources.find(s => s.azureResourceId === azureResource);
 		if (!resource) {
 			return undefined;
@@ -199,8 +205,13 @@ export abstract class AzureAuth implements vscode.Disposable {
 				if (!baseToken) {
 					return undefined;
 				}
+				try {
+					await this.refreshAccessToken(account.key, baseToken.refreshToken, tenant, resource);
+				} catch (ex) {
+					account.isStale = true;
+					return undefined;
+				}
 
-				await this.refreshAccessToken(account.key, baseToken.refreshToken, tenant, resource);
 				cachedTokens = await this.getCachedToken(account.key, resource.id, tenant.id);
 				if (!cachedTokens) {
 					return undefined;
@@ -349,8 +360,7 @@ export abstract class AzureAuth implements vscode.Disposable {
 			return { accessToken, refreshToken, tokenClaims };
 
 		} catch (err) {
-			console.dir(err);
-			const msg = localize('azure.noToken', "Retrieving the token failed.");
+			const msg = localize('azure.noToken', "Retrieving the Azure token failed. Please sign in again.");
 			vscode.window.showErrorMessage(msg);
 			throw new Error(err);
 		}

extensions/azurecore/src/account-provider/auths/azureAuthCodeGrant.ts

@@ -69,7 +69,7 @@ export class AzureAuthCodeGrant extends AzureAuth {
 			serverPort = await this.server.startup();
 		} catch (err) {
 			const msg = localize('azure.serverCouldNotStart', 'Server could not start. This could be a permissions error or an incompatibility on your system. You can try enabling device code authentication from settings.');
-			await vscode.window.showErrorMessage(msg);
+			vscode.window.showErrorMessage(msg);
 			console.dir(err);
 			return undefined;
 		}
@@ -181,7 +181,7 @@ export class AzureAuthCodeGrant extends AzureAuth {
 			refreshToken = rt;
 		} catch (ex) {
 			if (ex.msg) {
-				await vscode.window.showErrorMessage(ex.msg);
+				vscode.window.showErrorMessage(ex.msg);
 			}
 			console.log(ex);
 		}
@@ -199,7 +199,7 @@ export class AzureAuthCodeGrant extends AzureAuth {
 		} catch (ex) {
 			console.log(ex);
 			if (ex.msg) {
-				await vscode.window.showErrorMessage(ex.msg);
+				vscode.window.showErrorMessage(ex.msg);
 				authCompleteDeferred.reject(ex);
 			} else {
 				authCompleteDeferred.reject(new Error('There was an issue when storing the cache.'));

extensions/azurecore/src/account-provider/azureAccountProvider.ts

@@ -128,7 +128,7 @@ export class AzureAccountProvider implements azdata.AccountProvider, vscode.Disp
 
 		if (this.authMappings.size === 0) {
 			console.log('No auth method was enabled.');
-			await vscode.window.showErrorMessage(noAuthAvailable);
+			vscode.window.showErrorMessage(noAuthAvailable);
 			return { canceled: true };
 		}
 
@@ -145,7 +145,7 @@ export class AzureAccountProvider implements azdata.AccountProvider, vscode.Disp
 
 		if (!pick) {
 			console.log('No auth method was selected.');
-			await vscode.window.showErrorMessage(noAuthSelected);
+			vscode.window.showErrorMessage(noAuthSelected);
 			return { canceled: true };
 		}
 

extensions/azurecore/src/account-provider/interfaces.ts

@@ -69,6 +69,11 @@ interface Settings {
 	 */
 	signInResourceId?: string;
 
+	/**
+	 * Information that describes the Microsoft resource management resource
+	 */
+	microsoftResource?: Resource
+
 	/**
 	 * Information that describes the AAD graph resource
 	 */

extensions/azurecore/src/account-provider/providerSettings.ts

@@ -18,6 +18,11 @@ const publicAzureSettings: ProviderSettings = {
 			host: 'https://login.microsoftonline.com/',
 			clientId: 'a69788c6-1d43-44ed-9ca3-b83e194da255',
 			signInResourceId: 'https://management.core.windows.net/',
+			microsoftResource: {
+				id: 'marm',
+				endpoint: 'https://management.core.windows.net/',
+				azureResourceId: AzureResource.MicrosoftResourceManagement
+			},
 			graphResource: {
 				id: 'graph',
 				endpoint: 'https://graph.microsoft.com',
@@ -62,6 +67,11 @@ const usGovAzureSettings: ProviderSettings = {
 			host: 'https://login.microsoftonline.us/',
 			clientId: 'a69788c6-1d43-44ed-9ca3-b83e194da255',
 			signInResourceId: 'https://management.core.usgovcloudapi.net/',
+			microsoftResource: {
+				id: 'marm',
+				endpoint: 'https://management.core.usgovcloudapi.net/',
+				azureResourceId: AzureResource.MicrosoftResourceManagement
+			},
 			graphResource: {
 				id: 'graph',
 				endpoint: 'https://graph.windows.net',

extensions/azurecore/src/azureResource/commands.ts

@@ -16,13 +16,58 @@ import { TreeNode } from './treeNode';
 import { AzureResourceCredentialError } from './errors';
 import { AzureResourceTreeProvider } from './tree/treeProvider';
 import { AzureResourceAccountTreeNode } from './tree/accountTreeNode';
-import { IAzureResourceSubscriptionService, IAzureResourceSubscriptionFilterService } from '../azureResource/interfaces';
+import { IAzureResourceSubscriptionService, IAzureResourceSubscriptionFilterService, IAzureTerminalService } from '../azureResource/interfaces';
 import { AzureResourceServiceNames } from './constants';
 import { AzureResourceGroupService } from './providers/resourceGroup/resourceGroupService';
 import { GetSubscriptionsResult, GetResourceGroupsResult } from '../azurecore';
 import { isArray } from 'util';
+import { AzureAccount, Tenant } from '../account-provider/interfaces';
 
 export function registerAzureResourceCommands(appContext: AppContext, tree: AzureResourceTreeProvider): void {
+	appContext.apiWrapper.registerCommand('azure.resource.startterminal', async (node?: TreeNode) => {
+		try {
+			if (!node || !(node instanceof AzureResourceAccountTreeNode)) {
+				return;
+			}
+
+			const accountNode = node as AzureResourceAccountTreeNode;
+			const azureAccount = accountNode.account as AzureAccount;
+
+			const tokens = await appContext.apiWrapper.getSecurityToken(azureAccount, azdata.AzureResource.MicrosoftResourceManagement);
+
+			const terminalService = appContext.getService<IAzureTerminalService>(AzureResourceServiceNames.terminalService);
+
+			const listOfTenants = azureAccount.properties.tenants.map(t => t.displayName);
+
+			if (listOfTenants.length === 0) {
+				window.showErrorMessage(localize('azure.noTenants', "A tenant is required for this feature. Your Azure subscription seems to have no tenants."));
+				return;
+			}
+
+			let tenant: Tenant;
+			window.setStatusBarMessage(localize('azure.startingCloudShell', "Starting cloud shell…"), 5000);
+
+			if (listOfTenants.length === 1) {
+				// Don't show quickpick for a single option
+				tenant = azureAccount.properties.tenants[0];
+			} else {
+				const pickedTenant = await window.showQuickPick(listOfTenants, { canPickMany: false });
+
+				if (!pickedTenant) {
+					window.showErrorMessage(localize('azure.mustPickTenant', "You must select a tenant for this feature to work."));
+					return;
+				}
+
+				// The tenant the user picked
+				tenant = azureAccount.properties.tenants[listOfTenants.indexOf(pickedTenant)];
+			}
+
+			await terminalService.getOrCreateCloudConsole(azureAccount, tenant, tokens);
+		} catch (ex) {
+			console.error(ex);
+			window.showErrorMessage(ex);
+		}
+	});
 
 	// Resource Management commands
 	appContext.apiWrapper.registerCommand('azure.accounts.getSubscriptions', async (account?: azdata.Account, ignoreErrors: boolean = false): Promise<GetSubscriptionsResult> => {
@@ -98,6 +143,7 @@ export function registerAzureResourceCommands(appContext: AppContext, tree: Azur
 	});
 
 	// Resource Tree commands
+
 	appContext.apiWrapper.registerCommand('azure.resource.selectsubscriptions', async (node?: TreeNode) => {
 		if (!(node instanceof AzureResourceAccountTreeNode)) {
 			return;

extensions/azurecore/src/azureResource/constants.ts

@@ -20,5 +20,6 @@ export enum AzureResourceServiceNames {
 	accountService = 'AzureResourceAccountService',
 	subscriptionService = 'AzureResourceSubscriptionService',
 	subscriptionFilterService = 'AzureResourceSubscriptionFilterService',
-	tenantService = 'AzureResourceTenantService'
+	tenantService = 'AzureResourceTenantService',
+	terminalService = 'AzureTerminalService',
 }

extensions/azurecore/src/azureResource/interfaces.ts

@@ -9,6 +9,7 @@ import { Account, DidChangeAccountsParams } from 'azdata';
 import { Event } from 'vscode';
 
 import { azureResource } from './azure-resource';
+import { AzureAccount, AzureAccountSecurityToken, Tenant } from '../account-provider/interfaces';
 
 export interface IAzureResourceAccountService {
 	getAccounts(): Promise<Account[]>;
@@ -24,6 +25,10 @@ export interface IAzureResourceSubscriptionFilterService {
 	saveSelectedSubscriptions(account: Account, selectedSubscriptions: azureResource.AzureResourceSubscription[]): Promise<void>;
 }
 
+export interface IAzureTerminalService {
+	getOrCreateCloudConsole(account: AzureAccount, tenant: Tenant, tokens: { [key: string]: AzureAccountSecurityToken }): Promise<void>;
+}
+
 export interface IAzureResourceCacheService {
 	generateKey(id: string): string;
 

extensions/azurecore/src/azureResource/services/terminalService.ts

@@ -0,0 +1,196 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as vscode from 'vscode';
+import * as nls from 'vscode-nls';
+import axios, { AxiosRequestConfig } from 'axios';
+import * as WS from 'ws';
+
+import { IAzureTerminalService } from '../interfaces';
+import { AzureAccount, AzureAccountSecurityToken, Tenant } from '../../account-provider/interfaces';
+
+const localize = nls.loadMessageBundle();
+export class AzureTerminalService implements IAzureTerminalService {
+	private readonly apiVersion = '?api-version=2018-10-01';
+
+	public constructor(context: vscode.ExtensionContext) {
+
+	}
+
+	public async getOrCreateCloudConsole(account: AzureAccount, tenant: Tenant, tokens: { [key: string]: AzureAccountSecurityToken }): Promise<void> {
+		const token = tokens[tenant.id].token;
+		const settings: AxiosRequestConfig = {
+			headers: {
+				'Accept': 'application/json',
+				'Content-Type': 'application/json',
+				'Authorization': `Bearer ${token}`
+			},
+			validateStatus: () => true
+		};
+
+		const metadata = account.properties.providerSettings;
+
+		const userSettingsUri = this.getConsoleUserSettingsUri(metadata.settings.armResource.endpoint);
+		const userSettingsResult = await axios.get(userSettingsUri, settings);
+
+		const preferredShell = userSettingsResult.data?.properties?.preferredShellType ?? 'bash';
+		const preferredLocation = userSettingsResult.data?.properties?.preferredLocation;
+
+		const consoleRequestUri = this.getConsoleRequestUri(metadata.settings.armResource.endpoint);
+		if (preferredLocation) {
+			settings.headers['x-ms-console-preferred-location'] = preferredLocation;
+		}
+
+		const provisionResult = await axios.put(consoleRequestUri, {}, settings);
+
+		if (provisionResult.data?.properties?.provisioningState !== 'Succeeded') {
+			throw new Error(provisionResult.data);
+		}
+		const consoleUri = provisionResult.data.properties.uri;
+
+		return this.createTerminal(consoleUri, token, account.displayInfo.displayName, preferredShell);
+	}
+
+
+	private async createTerminal(provisionedUri: string, token: string, accountDisplayName: string, preferredShell: string): Promise<void> {
+		class ShellType implements vscode.QuickPickItem {
+			constructor(public readonly label: string, public readonly value: string) {
+			}
+		}
+
+		const shells = [new ShellType('PowerShell', 'pwsh'), new ShellType('Bash', 'bash'),];
+		const idx = shells.findIndex(s => s.value === preferredShell);
+
+		const prefShell = shells.splice(idx, 1);
+		shells.unshift(prefShell[0]);
+
+		let shell = await vscode.window.showQuickPick(shells, {
+			canPickMany: false,
+			placeHolder: localize('azure.selectShellType', "Select Bash or PowerShell for Azure Cloud Shell")
+		});
+
+		if (!shell) {
+			vscode.window.showErrorMessage(localize('azure.shellTypeRequired', "You must pick a shell type"));
+			return;
+		}
+
+		const terminalName = localize('azure.cloudShell', "Azure Cloud Shell (Preview)") + ` ${shell} (${accountDisplayName})`;
+
+		const azureTerminal = new AzureTerminal(provisionedUri, token, shell.value);
+		const terminal = vscode.window.createTerminal({
+			name: terminalName,
+			pty: azureTerminal
+		});
+
+		terminal.show();
+	}
+
+	public getConsoleRequestUri(armEndpoint: string): string {
+		return `${armEndpoint}/providers/Microsoft.Portal/consoles/default${this.apiVersion}`;
+	}
+
+	public getConsoleUserSettingsUri(armEndpoint: string): string {
+		return `${armEndpoint}/providers/Microsoft.Portal/userSettings/cloudconsole${this.apiVersion}`;
+	}
+}
+
+class AzureTerminal implements vscode.Pseudoterminal {
+	private readonly writeEmitter: vscode.EventEmitter<string>;
+	public readonly onDidWrite: vscode.Event<string>;
+
+	private socket: WS;
+	private intervalTimer: NodeJS.Timer;
+	private terminalDimensions: vscode.TerminalDimensions;
+
+	constructor(private readonly consoleUri: string, private readonly token: string, private shell: string) {
+		this.writeEmitter = new vscode.EventEmitter<string>();
+		this.onDidWrite = this.writeEmitter.event;
+	}
+
+	handleInput(data: string): void {
+		this.socket?.send(data);
+	}
+
+	async open(initialDimensions: vscode.TerminalDimensions): Promise<void> {
+		return this.resetTerminalSize(initialDimensions);
+	}
+
+	close(): void {
+		if (!this.socket) { return; }
+
+		this.socket.removeAllListeners('open');
+		this.socket.removeAllListeners('message');
+		this.socket.removeAllListeners('close');
+
+		this.socket.terminate();
+		if (this.intervalTimer) {
+			clearInterval(this.intervalTimer);
+		}
+	}
+
+	async setDimensions(dimensions: vscode.TerminalDimensions): Promise<void> {
+		return this.resetTerminalSize(dimensions);
+	}
+
+	private async resetTerminalSize(dimensions: vscode.TerminalDimensions): Promise<void> {
+		try {
+
+			if (!this.terminalDimensions) { // first time
+				this.writeEmitter.fire(localize('azure.connectingShellTerminal', "Connecting terminal...\n"));
+			}
+
+			if (dimensions) {
+				this.terminalDimensions = dimensions;
+			}
+
+			// Close the shell before this and restablish a new connection
+			this.close();
+
+			const terminalUri = await this.establishTerminal(this.terminalDimensions);
+			this.socket = new WS(terminalUri);
+
+			this.socket.on('message', (data: WS.Data) => {
+				// Write to the console
+				this.writeEmitter.fire(data.toString());
+			});
+
+			this.socket.on('close', () => {
+				this.writeEmitter.fire(localize('azure.shellClosed', "Shell closed.\n"));
+				this.close();
+			});
+
+			// Keep alives
+			this.intervalTimer = setInterval(() => {
+				this.socket.ping();
+			}, 5000);
+		} catch (ex) {
+			console.log(ex);
+		}
+	}
+
+
+	private async establishTerminal(dimensions: vscode.TerminalDimensions): Promise<string> {
+		const terminalResult = await axios.post(`${this.consoleUri}/terminals?rows=${dimensions.rows}&cols=${dimensions.columns}&shell=${this.shell}`, undefined, {
+			headers: {
+				'Accept': 'application/json',
+				'Content-Type': 'application/json',
+				'Authorization': `Bearer ${this.token}`
+			}
+		});
+
+		const terminalUri = terminalResult.data?.socketUri;
+
+		if (terminalResult.data.error) {
+			vscode.window.showErrorMessage(terminalResult.data.error.message);
+		}
+
+		if (!terminalUri) {
+			console.log(terminalResult);
+			throw new Error(terminalResult.data);
+		}
+
+		return terminalUri;
+	}
+}

extensions/azurecore/src/extension.ts

@@ -17,7 +17,7 @@ import { AzureResourceDatabaseServerService } from './azureResource/providers/da
 import { AzureResourceDatabaseProvider } from './azureResource/providers/database/databaseProvider';
 import { AzureResourceDatabaseService } from './azureResource/providers/database/databaseService';
 import { AzureResourceService } from './azureResource/resourceService';
-import { IAzureResourceCacheService, IAzureResourceAccountService, IAzureResourceSubscriptionService, IAzureResourceSubscriptionFilterService, IAzureResourceTenantService } from './azureResource/interfaces';
+import { IAzureResourceCacheService, IAzureResourceAccountService, IAzureResourceSubscriptionService, IAzureResourceSubscriptionFilterService, IAzureResourceTenantService, IAzureTerminalService } from './azureResource/interfaces';
 import { AzureResourceServiceNames } from './azureResource/constants';
 import { AzureResourceAccountService } from './azureResource/services/accountService';
 import { AzureResourceSubscriptionService } from './azureResource/services/subscriptionService';
@@ -30,6 +30,7 @@ import { SqlInstanceResourceService } from './azureResource/providers/sqlinstanc
 import { SqlInstanceProvider } from './azureResource/providers/sqlinstance/sqlInstanceProvider';
 import { PostgresServerProvider } from './azureResource/providers/postgresServer/postgresServerProvider';
 import { PostgresServerService } from './azureResource/providers/postgresServer/postgresServerService';
+import { AzureTerminalService } from './azureResource/services/terminalService';
 import { SqlInstanceArcProvider } from './azureResource/providers/sqlinstanceArc/sqlInstanceArcProvider';
 import { SqlInstanceArcResourceService } from './azureResource/providers/sqlinstanceArc/sqlInstanceArcService';
 import { PostgresServerArcProvider } from './azureResource/providers/postgresArcServer/postgresServerProvider';
@@ -145,6 +146,7 @@ function registerAzureServices(appContext: AppContext): void {
 	appContext.registerService<IAzureResourceSubscriptionService>(AzureResourceServiceNames.subscriptionService, new AzureResourceSubscriptionService());
 	appContext.registerService<IAzureResourceSubscriptionFilterService>(AzureResourceServiceNames.subscriptionFilterService, new AzureResourceSubscriptionFilterService(new AzureResourceCacheService(extensionContext)));
 	appContext.registerService<IAzureResourceTenantService>(AzureResourceServiceNames.tenantService, new AzureResourceTenantService());
+	appContext.registerService<IAzureTerminalService>(AzureResourceServiceNames.terminalService, new AzureTerminalService(extensionContext));
 }
 
 async function onDidChangeConfiguration(e: vscode.ConfigurationChangeEvent, apiWrapper: ApiWrapper): Promise<void> {