Migrate old accounts to MSAL (#234147)

Bascally, we reach into the old location in secret storage and if we find sessions (with a refresh token) we seed that in the MSAL world.

We do this one time... unless they switch back to the old world and then switch to the new world.

This has two different behaviors depending on if the Broker is used:
* If the broker is not used, this does what you might expect. It makes it seem totally transparent to the user that something has changed. All sessions get migrated over and the user is still logged in to what they were previously.
* If the broker is used... you don't get automatically logged in _unless_ you have already logged in to that account at the OS level. So this helps skip the "VS Code access layer" outlined in `accountAccess.ts`. Not as good as the previous bullet, but this is the best we can do in the broker world.

In time, we can remove this migration along with the old way of doing things.

extensions/microsoft-authentication/src/extensionV1.ts

@@ -105,6 +105,10 @@ async function initMicrosoftSovereignCloudAuthProvider(context: vscode.Extension
 }
 
 export async function activate(context: vscode.ExtensionContext, telemetryReporter: TelemetryReporter) {
+	// If we ever activate the old flow, then mark that we will need to migrate when the user upgrades to v2.
+	// TODO: MSAL Migration. Remove this when we remove the old flow.
+	context.globalState.update('msalMigration', false);
+
 	const uriHandler = new UriEventHandler();
 	context.subscriptions.push(uriHandler);
 	const betterSecretStorage = new BetterTokenStorage<IStoredSession>('microsoft.login.keylist', context);

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

@@ -13,6 +13,8 @@ import { MicrosoftAccountType, MicrosoftAuthenticationTelemetryReporter } from '
 import { loopbackTemplate } from './loopbackTemplate';
 import { ScopeData } from '../common/scopeData';
 import { EventBufferer } from '../common/event';
+import { BetterTokenStorage } from '../betterSecretStorage';
+import { IStoredSession } from '../AADHelper';
 
 const redirectUri = 'https://vscode.dev/redirect';
 const MSA_TID = '9188040d-6c67-4c5b-b112-36a304b66dad';
@@ -43,16 +45,16 @@ export class MsalAuthProvider implements AuthenticationProvider {
 	onDidChangeSessions = this._onDidChangeSessionsEmitter.event;
 
 	constructor(
-		context: ExtensionContext,
+		private readonly _context: ExtensionContext,
 		private readonly _telemetryReporter: MicrosoftAuthenticationTelemetryReporter,
 		private readonly _logger: LogOutputChannel,
 		private readonly _uriHandler: UriEventHandler,
 		private readonly _env: Environment = Environment.AzureCloud
 	) {
-		this._disposables = context.subscriptions;
+		this._disposables = _context.subscriptions;
 		this._publicClientManager = new CachedPublicClientApplicationManager(
-			context.globalState,
-			context.secrets,
+			_context.globalState,
+			_context.secrets,
 			this._logger,
 			this._env.name
 		);
@@ -85,9 +87,41 @@ export class MsalAuthProvider implements AuthenticationProvider {
 		);
 	}
 
+	/**
+	 * Migrate sessions from the old secret storage to MSAL.
+	 * TODO: MSAL Migration. Remove this when we remove the old flow.
+	 */
+	private async _migrateSessions() {
+		const betterSecretStorage = new BetterTokenStorage<IStoredSession>('microsoft.login.keylist', this._context);
+		const sessions = await betterSecretStorage.getAll(item => {
+			item.endpoint ||= Environment.AzureCloud.activeDirectoryEndpointUrl;
+			return item.endpoint === this._env.activeDirectoryEndpointUrl;
+		});
+		this._context.globalState.update('msalMigration', true);
+
+		const clientTenantMap = new Map<string, { clientId: string; tenant: string; refreshTokens: string[] }>();
+
+		for (const session of sessions) {
+			const scopeData = new ScopeData(session.scope.split(' '));
+			const key = `${scopeData.clientId}:${scopeData.tenant}`;
+			if (!clientTenantMap.has(key)) {
+				clientTenantMap.set(key, { clientId: scopeData.clientId, tenant: scopeData.tenant, refreshTokens: [] });
+			}
+			clientTenantMap.get(key)!.refreshTokens.push(session.refreshToken);
+		}
+
+		for (const { clientId, tenant, refreshTokens } of clientTenantMap.values()) {
+			await this.getOrCreatePublicClientApplication(clientId, tenant, refreshTokens);
+		}
+	}
+
 	async initialize(): Promise<void> {
 		await this._eventBufferer.bufferEventsAsync(() => this._publicClientManager.initialize());
 
+		if (!this._context.globalState.get('msalMigration', false)) {
+			await this._migrateSessions();
+		}
+
 		// Send telemetry for existing accounts
 		for (const cachedPca of this._publicClientManager.getAll()) {
 			for (const account of cachedPca.accounts) {
@@ -259,9 +293,9 @@ export class MsalAuthProvider implements AuthenticationProvider {
 
 	//#endregion
 
-	private async getOrCreatePublicClientApplication(clientId: string, tenant: string): Promise<ICachedPublicClientApplication> {
+	private async getOrCreatePublicClientApplication(clientId: string, tenant: string, refreshTokensToMigrate?: string[]): Promise<ICachedPublicClientApplication> {
 		const authority = new URL(tenant, this._env.activeDirectoryEndpointUrl).toString();
-		return await this._publicClientManager.getOrCreate(clientId, authority);
+		return await this._publicClientManager.getOrCreate(clientId, authority, refreshTokensToMigrate);
 	}
 
 	private async getAllSessionsForPca(

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

@@ -3,7 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { PublicClientApplication, AccountInfo, Configuration, SilentFlowRequest, AuthenticationResult, InteractiveRequest, LogLevel } from '@azure/msal-node';
+import { PublicClientApplication, AccountInfo, Configuration, SilentFlowRequest, AuthenticationResult, InteractiveRequest, LogLevel, RefreshTokenRequest } from '@azure/msal-node';
 import { NativeBrokerPlugin } from '@azure/msal-node-extensions';
 import { Disposable, Memento, SecretStorage, LogOutputChannel, window, ProgressLocation, l10n, EventEmitter } from 'vscode';
 import { Delayer, raceCancellationAndTimeoutError } from '../common/async';
@@ -128,6 +128,24 @@ export class CachedPublicClientApplication implements ICachedPublicClientApplica
 		return result;
 	}
 
+	/**
+	 * Allows for passing in a refresh token to get a new access token. This is the migration scenario.
+	 * TODO: MSAL Migration. Remove this when we remove the old flow.
+	 * @param request a {@link RefreshTokenRequest} object that contains the refresh token and other parameters.
+	 * @returns an {@link AuthenticationResult} object that contains the result of the token acquisition operation.
+	 */
+	async acquireTokenByRefreshToken(request: RefreshTokenRequest) {
+		this._logger.debug(`[acquireTokenByRefreshToken] [${this._clientId}] [${this._authority}] [${request.scopes.join(' ')}]`);
+		const result = await this._pca.acquireTokenByRefreshToken(request);
+		if (result) {
+			this._setupRefresh(result);
+			if (this._isBrokerAvailable && result.account) {
+				await this._accountAccess.setAllowedAccess(result.account, true);
+			}
+		}
+		return result;
+	}
+
 	removeAccount(account: AccountInfo): Promise<void> {
 		this._globalMemento.update(`lastRemoval:${this._clientId}:${this._authority}`, new Date());
 		if (this._isBrokerAvailable) {

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

@@ -94,19 +94,37 @@ export class CachedPublicClientApplicationManager implements ICachedPublicClient
 		Disposable.from(...this._pcaDisposables.values()).dispose();
 	}
 
-	async getOrCreate(clientId: string, authority: string): Promise<ICachedPublicClientApplication> {
+	async getOrCreate(clientId: string, authority: string, refreshTokensToMigrate?: string[]): Promise<ICachedPublicClientApplication> {
 		// Use the clientId and authority as the key
 		const pcasKey = JSON.stringify({ clientId, authority });
 		let pca = this._pcas.get(pcasKey);
 		if (pca) {
 			this._logger.debug(`[getOrCreate] [${clientId}] [${authority}] PublicClientApplicationManager cache hit`);
-			return pca;
+		} else {
+			this._logger.debug(`[getOrCreate] [${clientId}] [${authority}] PublicClientApplicationManager cache miss, creating new PCA...`);
+			pca = await this._doCreatePublicClientApplication(clientId, authority, pcasKey);
+			await this._storePublicClientApplications();
+			this._logger.debug(`[getOrCreate] [${clientId}] [${authority}] PCA created.`);
 		}
 
-		this._logger.debug(`[getOrCreate] [${clientId}] [${authority}] PublicClientApplicationManager cache miss, creating new PCA...`);
-		pca = await this._doCreatePublicClientApplication(clientId, authority, pcasKey);
-		await this._storePublicClientApplications();
-		this._logger.debug(`[getOrCreate] [${clientId}] [${authority}] PCA created.`);
+		// TODO: MSAL Migration. Remove this when we remove the old flow.
+		if (refreshTokensToMigrate?.length) {
+			this._logger.debug(`[getOrCreate] [${clientId}] [${authority}] Migrating refresh tokens to PCA...`);
+			for (const refreshToken of refreshTokensToMigrate) {
+				try {
+					// Use the refresh token to acquire a result. This will cache the refresh token for future operations.
+					// The scopes don't matter here since we can create any token from the refresh token.
+					const result = await pca.acquireTokenByRefreshToken({ refreshToken, forceCache: true, scopes: [] });
+					if (result?.account) {
+						this._logger.debug(`[getOrCreate] [${clientId}] [${authority}] Refresh token migrated to PCA.`);
+					}
+				} catch (e) {
+					this._logger.error(`[getOrCreate] [${clientId}] [${authority}] Error migrating refresh token:`, e);
+				}
+			}
+			// reinitialize the PCA so the account is properly cached
+			await pca.initialize();
+		}
 		return pca;
 	}