Ability to pass down WWW-Authenticate challenges down to Auth Providers (#261717)

* Initial plan

* Implement authentication challenges support for mandatory MFA

Co-authored-by: TylerLeonhardt <2644648+TylerLeonhardt@users.noreply.github.com>

* Add documentation and integration test for authentication challenges

Co-authored-by: TylerLeonhardt <2644648+TylerLeonhardt@users.noreply.github.com>

* Add validation script and finalize implementation

Co-authored-by: TylerLeonhardt <2644648+TylerLeonhardt@users.noreply.github.com>

* Update authentication challenges API to use AuthenticationConstraint interface

Co-authored-by: TylerLeonhardt <2644648+TylerLeonhardt@users.noreply.github.com>

* Get it compiling... who knows if it works

* New parseWWWAuthenticateHeader behavior

* works

* let's go with this for now

* Good shape

* bye

* final polish

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>

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

@@ -45,11 +45,17 @@ export class ScopeData {
 	 */
 	readonly tenantId: string | undefined;
 
-	constructor(readonly originalScopes: readonly string[] = [], authorizationServer?: Uri) {
+	/**
+	 * The claims to include in the token request.
+	 */
+	readonly claims?: string;
+
+	constructor(readonly originalScopes: readonly string[] = [], claims?: string, authorizationServer?: Uri) {
 		const modifiedScopes = [...originalScopes];
 		modifiedScopes.sort();
 		this.allScopes = modifiedScopes;
 		this.scopeStr = modifiedScopes.join(' ');
+		this.claims = claims;
 		this.scopesToSend = this.getScopesToSend(modifiedScopes);
 		this.clientId = this.getClientId(this.allScopes);
 		this.tenant = this.getTenant(this.allScopes, authorizationServer);

extensions/microsoft-authentication/src/common/test/scopeData.test.ts

@@ -75,21 +75,31 @@ suite('ScopeData', () => {
 		assert.strictEqual(scopeData.tenantId, 'some_guid');
 	});
 
+	test('should not return claims', () => {
+		const scopeData = new ScopeData(['custom_scope']);
+		assert.strictEqual(scopeData.claims, undefined);
+	});
+
+	test('should return claims', () => {
+		const scopeData = new ScopeData(['custom_scope'], 'test');
+		assert.strictEqual(scopeData.claims, 'test');
+	});
+
 	test('should extract tenant from authorization server URL path', () => {
 		const authorizationServer = Uri.parse('https://login.microsoftonline.com/tenant123/oauth2/v2.0');
-		const scopeData = new ScopeData(['custom_scope'], authorizationServer);
+		const scopeData = new ScopeData(['custom_scope'], undefined, authorizationServer);
 		assert.strictEqual(scopeData.tenant, 'tenant123');
 	});
 
 	test('should fallback to default tenant if authorization server URL has no path segments', () => {
 		const authorizationServer = Uri.parse('https://login.microsoftonline.com');
-		const scopeData = new ScopeData(['custom_scope'], authorizationServer);
+		const scopeData = new ScopeData(['custom_scope'], undefined, authorizationServer);
 		assert.strictEqual(scopeData.tenant, 'organizations');
 	});
 
 	test('should prioritize authorization server URL over VSCODE_TENANT scope', () => {
 		const authorizationServer = Uri.parse('https://login.microsoftonline.com/url_tenant/oauth2/v2.0');
-		const scopeData = new ScopeData(['custom_scope', 'VSCODE_TENANT:scope_tenant'], authorizationServer);
+		const scopeData = new ScopeData(['custom_scope', 'VSCODE_TENANT:scope_tenant'], undefined, authorizationServer);
 		assert.strictEqual(scopeData.tenant, 'url_tenant');
 	});
 });

extensions/microsoft-authentication/src/extensionV2.ts

@@ -60,7 +60,7 @@ async function initMicrosoftSovereignCloudAuthProvider(
 		'microsoft-sovereign-cloud',
 		authProviderName,
 		authProvider,
-		{ supportsMultipleAccounts: true }
+		{ supportsMultipleAccounts: true, supportsChallenges: true }
 	);
 	context.subscriptions.push(disposable);
 	return disposable;
@@ -81,6 +81,7 @@ export async function activate(context: ExtensionContext, mainTelemetryReporter:
 		authProvider,
 		{
 			supportsMultipleAccounts: true,
+			supportsChallenges: true,
 			supportedAuthorizationServers: [
 				Uri.parse('https://login.microsoftonline.com/*/v2.0')
 			]

extensions/microsoft-authentication/src/node/authProvider.ts

@@ -2,8 +2,8 @@
  *  Copyright (c) Microsoft Corporation. All rights reserved.
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
-import { AccountInfo, AuthenticationResult, ClientAuthError, ClientAuthErrorCodes, ServerError } from '@azure/msal-node';
-import { AuthenticationGetSessionOptions, AuthenticationProvider, AuthenticationProviderAuthenticationSessionsChangeEvent, AuthenticationProviderSessionOptions, AuthenticationSession, AuthenticationSessionAccountInformation, CancellationError, EventEmitter, ExtensionContext, ExtensionKind, l10n, LogOutputChannel, Uri, window } from 'vscode';
+import { AccountInfo, AuthenticationResult, ClientAuthError, ClientAuthErrorCodes, ServerError, SilentFlowRequest } from '@azure/msal-node';
+import { AuthenticationChallenge, AuthenticationConstraint, AuthenticationGetSessionOptions, AuthenticationProvider, AuthenticationProviderAuthenticationSessionsChangeEvent, AuthenticationProviderSessionOptions, AuthenticationSession, AuthenticationSessionAccountInformation, CancellationError, EventEmitter, ExtensionContext, ExtensionKind, l10n, LogOutputChannel, Uri, window } from 'vscode';
 import { Environment } from '@azure/ms-rest-azure-env';
 import { CachedPublicClientApplicationManager } from './publicClientCache';
 import { UriEventHandler } from '../UriEventHandler';
@@ -14,6 +14,7 @@ import { EventBufferer } from '../common/event';
 import { BetterTokenStorage } from '../betterSecretStorage';
 import { IStoredSession } from '../AADHelper';
 import { ExtensionHost, getMsalFlows } from './flows';
+import { base64Decode } from './buffer';
 
 const redirectUri = 'https://vscode.dev/redirect';
 const MSA_TID = '9188040d-6c67-4c5b-b112-36a304b66dad';
@@ -153,7 +154,7 @@ export class MsalAuthProvider implements AuthenticationProvider {
 
 	async getSessions(scopes: string[] | undefined, options: AuthenticationGetSessionOptions = {}): Promise<AuthenticationSession[]> {
 		const askingForAll = scopes === undefined;
-		const scopeData = new ScopeData(scopes, options?.authorizationServer);
+		const scopeData = new ScopeData(scopes, undefined, options?.authorizationServer);
 		// Do NOT use `scopes` beyond this place in the code. Use `scopeData` instead.
 		this._logger.info('[getSessions]', askingForAll ? '[all]' : `[${scopeData.scopeStr}]`, 'starting');
 
@@ -183,7 +184,7 @@ export class MsalAuthProvider implements AuthenticationProvider {
 	}
 
 	async createSession(scopes: readonly string[], options: AuthenticationProviderSessionOptions): Promise<AuthenticationSession> {
-		const scopeData = new ScopeData(scopes, options.authorizationServer);
+		const scopeData = new ScopeData(scopes, undefined, options.authorizationServer);
 		// Do NOT use `scopes` beyond this place in the code. Use `scopeData` instead.
 
 		this._logger.info('[createSession]', `[${scopeData.scopeStr}]`, 'starting');
@@ -284,6 +285,139 @@ export class MsalAuthProvider implements AuthenticationProvider {
 		this._logger.info('[removeSession]', sessionId, `attempted to remove ${promises.length} sessions`);
 	}
 
+	async getSessionsFromChallenges(constraint: AuthenticationConstraint, options: AuthenticationProviderSessionOptions): Promise<readonly AuthenticationSession[]> {
+		this._logger.info('[getSessionsFromChallenges]', 'starting with', constraint.challenges.length, 'challenges');
+
+		// Use scopes from constraint if provided, otherwise extract from challenges
+		const scopes = constraint.scopes?.length ? [...constraint.scopes] : this.extractScopesFromChallenges(constraint.challenges);
+		const claims = this.extractClaimsFromChallenges(constraint.challenges);
+		if (!claims) {
+			throw new Error('No claims found in authentication challenges');
+		}
+		const scopeData = new ScopeData(scopes, claims, options?.authorizationServer);
+		this._logger.info('[getSessionsFromChallenges]', `[${scopeData.scopeStr}]`, 'with claims:', scopeData.claims);
+
+		const cachedPca = await this._publicClientManager.getOrCreate(scopeData.clientId);
+		const sessions = await this.getAllSessionsForPca(cachedPca, scopeData, options?.account);
+
+		this._logger.info('[getSessionsFromChallenges]', 'returning', sessions.length, 'sessions');
+		return sessions;
+	}
+
+	async createSessionFromChallenges(constraint: AuthenticationConstraint, options: AuthenticationProviderSessionOptions): Promise<AuthenticationSession> {
+		this._logger.info('[createSessionFromChallenges]', 'starting with', constraint.challenges.length, 'challenges');
+
+		// Use scopes from constraint if provided, otherwise extract from challenges
+		const scopes = constraint.scopes?.length ? [...constraint.scopes] : this.extractScopesFromChallenges(constraint.challenges);
+		const claims = this.extractClaimsFromChallenges(constraint.challenges);
+
+		// Use scopes if available, otherwise fall back to default scopes
+		const effectiveScopes = scopes.length > 0 ? scopes : ['https://graph.microsoft.com/User.Read'];
+
+		const scopeData = new ScopeData(effectiveScopes, claims, options.authorizationServer);
+		this._logger.info('[createSessionFromChallenges]', `[${scopeData.scopeStr}]`, 'starting with claims:', claims);
+
+		const cachedPca = await this._publicClientManager.getOrCreate(scopeData.clientId);
+
+		// Used for showing a friendlier message to the user when the explicitly cancel a flow.
+		let userCancelled: boolean | undefined;
+		const yes = l10n.t('Yes');
+		const no = l10n.t('No');
+		const promptToContinue = async (mode: string) => {
+			if (userCancelled === undefined) {
+				// We haven't had a failure yet so wait to prompt
+				return;
+			}
+			const message = userCancelled
+				? l10n.t('Having trouble logging in? Would you like to try a different way? ({0})', mode)
+				: l10n.t('You have not yet finished authorizing this extension to use your Microsoft Account. Would you like to try a different way? ({0})', mode);
+			const result = await window.showWarningMessage(message, yes, no);
+			if (result !== yes) {
+				throw new CancellationError();
+			}
+		};
+
+		const isNodeEnvironment = typeof process !== 'undefined' && typeof process?.versions?.node === 'string';
+		const flows = getMsalFlows({
+			extensionHost: isNodeEnvironment
+				? this._context.extension.extensionKind === ExtensionKind.UI ? ExtensionHost.Local : ExtensionHost.Remote
+				: ExtensionHost.WebWorker,
+		});
+
+		const authority = new URL(scopeData.tenant, this._env.activeDirectoryEndpointUrl).toString();
+		let lastError: Error | undefined;
+		for (const flow of flows) {
+			if (flow !== flows[0]) {
+				try {
+					await promptToContinue(flow.label);
+				} finally {
+					this._telemetryReporter.sendLoginFailedEvent();
+				}
+			}
+			try {
+				// Create the authentication request with claims if provided
+				const authRequest = {
+					cachedPca,
+					authority,
+					scopes: scopeData.scopesToSend,
+					loginHint: options.account?.label,
+					windowHandle: window.nativeHandle ? Buffer.from(window.nativeHandle) : undefined,
+					logger: this._logger,
+					uriHandler: this._uriHandler,
+					claims: scopeData.claims
+				};
+
+				const result = await flow.trigger(authRequest);
+
+				const session = this.sessionFromAuthenticationResult(result, scopeData.originalScopes);
+				this._telemetryReporter.sendLoginEvent(session.scopes);
+				this._logger.info('[createSessionFromChallenges]', `[${scopeData.scopeStr}]`, 'returned session');
+				return session;
+			} catch (e) {
+				lastError = e as Error;
+				if (e instanceof ClientAuthError && e.errorCode === ClientAuthErrorCodes.userCanceled) {
+					this._logger.info('[createSessionFromChallenges]', `[${scopeData.scopeStr}]`, 'user cancelled');
+					userCancelled = true;
+					continue;
+				}
+				this._logger.error('[createSessionFromChallenges]', `[${scopeData.scopeStr}]`, 'error', e);
+				throw e;
+			}
+		}
+
+		this._telemetryReporter.sendLoginFailedEvent();
+		throw lastError ?? new Error('No auth flow succeeded');
+	}
+
+	private extractScopesFromChallenges(challenges: readonly AuthenticationChallenge[]): string[] {
+		const scopes: string[] = [];
+		for (const challenge of challenges) {
+			if (challenge.scheme.toLowerCase() === 'bearer' && challenge.params.scope) {
+				scopes.push(...challenge.params.scope.split(' '));
+			}
+		}
+		return scopes;
+	}
+
+	private extractClaimsFromChallenges(challenges: readonly AuthenticationChallenge[]): string | undefined {
+		for (const challenge of challenges) {
+			if (challenge.scheme.toLowerCase() === 'bearer' && challenge.params.claims) {
+				try {
+					return base64Decode(challenge.params.claims);
+				} catch (e) {
+					this._logger.warn('[extractClaimsFromChallenges]', 'failed to decode claims... checking if it is already JSON', e);
+					try {
+						JSON.parse(challenge.params.claims);
+						return challenge.params.claims;
+					} catch (e) {
+						this._logger.error('[extractClaimsFromChallenges]', 'failed to parse claims as JSON... returning undefined', e);
+					}
+				}
+			}
+		}
+		return undefined;
+	}
+
 	//#endregion
 
 	private async getAllSessionsForPca(
@@ -341,11 +475,18 @@ export class MsalAuthProvider implements AuthenticationProvider {
 							forceRefresh = true;
 						}
 					}
+					// When claims are present, force refresh to ensure we get a token that satisfies the claims
+					let claims: string | undefined;
+					if (scopeData.claims) {
+						forceRefresh = true;
+						claims = scopeData.claims;
+					}
 					const result = await cachedPca.acquireTokenSilent({
 						account,
 						authority,
 						scopes: scopeData.scopesToSend,
 						redirectUri,
+						claims,
 						forceRefresh
 					});
 					sessions.push(this.sessionFromAuthenticationResult(result, scopeData.originalScopes));

extensions/microsoft-authentication/src/node/cachedPublicClientApplication.ts

@@ -112,7 +112,7 @@ export class CachedPublicClientApplication implements ICachedPublicClientApplica
 				this._logger.debug(`[acquireTokenSilent] [${this._clientId}] [${request.authority}] [${request.scopes.join(' ')}] [${request.account.username}] id token is expired or about to expire. Forcing refresh...`);
 				const newRequest = this._isBrokerAvailable
 					// HACK: Broker doesn't support forceRefresh so we need to pass in claims which will force a refresh
-					? { ...request, claims: '{ "id_token": {}}' }
+					? { ...request, claims: request.claims ?? '{ "id_token": {}}' }
 					: { ...request, forceRefresh: true };
 				result = await this._sequencer.queue(() => this._pca.acquireTokenSilent(newRequest));
 				this._logger.debug(`[acquireTokenSilent] [${this._clientId}] [${request.authority}] [${request.scopes.join(' ')}] [${request.account.username}] got forced result`);
@@ -130,7 +130,7 @@ export class CachedPublicClientApplication implements ICachedPublicClientApplica
 					// there has been a situation where both tokens are expired.
 					if (this._isBrokerAvailable) {
 						this._logger.error(`[acquireTokenSilent] [${this._clientId}] [${request.authority}] [${request.scopes.join(' ')}] [${request.account.username}] forcing refresh with different claims...`);
-						const newRequest = { ...request, claims: '{ "access_token": {}}' };
+						const newRequest = { ...request, claims: request.claims ?? '{ "access_token": {}}' };
 						result = await this._sequencer.queue(() => this._pca.acquireTokenSilent(newRequest));
 						this._logger.debug(`[acquireTokenSilent] [${this._clientId}] [${request.authority}] [${request.scopes.join(' ')}] [${request.account.username}] got forced result with different claims`);
 						const newIdTokenExpirationInSecs = (result.idTokenClaims as { exp?: number }).exp;

extensions/microsoft-authentication/src/node/flows.ts

@@ -31,6 +31,7 @@ interface IMsalFlowTriggerOptions {
 	windowHandle?: Buffer;
 	logger: LogOutputChannel;
 	uriHandler: UriEventHandler;
+	claims?: string;
 }
 
 interface IMsalFlow {
@@ -46,7 +47,7 @@ class DefaultLoopbackFlow implements IMsalFlow {
 		supportsWebWorkerExtensionHost: false
 	};
 
-	async trigger({ cachedPca, authority, scopes, loginHint, windowHandle, logger }: IMsalFlowTriggerOptions): Promise<AuthenticationResult> {
+	async trigger({ cachedPca, authority, scopes, claims, loginHint, windowHandle, logger }: IMsalFlowTriggerOptions): Promise<AuthenticationResult> {
 		logger.info('Trying default msal flow...');
 		return await cachedPca.acquireTokenInteractive({
 			openBrowser: async (url: string) => { await env.openExternal(Uri.parse(url)); },
@@ -56,7 +57,8 @@ class DefaultLoopbackFlow implements IMsalFlow {
 			errorTemplate: loopbackTemplate,
 			loginHint,
 			prompt: loginHint ? undefined : 'select_account',
-			windowHandle
+			windowHandle,
+			claims
 		});
 	}
 }
@@ -68,7 +70,7 @@ class UrlHandlerFlow implements IMsalFlow {
 		supportsWebWorkerExtensionHost: false
 	};
 
-	async trigger({ cachedPca, authority, scopes, loginHint, windowHandle, logger, uriHandler }: IMsalFlowTriggerOptions): Promise<AuthenticationResult> {
+	async trigger({ cachedPca, authority, scopes, claims, loginHint, windowHandle, logger, uriHandler }: IMsalFlowTriggerOptions): Promise<AuthenticationResult> {
 		logger.info('Trying protocol handler flow...');
 		const loopbackClient = new UriHandlerLoopbackClient(uriHandler, redirectUri, logger);
 		return await cachedPca.acquireTokenInteractive({
@@ -78,7 +80,8 @@ class UrlHandlerFlow implements IMsalFlow {
 			loopbackClient,
 			loginHint,
 			prompt: loginHint ? undefined : 'select_account',
-			windowHandle
+			windowHandle,
+			claims
 		});
 	}
 }