Add cancellable promise to Microsoft auth flows

Fixes #211406

extensions/microsoft-authentication/src/AADHelper.ts

@@ -6,7 +6,7 @@
 import * as vscode from 'vscode';
 import * as path from 'path';
 import { isSupportedEnvironment } from './common/uri';
-import { IntervalTimer, SequencerByKey } from './common/async';
+import { IntervalTimer, raceCancellationAndTimeoutError, SequencerByKey } from './common/async';
 import { generateCodeChallenge, generateCodeVerifier, randomUUID } from './cryptoUtils';
 import { BetterTokenStorage, IDidChangeInOtherWindowEvent } from './betterSecretStorage';
 import { LoopbackAuthServer } from './node/authServer';
@@ -314,25 +314,27 @@ export class AzureActiveDirectoryService {
 			throw new Error('Sign in to non-public clouds is not supported on the web.');
 		}
 
-		if (runsRemote || runsServerless) {
-			return this.createSessionWithoutLocalServer(scopeData);
-		}
+		return await vscode.window.withProgress({ location: vscode.ProgressLocation.Notification, title: vscode.l10n.t('Signing in to your account...'), cancellable: true }, async (_progress, token) => {
+			if (runsRemote || runsServerless) {
+				return await this.createSessionWithoutLocalServer(scopeData, token);
+			}
 
-		try {
-			return await this.createSessionWithLocalServer(scopeData);
-		} catch (e) {
-			this._logger.error(`[${scopeData.scopeStr}] Error creating session: ${e}`);
+			try {
+				return await this.createSessionWithLocalServer(scopeData, token);
+			} catch (e) {
+				this._logger.error(`[${scopeData.scopeStr}] Error creating session: ${e}`);
 
-			// If the error was about starting the server, try directly hitting the login endpoint instead
-			if (e.message === 'Error listening to server' || e.message === 'Closed' || e.message === 'Timeout waiting for port') {
-				return this.createSessionWithoutLocalServer(scopeData);
-			}
+				// If the error was about starting the server, try directly hitting the login endpoint instead
+				if (e.message === 'Error listening to server' || e.message === 'Closed' || e.message === 'Timeout waiting for port') {
+					return this.createSessionWithoutLocalServer(scopeData, token);
+				}
 
-			throw e;
-		}
+				throw e;
+			}
+		});
 	}
 
