Support sovereign/custom clouds in microsoft authentication provider (#178725)

extensions/microsoft-authentication/README.md

@@ -5,3 +5,5 @@
 ## Features
 
 This extension provides support for authenticating to Microsoft. It registers the `microsoft` Authentication Provider that can be leveraged by other extensions. This also provides the Microsoft authentication used by Settings Sync.
+
+Additionally, it provides the `microsoft-sovereign-cloud` Authentication Provider that can be used to sign in to other Azure clouds like Azure for US Government or Azure China. Use the setting `microsoft-sovereign-cloud.endpoint` to select the authentication endpoint the provider should use. Please note that different scopes may also be required in different environments.

extensions/microsoft-authentication/package.json

@@ -31,6 +31,32 @@
       {
         "label": "Microsoft",
         "id": "microsoft"
+      },
+      {
+        "label": "Microsoft Sovereign Cloud",
+        "id": "microsoft-sovereign-cloud"
+      }
+    ],
+    "configuration": [
+      {
+        "title": "Microsoft Sovereign Cloud",
+        "properties": {
+          "microsoft-sovereign-cloud.endpoint": {
+            "anyOf": [
+              {
+                "type": "string"
+              },
+              {
+                "type": "string",
+                "enum": [
+                  "Azure China",
+                  "Azure US Government"
+                ]
+              }
+            ],
+            "description": "%microsoft-sovereign-cloud.endpoint.description%"
+          }
+        }
       }
     ]
   },

extensions/microsoft-authentication/package.nls.json

@@ -2,5 +2,6 @@
 	"displayName": "Microsoft Account",
 	"description": "Microsoft authentication provider",
 	"signIn": "Sign In",
-	"signOut": "Sign Out"
+	"signOut": "Sign Out",
+	"microsoft-sovereign-cloud.endpoint.description": "Login endpoint for Azure authentication. Select a national cloud or enter the login URL for a custom Azure cloud."
 }

extensions/microsoft-authentication/src/AADHelper.ts

@@ -13,9 +13,10 @@ import { BetterTokenStorage, IDidChangeInOtherWindowEvent } from './betterSecret
 import { LoopbackAuthServer } from './node/authServer';
 import { base64Decode } from './node/buffer';
 import { fetching } from './node/fetch';
+import { UriEventHandler } from './UriEventHandler';
 
 const redirectUrl = 'https://vscode.dev/redirect';
-const loginEndpointUrl = 'https://login.microsoftonline.com/';
+const defaultLoginEndpointUrl = 'https://login.microsoftonline.com/';
 const DEFAULT_CLIENT_ID = 'aebc6443-996d-45c2-90f0-388ff96faa56';
 const DEFAULT_TENANT = 'organizations';
 
@@ -35,7 +36,7 @@ interface IToken {
 	sessionId: string; // The account id + the scope
 }
 