-	private async createSessionWithLocalServer(scopeData: IScopeData) {
+	private async createSessionWithLocalServer(scopeData: IScopeData, token: vscode.CancellationToken): Promise<vscode.AuthenticationSession> {
 		this._logger.trace(`[${scopeData.scopeStr}] Starting login flow with local server`);
 		const codeVerifier = generateCodeVerifier();
 		const codeChallenge = await generateCodeChallenge(codeVerifier);
@@ -353,7 +355,7 @@ export class AzureActiveDirectoryService {
 		let codeToExchange;
 		try {
 			vscode.env.openExternal(vscode.Uri.parse(`http://127.0.0.1:${server.port}/signin?nonce=${encodeURIComponent(server.nonce)}`));
-			const { code } = await server.waitForOAuthResponse();
+			const { code } = await raceCancellationAndTimeoutError(server.waitForOAuthResponse(), token, 1000 * 60 * 5); // 5 minutes
 			codeToExchange = code;
 		} finally {
 			setTimeout(() => {
@@ -368,7 +370,7 @@ export class AzureActiveDirectoryService {
 		return session;
 	}
 
-	private async createSessionWithoutLocalServer(scopeData: IScopeData): Promise<vscode.AuthenticationSession> {
+	private async createSessionWithoutLocalServer(scopeData: IScopeData, token: vscode.CancellationToken): Promise<vscode.AuthenticationSession> {
 		this._logger.trace(`[${scopeData.scopeStr}] Starting login flow without local server`);
 		let callbackUri = await vscode.env.asExternalUri(vscode.Uri.parse(`${vscode.env.uriScheme}://vscode.microsoft-authentication`));
 		const nonce = generateCodeVerifier();
@@ -395,21 +397,14 @@ export class AzureActiveDirectoryService {
 		const uri = vscode.Uri.parse(signInUrl.toString());
 		vscode.env.openExternal(uri);
 
-		let inputBox: vscode.InputBox | undefined;
-		const timeoutPromise = new Promise((_: (value: vscode.AuthenticationSession) => void, reject) => {
-			const wait = setTimeout(() => {
-				clearTimeout(wait);
-				inputBox?.dispose();
-				reject('Login timed out.');
-			}, 1000 * 60 * 5);
-		});
 
 		const existingNonces = this._pendingNonces.get(scopeData.scopeStr) || [];
 		this._pendingNonces.set(scopeData.scopeStr, [...existingNonces, nonce]);
 
 		// Register a single listener for the URI callback, in case the user starts the login process multiple times
 		// before completing it.
 		let existingPromise = this._codeExchangePromises.get(scopeData.scopeStr);
+		let inputBox: vscode.InputBox | undefined;
 		if (!existingPromise) {
 			if (isSupportedEnvironment(callbackUri)) {
 				existingPromise = this.handleCodeResponse(scopeData);
@@ -422,11 +417,12 @@ export class AzureActiveDirectoryService {
 
 		this._codeVerfifiers.set(nonce, codeVerifier);
 
-		return Promise.race([existingPromise, timeoutPromise])
+		return await raceCancellationAndTimeoutError(existingPromise, token, 1000 * 60 * 5) // 5 minutes
 			.finally(() => {
 				this._pendingNonces.delete(scopeData.scopeStr);
 				this._codeExchangePromises.delete(scopeData.scopeStr);
 				this._codeVerfifiers.delete(nonce);
+				inputBox?.dispose();
 			});
 	}
 

extensions/microsoft-authentication/src/common/async.ts

@@ -3,7 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { Disposable } from 'vscode';
+import { CancellationError, CancellationToken, Disposable } from 'vscode';
 
 export class SequencerByKey<TKey> {
 
@@ -47,3 +47,36 @@ export class IntervalTimer extends Disposable {
 		}, interval);
 	}
 }
+
+/**
+ * Returns a promise that rejects with an {@CancellationError} as soon as the passed token is cancelled.
+ * @see {@link raceCancellation}
+ */
+export function raceCancellationError<T>(promise: Promise<T>, token: CancellationToken): Promise<T> {
+	return new Promise((resolve, reject) => {
+		const ref = token.onCancellationRequested(() => {
+			ref.dispose();
+			reject(new CancellationError());
+		});
+		promise.then(resolve, reject).finally(() => ref.dispose());
+	});
+}
+
+export class TimeoutError extends Error {
+	constructor() {
+		super('Timed out');
+	}
+}
+
+export function raceTimeoutError<T>(promise: Promise<T>, timeout: number): Promise<T> {
+	return new Promise((resolve, reject) => {
+		const ref = setTimeout(() => {
+			reject(new CancellationError());
+		}, timeout);
+		promise.then(resolve, reject).finally(() => clearTimeout(ref));
+	});
+}
+
+export function raceCancellationAndTimeoutError<T>(promise: Promise<T>, token: CancellationToken, timeout: number): Promise<T> {
+	return raceCancellationError(raceTimeoutError(promise, timeout), token);
+}

src/vs/workbench/api/common/extHostAuthentication.ts

@@ -89,28 +89,28 @@ export class ExtHostAuthentication implements ExtHostAuthenticationShape {
 		});
 	}
 
-	$createSession(providerId: string, scopes: string[], options: vscode.AuthenticationProviderCreateSessionOptions): Promise<vscode.AuthenticationSession> {
+	async $createSession(providerId: string, scopes: string[], options: vscode.AuthenticationProviderCreateSessionOptions): Promise<vscode.AuthenticationSession> {
 		const providerData = this._authenticationProviders.get(providerId);
 		if (providerData) {
-			return Promise.resolve(providerData.provider.createSession(scopes, options));
+			return await providerData.provider.createSession(scopes, options);
 		}
 
 		throw new Error(`Unable to find authentication provider with handle: ${providerId}`);
 	}
 
-	$removeSession(providerId: string, sessionId: string): Promise<void> {
+	async $removeSession(providerId: string, sessionId: string): Promise<void> {
 		const providerData = this._authenticationProviders.get(providerId);
 		if (providerData) {
-			return Promise.resolve(providerData.provider.removeSession(sessionId));
+			return await providerData.provider.removeSession(sessionId);
 		}
 
 		throw new Error(`Unable to find authentication provider with handle: ${providerId}`);
 	}
 
-	$getSessions(providerId: string, scopes?: string[]): Promise<ReadonlyArray<vscode.AuthenticationSession>> {
+	async $getSessions(providerId: string, scopes?: string[]): Promise<ReadonlyArray<vscode.AuthenticationSession>> {
 		const providerData = this._authenticationProviders.get(providerId);
 		if (providerData) {
-			return Promise.resolve(providerData.provider.getSessions(scopes));
+			return await providerData.provider.getSessions(scopes);
 		}
 
 		throw new Error(`Unable to find authentication provider with handle: ${providerId}`);

src/vs/workbench/services/authentication/browser/authenticationExtensionsService.ts

@@ -246,23 +246,30 @@ export class AuthenticationExtensionsService extends Disposable implements IAuth
 			quickPick.placeholder = nls.localize('getSessionPlateholder', "Select an account for '{0}' to use or Esc to cancel", extensionName);
 
 			quickPick.onDidAccept(async _ => {
-				const session = quickPick.selectedItems[0].session ?? await this._authenticationService.createSession(providerId, scopes);
+				quickPick.dispose();
+				let session = quickPick.selectedItems[0].session;
+				if (!session) {
+					try {
+						session = await this._authenticationService.createSession(providerId, scopes);
+					} catch (e) {
+						reject(e);
+						return;
+					}
+				}
 				const accountName = session.account.label;
 
 				this._authenticationAccessService.updateAllowedExtensions(providerId, accountName, [{ id: extensionId, name: extensionName, allowed: true }]);
 				this.updateSessionPreference(providerId, extensionId, session);
 				this.removeAccessRequest(providerId, extensionId);
 
-				quickPick.dispose();
 				resolve(session);
 			});
 
 			quickPick.onDidHide(_ => {
+				quickPick.dispose();
 				if (!quickPick.selectedItems[0]) {
 					reject('User did not consent to account access');
 				}
-
-				quickPick.dispose();
 			});
 
 			quickPick.show();