-interface IStoredSession {
+export interface IStoredSession {
 	id: string;
 	refreshToken: string;
 	scope: string; // Scopes are alphabetized and joined with a space
@@ -44,6 +45,7 @@ interface IStoredSession {
 		displayName?: string;
 		id: string;
 	};
+	endpoint: string | undefined;
 }
 
 export interface ITokenResponse {
@@ -70,16 +72,8 @@ interface IScopeData {
 	tenant: string;
 }
 
-export const onDidChangeSessions = new vscode.EventEmitter<vscode.AuthenticationProviderAuthenticationSessionsChangeEvent>();
-
 export const REFRESH_NETWORK_FAILURE = 'Network failure';
 
-class UriEventHandler extends vscode.EventEmitter<vscode.Uri> implements vscode.UriHandler {
-	public handleUri(uri: vscode.Uri) {
-		this.fire(uri);
-	}
-}
-
 export class AzureActiveDirectoryService {
 	// For details on why this is set to 2/3... see https://github.com/microsoft/vscode/issues/133201#issuecomment-966668197
 	private static REFRESH_TIMEOUT_MODIFIER = 1000 * 2 / 3;
@@ -87,25 +81,25 @@ export class AzureActiveDirectoryService {
 	private _tokens: IToken[] = [];
 	private _refreshTimeouts: Map<string, NodeJS.Timeout> = new Map<string, NodeJS.Timeout>();
 	private _refreshingPromise: Promise<any> | undefined;
-	private _uriHandler: UriEventHandler;
+	private _sessionChangeEmitter: vscode.EventEmitter<vscode.AuthenticationProviderAuthenticationSessionsChangeEvent> = new vscode.EventEmitter<vscode.AuthenticationProviderAuthenticationSessionsChangeEvent>();
 
 	// Used to keep track of current requests when not using the local server approach.
 	private _pendingNonces = new Map<string, string[]>();
 	private _codeExchangePromises = new Map<string, Promise<vscode.AuthenticationSession>>();
 	private _codeVerfifiers = new Map<string, string>();
 
-	private readonly _tokenStorage: BetterTokenStorage<IStoredSession>;
-
-	constructor(private _context: vscode.ExtensionContext) {
-		this._tokenStorage = new BetterTokenStorage('microsoft.login.keylist', _context);
-		this._uriHandler = new UriEventHandler();
-		_context.subscriptions.push(vscode.window.registerUriHandler(this._uriHandler));
+	constructor(
+		private readonly _context: vscode.ExtensionContext,
+		private readonly _uriHandler: UriEventHandler,
+		private readonly _tokenStorage: BetterTokenStorage<IStoredSession>,
+		private readonly _loginEndpointUrl: string = defaultLoginEndpointUrl
+	) {
 		_context.subscriptions.push(this._tokenStorage.onDidChangeInOtherWindow((e) => this.checkForUpdates(e)));
 	}
 
 	public async initialize(): Promise<void> {
 		Logger.info('Reading sessions from secret storage...');
-		const sessions = await this._tokenStorage.getAll();
+		const sessions = await this._tokenStorage.getAll(item => this.sessionMatchesEndpoint(item));
 		Logger.info(`Got ${sessions.length} stored sessions`);
 
 		const refreshes = sessions.map(async session => {
@@ -162,6 +156,10 @@ export class AzureActiveDirectoryService {
 
 	//#region session operations
 
+	public get onDidChangeSessions(): vscode.Event<vscode.AuthenticationProviderAuthenticationSessionsChangeEvent> {
+		return this._sessionChangeEmitter.event;
+	}
+
 	async getSessions(scopes?: string[]): Promise<vscode.AuthenticationSession[]> {
 		if (!scopes) {
 			Logger.info('Getting sessions for all scopes...');
@@ -246,7 +244,7 @@ export class AzureActiveDirectoryService {
 		return Promise.all(matchingTokens.map(token => this.convertToSession(token, scopeData)));
 	}
 
-	public createSession(scopes: string[]): Promise<vscode.AuthenticationSession> {
+	public async createSession(scopes: string[]): Promise<vscode.AuthenticationSession> {
 		let modifiedScopes = [...scopes];
 		if (!modifiedScopes.includes('openid')) {
 			modifiedScopes.push('openid');
@@ -275,12 +273,19 @@ export class AzureActiveDirectoryService {
 
 		const runsRemote = vscode.env.remoteName !== undefined;
 		const runsServerless = vscode.env.remoteName === undefined && vscode.env.uiKind === vscode.UIKind.Web;
+
+		if (runsServerless && this._loginEndpointUrl !== defaultLoginEndpointUrl) {
+			throw new Error('Sign in to non-public Azure clouds is not supported on the web.');
+		}
+
 		if (runsRemote || runsServerless) {
 			return this.createSessionWithoutLocalServer(scopeData);
 		}
 
 		try {
-			return this.createSessionWithLocalServer(scopeData);
+			const session = await this.createSessionWithLocalServer(scopeData);
+			this._sessionChangeEmitter.fire({ added: [session], removed: [], changed: [] });
+			return session;
 		} catch (e) {
 			Logger.error(`Error creating session for scopes: ${scopeData.scopeStr} Error: ${e}`);
 
@@ -306,7 +311,7 @@ export class AzureActiveDirectoryService {
 			code_challenge_method: 'S256',
 			code_challenge: codeChallenge,
 		}).toString();
-		const loginUrl = `${loginEndpointUrl}${scopeData.tenant}/oauth2/v2.0/authorize?${qs}`;
+		const loginUrl = `${this._loginEndpointUrl}${scopeData.tenant}/oauth2/v2.0/authorize?${qs}`;
 		const server = new LoopbackAuthServer(path.join(__dirname, '../media'), loginUrl);
 		await server.start();
 
@@ -336,7 +341,7 @@ export class AzureActiveDirectoryService {
 		const state = encodeURIComponent(callbackUri.toString(true));
 		const codeVerifier = generateCodeVerifier();
 		const codeChallenge = await generateCodeChallenge(codeVerifier);
-		const signInUrl = `${loginEndpointUrl}${scopeData.tenant}/oauth2/v2.0/authorize`;
+		const signInUrl = `${this._loginEndpointUrl}${scopeData.tenant}/oauth2/v2.0/authorize`;
 		const oauthStartQuery = new URLSearchParams({
 			response_type: 'code',
 			client_id: encodeURIComponent(scopeData.clientId),
@@ -386,7 +391,7 @@ export class AzureActiveDirectoryService {
 			});
 	}
 
-	public removeSessionById(sessionId: string, writeToDisk: boolean = true): Promise<vscode.AuthenticationSession | undefined> {
+	public async removeSessionById(sessionId: string, writeToDisk: boolean = true): Promise<vscode.AuthenticationSession | undefined> {
 		Logger.info(`Logging out of session '${sessionId}'`);
 		const tokenIndex = this._tokens.findIndex(token => token.sessionId === sessionId);
 		if (tokenIndex === -1) {
@@ -395,13 +400,19 @@ export class AzureActiveDirectoryService {
 		}
 
 		const token = this._tokens.splice(tokenIndex, 1)[0];
-		return this.removeSessionByIToken(token, writeToDisk);
+		const session = await this.removeSessionByIToken(token, writeToDisk);
+
+		if (session) {
+			this._sessionChangeEmitter.fire({ added: [], removed: [session], changed: [] });
+		}
+
+		return session;
 	}
 
 	public async clearSessions() {
 		Logger.info('Logging out of all sessions');
 		this._tokens = [];
-		await this._tokenStorage.deleteAll();
+		await this._tokenStorage.deleteAll(item => this.sessionMatchesEndpoint(item));
 
 		this._refreshTimeouts.forEach(timeout => {
 			clearTimeout(timeout);
@@ -424,7 +435,7 @@ export class AzureActiveDirectoryService {
 
 		const session = this.convertToSessionSync(token);
 		Logger.info(`Sending change event for session that was removed with scopes: ${token.scope}`);
-		onDidChangeSessions.fire({ added: [], removed: [session], changed: [] });
+		this._sessionChangeEmitter.fire({ added: [], removed: [session], changed: [] });
 		Logger.info(`Logged out of session '${token.sessionId}' with scopes: ${token.scope}`);
 		return session;
 	}
@@ -439,7 +450,7 @@ export class AzureActiveDirectoryService {
 			try {
 				const refreshedToken = await this.refreshToken(refreshToken, scopeData, sessionId);
 				Logger.info('Triggering change session event...');
-				onDidChangeSessions.fire({ added: [], removed: [], changed: [this.convertToSessionSync(refreshedToken)] });
+				this._sessionChangeEmitter.fire({ added: [], removed: [], changed: [this.convertToSessionSync(refreshedToken)] });
 			} catch (e) {
 				if (e.message !== REFRESH_NETWORK_FAILURE) {
 					vscode.window.showErrorMessage(vscode.l10n.t('You have been signed out because reading stored authentication information failed.'));
@@ -570,7 +581,7 @@ export class AzureActiveDirectoryService {
 		});
 
 		const proxyEndpoints: { [providerId: string]: string } | undefined = await vscode.commands.executeCommand('workbench.getCodeExchangeProxyEndpoints');
-		const endpointUrl = proxyEndpoints?.microsoft || loginEndpointUrl;
+		const endpointUrl = proxyEndpoints?.microsoft || this._loginEndpointUrl;
 		const endpoint = `${endpointUrl}${scopeData.tenant}/oauth2/v2.0/token`;
 
 		try {
@@ -705,8 +716,16 @@ export class AzureActiveDirectoryService {
 				redirect_uri: redirectUrl
 			});
 
-			const proxyEndpoints: { [providerId: string]: string } | undefined = await vscode.commands.executeCommand('workbench.getCodeExchangeProxyEndpoints');
-			const endpointUrl = proxyEndpoints?.microsoft || loginEndpointUrl;
+			let endpointUrl: string;
+
+			if (this._loginEndpointUrl !== defaultLoginEndpointUrl) {
+				// If this is for sovereign clouds, don't try using the proxy endpoint, which supports only public cloud
+				endpointUrl = this._loginEndpointUrl;
+			} else {
+				const proxyEndpoints: { [providerId: string]: string } | undefined = await vscode.commands.executeCommand('workbench.getCodeExchangeProxyEndpoints');
+				endpointUrl = proxyEndpoints?.microsoft || this._loginEndpointUrl;
+			}
+
 			const endpoint = `${endpointUrl}${scopeData.tenant}/oauth2/v2.0/token`;
 
 			const json = await this.fetchTokenResponse(endpoint, postData, scopeData);
@@ -821,7 +840,8 @@ export class AzureActiveDirectoryService {
 			id: token.sessionId,
 			refreshToken: token.refreshToken,
 			scope: token.scope,
-			account: token.account
+			account: token.account,
+			endpoint: this._loginEndpointUrl,
 		});
 		Logger.info(`Stored token for scopes: ${scopeData.scopeStr}`);
 	}
@@ -833,6 +853,12 @@ export class AzureActiveDirectoryService {
 				Logger.error('session not found that was apparently just added');
 				return;
 			}
+
+			if (!this.sessionMatchesEndpoint(session)) {
+				// If the session wasn't made for this login endpoint, ignore this update
+				continue;
+			}
+
 			const matchesExisting = this._tokens.some(token => token.scope === session.scope && token.sessionId === session.id);
 			if (!matchesExisting && session.refreshToken) {
 				try {
@@ -848,7 +874,7 @@ export class AzureActiveDirectoryService {
 					Logger.info(`Session added in another window with scopes: ${session.scope}`);
 					const token = await this.refreshToken(session.refreshToken, scopeData, session.id);
 					Logger.info(`Sending change event for session that was added with scopes: ${scopeData.scopeStr}`);
-					onDidChangeSessions.fire({ added: [this.convertToSessionSync(token)], removed: [], changed: [] });
+					this._sessionChangeEmitter.fire({ added: [this.convertToSessionSync(token)], removed: [], changed: [] });
 					return;
 				} catch (e) {
 					// Network failures will automatically retry on next poll.
@@ -862,6 +888,11 @@ export class AzureActiveDirectoryService {
 		}
 
 		for (const { value } of e.removed) {
+			if (!this.sessionMatchesEndpoint(value)) {
+				// If the session wasn't made for this login endpoint, ignore this update
+				continue;
+			}
+
 			Logger.info(`Session removed in another window with scopes: ${value.scope}`);
 			await this.removeSessionById(value.id, false);
 		}
@@ -872,6 +903,13 @@ export class AzureActiveDirectoryService {
 		// are already managing (see usages of `setSessionTimeout`).
 	}
 
+	private sessionMatchesEndpoint(session: IStoredSession): boolean {
+		// For older sessions with no endpoint set, it can be assumed to be the default endpoint
+		session.endpoint ||= defaultLoginEndpointUrl;
+
+		return session.endpoint === this._loginEndpointUrl;
+	}
+
 	//#endregion
 
 	//#region static methods

extensions/microsoft-authentication/src/UriEventHandler.ts

@@ -0,0 +1,12 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as vscode from 'vscode';
+
+export class UriEventHandler extends vscode.EventEmitter<vscode.Uri> implements vscode.UriHandler {
+	public handleUri(uri: vscode.Uri) {
+		this.fire(uri);
+	}
+}

extensions/microsoft-authentication/src/betterSecretStorage.ts

@@ -81,11 +81,13 @@ export class BetterTokenStorage<T> {
 		return tokens.get(key);
 	}
 
-	async getAll(): Promise<T[]> {
+	async getAll(predicate?: (item: T) => boolean): Promise<T[]> {
 		const tokens = await this.getTokens();
 		const values = new Array<T>();
 		for (const [_, value] of tokens) {
-			values.push(value);
+			if (!predicate || predicate(value)) {
+				values.push(value);
+			}
 		}
 		return values;
 	}
@@ -141,11 +143,13 @@ export class BetterTokenStorage<T> {
 		this._operationInProgress = false;
 	}
 
-	async deleteAll(): Promise<void> {
+	async deleteAll(predicate?: (item: T) => boolean): Promise<void> {
 		const tokens = await this.getTokens();
 		const promises = [];
-		for (const [key] of tokens) {
-			promises.push(this.delete(key));
+		for (const [key, value] of tokens) {
+			if (!predicate || predicate(value)) {
+				promises.push(this.delete(key));
+			}
 		}
 		await Promise.all(promises);
 	}

extensions/microsoft-authentication/src/extension.ts

@@ -4,18 +4,105 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as vscode from 'vscode';
-import { AzureActiveDirectoryService, onDidChangeSessions } from './AADHelper';
+import { AzureActiveDirectoryService, IStoredSession } from './AADHelper';
+import { BetterTokenStorage } from './betterSecretStorage';
+import { UriEventHandler } from './UriEventHandler';
 import TelemetryReporter from '@vscode/extension-telemetry';
 
+async function initAzureCloudAuthProvider(context: vscode.ExtensionContext, telemetryReporter: TelemetryReporter, uriHandler: UriEventHandler, tokenStorage: BetterTokenStorage<IStoredSession>): Promise<vscode.Disposable | undefined> {
+	let settingValue = vscode.workspace.getConfiguration('microsoft-sovereign-cloud').get<string | undefined>('endpoint');
+	let authProviderName: string | undefined;
+	if (!settingValue) {
+		return undefined;
+	} else if (settingValue === 'Azure China') {
+		authProviderName = settingValue;
+		settingValue = 'https://login.chinacloudapi.cn/';
+	} else if (settingValue === 'Azure US Government') {
+		authProviderName = settingValue;
+		settingValue = 'https://login.microsoftonline.us/';
+	}
+
+	// validate user value
+	let uri: vscode.Uri;
+	try {
+		uri = vscode.Uri.parse(settingValue, true);
+	} catch (e) {
+		vscode.window.showErrorMessage(vscode.l10n.t('Azure Cloud login URI is not a valid URI: {0}', e.message ?? e));
+		return;
+	}
+
+	// Add trailing slash if needed
+	if (!settingValue.endsWith('/')) {
+		settingValue += '/';
+	}
+
+	const azureEnterpriseAuthProvider = new AzureActiveDirectoryService(context, uriHandler, tokenStorage, settingValue);
+	await azureEnterpriseAuthProvider.initialize();
+
+	authProviderName ||= uri.authority;
+	const disposable = vscode.authentication.registerAuthenticationProvider('microsoft-sovereign-cloud', authProviderName, {
+		onDidChangeSessions: azureEnterpriseAuthProvider.onDidChangeSessions,
+		getSessions: (scopes: string[]) => azureEnterpriseAuthProvider.getSessions(scopes),
+		createSession: async (scopes: string[]) => {
+			try {
+				/* __GDPR__
+					"login" : {
+						"owner": "TylerLeonhardt",
+						"comment": "Used to determine the usage of the Azure Cloud Auth Provider.",
+						"scopes": { "classification": "PublicNonPersonalData", "purpose": "FeatureInsight", "comment": "Used to determine what scope combinations are being requested." }
+					}
+				*/
+				telemetryReporter.sendTelemetryEvent('loginAzureCloud', {
+					// Get rid of guids from telemetry.
+					scopes: JSON.stringify(scopes.map(s => s.replace(/[0-9A-F]{8}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{12}/i, '{guid}'))),
+				});
+
+				return await azureEnterpriseAuthProvider.createSession(scopes.sort());
+			} catch (e) {
+				/* __GDPR__
+					"loginFailed" : { "owner": "TylerLeonhardt", "comment": "Used to determine how often users run into issues with the login flow." }
+				*/
+				telemetryReporter.sendTelemetryEvent('loginAzureCloudFailed');
+
+				throw e;
+			}
+		},
+		removeSession: async (id: string) => {
+			try {
+				/* __GDPR__
+					"logout" : { "owner": "TylerLeonhardt", "comment": "Used to determine how often users log out." }
+				*/
+				telemetryReporter.sendTelemetryEvent('logoutAzureCloud');
+
+				await azureEnterpriseAuthProvider.removeSessionById(id);
+			} catch (e) {
+				/* __GDPR__
+					"logoutFailed" : { "owner": "TylerLeonhardt", "comment": "Used to determine how often fail to log out." }
+				*/
+				telemetryReporter.sendTelemetryEvent('logoutAzureCloudFailed');
+			}
+		}
+	}, { supportsMultipleAccounts: true });
+
+	context.subscriptions.push(disposable);
+	return disposable;
+}
+
 export async function activate(context: vscode.ExtensionContext) {
 	const { name, version, aiKey } = context.extension.packageJSON as { name: string; version: string; aiKey: string };
 	const telemetryReporter = new TelemetryReporter(aiKey);
 
-	const loginService = new AzureActiveDirectoryService(context);
+	const uriHandler = new UriEventHandler();
+	context.subscriptions.push(uriHandler);
+	context.subscriptions.push(vscode.window.registerUriHandler(uriHandler));
+
+	const betterSecretStorage = new BetterTokenStorage<IStoredSession>('microsoft.login.keylist', context);
+
+	const loginService = new AzureActiveDirectoryService(context, uriHandler, betterSecretStorage);
 	await loginService.initialize();
 
 	context.subscriptions.push(vscode.authentication.registerAuthenticationProvider('microsoft', 'Microsoft', {
-		onDidChangeSessions: onDidChangeSessions.event,
+		onDidChangeSessions: loginService.onDidChangeSessions,
 		getSessions: (scopes: string[]) => loginService.getSessions(scopes),
 		createSession: async (scopes: string[]) => {
 			try {
@@ -31,9 +118,7 @@ export async function activate(context: vscode.ExtensionContext) {
 					scopes: JSON.stringify(scopes.map(s => s.replace(/[0-9A-F]{8}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{12}/i, '{guid}'))),
 				});
 
-				const session = await loginService.createSession(scopes.sort());
-				onDidChangeSessions.fire({ added: [session], removed: [], changed: [] });
-				return session;
+				return await loginService.createSession(scopes.sort());
 			} catch (e) {
 				/* __GDPR__
 					"loginFailed" : { "owner": "TylerLeonhardt", "comment": "Used to determine how often users run into issues with the login flow." }
@@ -50,10 +135,7 @@ export async function activate(context: vscode.ExtensionContext) {
 				*/
 				telemetryReporter.sendTelemetryEvent('logout');
 
-				const session = await loginService.removeSessionById(id);
-				if (session) {
-					onDidChangeSessions.fire({ added: [], removed: [session], changed: [] });
-				}
+				await loginService.removeSessionById(id);
 			} catch (e) {
 				/* __GDPR__
 					"logoutFailed" : { "owner": "TylerLeonhardt", "comment": "Used to determine how often fail to log out." }
@@ -63,6 +145,15 @@ export async function activate(context: vscode.ExtensionContext) {
 		}
 	}, { supportsMultipleAccounts: true }));
 
+	let azureCloudAuthProviderDisposable = await initAzureCloudAuthProvider(context, telemetryReporter, uriHandler, betterSecretStorage);
+
+	context.subscriptions.push(vscode.workspace.onDidChangeConfiguration(async e => {
+		if (e.affectsConfiguration('microsoft-sovereign-cloud.endpoint')) {
+			azureCloudAuthProviderDisposable?.dispose();
+			azureCloudAuthProviderDisposable = await initAzureCloudAuthProvider(context, telemetryReporter, uriHandler, betterSecretStorage);
+		}
+	}));
+
 	return;
 }