From eab6f90c72520928256dfde94792a272e082f0a0 Mon Sep 17 00:00:00 2001 From: Tyler James Leonhardt <2644648+TylerLeonhardt@users.noreply.github.com> Date: Wed, 5 Mar 2025 17:50:14 -0800 Subject: [PATCH] Better lifecycle handling (#242758) I moved to a factory model because there was just so much that needed to be async. I think the amount of async code will be reduced in the future as we remove some migration logic, but this makes sure we don't accidentally create instances without awaiting their initialization. --- .../src/common/accountAccess.ts | 63 ++++++++++++++----- .../src/common/cachePlugin.ts | 2 +- .../src/common/publicClientCache.ts | 3 +- .../src/extensionV2.ts | 6 +- .../src/node/authProvider.ts | 26 +++++--- .../src/node/cachedPublicClientApplication.ts | 33 +++++++--- .../src/node/publicClientCache.ts | 61 ++++++++++++------ 7 files changed, 137 insertions(+), 57 deletions(-) diff --git a/extensions/microsoft-authentication/src/common/accountAccess.ts b/extensions/microsoft-authentication/src/common/accountAccess.ts index ced05311042..27d56b0bc1b 100644 --- a/extensions/microsoft-authentication/src/common/accountAccess.ts +++ b/extensions/microsoft-authentication/src/common/accountAccess.ts @@ -7,32 +7,47 @@ import { Disposable, Event, EventEmitter, LogOutputChannel, SecretStorage } from import { AccountInfo } from '@azure/msal-node'; export interface IAccountAccess { - initialize(): Promise; onDidAccountAccessChange: Event; isAllowedAccess(account: AccountInfo): boolean; setAllowedAccess(account: AccountInfo, allowed: boolean): Promise; } -export class ScopedAccountAccess implements IAccountAccess { +export class ScopedAccountAccess implements IAccountAccess, Disposable { private readonly _onDidAccountAccessChangeEmitter = new EventEmitter(); readonly onDidAccountAccessChange = this._onDidAccountAccessChangeEmitter.event; - private readonly _accountAccessSecretStorage: AccountAccessSecretStorage; - private value = new Array(); - constructor( + private readonly _disposable: Disposable; + + private constructor( + private readonly _accountAccessSecretStorage: IAccountAccessSecretStorage, + disposables: Disposable[] = [] + ) { + this._disposable = Disposable.from( + ...disposables, + this._onDidAccountAccessChangeEmitter, + this._accountAccessSecretStorage.onDidChange(() => this.update()) + ); + } + + static async create( secretStorage: SecretStorage, cloudName: string, logger: LogOutputChannel, - migrations?: { clientId: string; authority: string }[], - ) { - this._accountAccessSecretStorage = new AccountAccessSecretStorage(secretStorage, cloudName, logger, migrations); - this._accountAccessSecretStorage.onDidChange(() => this.update()); + migrations: { clientId: string; authority: string }[] | undefined, + ): Promise { + const storage = await AccountAccessSecretStorage.create(secretStorage, cloudName, logger, migrations); + const access = new ScopedAccountAccess(storage, [storage]); + await access.initialize(); + return access; } - async initialize() { - await this._accountAccessSecretStorage.initialize(); + dispose() { + this._disposable.dispose(); + } + + private async initialize(): Promise { await this.update(); } @@ -62,15 +77,22 @@ export class ScopedAccountAccess implements IAccountAccess { } } -export class AccountAccessSecretStorage { +interface IAccountAccessSecretStorage { + get(): Promise; + store(value: string[]): Thenable; + delete(): Thenable; + onDidChange: Event; +} + +class AccountAccessSecretStorage implements IAccountAccessSecretStorage, Disposable { private _disposable: Disposable; - private readonly _onDidChangeEmitter = new EventEmitter; + private readonly _onDidChangeEmitter = new EventEmitter(); readonly onDidChange: Event = this._onDidChangeEmitter.event; private readonly _key = `accounts-${this._cloudName}`; - constructor( + private constructor( private readonly _secretStorage: SecretStorage, private readonly _cloudName: string, private readonly _logger: LogOutputChannel, @@ -86,10 +108,21 @@ export class AccountAccessSecretStorage { ); } + static async create( + secretStorage: SecretStorage, + cloudName: string, + logger: LogOutputChannel, + migrations?: { clientId: string; authority: string }[], + ): Promise { + const storage = new AccountAccessSecretStorage(secretStorage, cloudName, logger, migrations); + await storage.initialize(); + return storage; + } + /** * TODO: Remove this method after a release with the migration */ - async initialize(): Promise { + private async initialize(): Promise { if (!this._migrations) { return; } diff --git a/extensions/microsoft-authentication/src/common/cachePlugin.ts b/extensions/microsoft-authentication/src/common/cachePlugin.ts index 91b4f0ee6a8..b87fdb78d9b 100644 --- a/extensions/microsoft-authentication/src/common/cachePlugin.ts +++ b/extensions/microsoft-authentication/src/common/cachePlugin.ts @@ -6,7 +6,7 @@ import { ICachePlugin, TokenCacheContext } from '@azure/msal-node'; import { Disposable, EventEmitter, SecretStorage } from 'vscode'; -export class SecretStorageCachePlugin implements ICachePlugin { +export class SecretStorageCachePlugin implements ICachePlugin, Disposable { private readonly _onDidChange: EventEmitter = new EventEmitter(); readonly onDidChange = this._onDidChange.event; diff --git a/extensions/microsoft-authentication/src/common/publicClientCache.ts b/extensions/microsoft-authentication/src/common/publicClientCache.ts index 9e54e0a2aad..acc8ba0d307 100644 --- a/extensions/microsoft-authentication/src/common/publicClientCache.ts +++ b/extensions/microsoft-authentication/src/common/publicClientCache.ts @@ -5,8 +5,7 @@ import type { AccountInfo, AuthenticationResult, InteractiveRequest, RefreshTokenRequest, SilentFlowRequest } from '@azure/msal-node'; import type { Disposable, Event } from 'vscode'; -export interface ICachedPublicClientApplication extends Disposable { - initialize(): Promise; +export interface ICachedPublicClientApplication { onDidAccountsChange: Event<{ added: AccountInfo[]; changed: AccountInfo[]; deleted: AccountInfo[] }>; onDidRemoveLastAccount: Event; acquireTokenSilent(request: SilentFlowRequest): Promise; diff --git a/extensions/microsoft-authentication/src/extensionV2.ts b/extensions/microsoft-authentication/src/extensionV2.ts index 9610af37977..978603ad132 100644 --- a/extensions/microsoft-authentication/src/extensionV2.ts +++ b/extensions/microsoft-authentication/src/extensionV2.ts @@ -49,14 +49,13 @@ async function initMicrosoftSovereignCloudAuthProvider( return undefined; } - const authProvider = new MsalAuthProvider( + const authProvider = await MsalAuthProvider.create( context, new MicrosoftSovereignCloudAuthenticationTelemetryReporter(context.extension.packageJSON.aiKey), window.createOutputChannel(l10n.t('Microsoft Sovereign Cloud Authentication'), { log: true }), uriHandler, env ); - await authProvider.initialize(); const disposable = authentication.registerAuthenticationProvider( 'microsoft-sovereign-cloud', authProviderName, @@ -70,13 +69,12 @@ async function initMicrosoftSovereignCloudAuthProvider( export async function activate(context: ExtensionContext, mainTelemetryReporter: MicrosoftAuthenticationTelemetryReporter) { const uriHandler = new UriEventHandler(); context.subscriptions.push(uriHandler); - const authProvider = new MsalAuthProvider( + const authProvider = await MsalAuthProvider.create( context, mainTelemetryReporter, Logger, uriHandler ); - await authProvider.initialize(); context.subscriptions.push(authentication.registerAuthenticationProvider( 'microsoft', 'Microsoft', diff --git a/extensions/microsoft-authentication/src/node/authProvider.ts b/extensions/microsoft-authentication/src/node/authProvider.ts index 1956def085c..a26008fb780 100644 --- a/extensions/microsoft-authentication/src/node/authProvider.ts +++ b/extensions/microsoft-authentication/src/node/authProvider.ts @@ -7,7 +7,7 @@ import { AuthenticationGetSessionOptions, AuthenticationProvider, Authentication import { Environment } from '@azure/ms-rest-azure-env'; import { CachedPublicClientApplicationManager } from './publicClientCache'; import { UriEventHandler } from '../UriEventHandler'; -import { ICachedPublicClientApplication } from '../common/publicClientCache'; +import { ICachedPublicClientApplication, ICachedPublicClientApplicationManager } from '../common/publicClientCache'; import { MicrosoftAccountType, MicrosoftAuthenticationTelemetryReporter } from '../common/telemetryReporter'; import { ScopeData } from '../common/scopeData'; import { EventBufferer } from '../common/event'; @@ -22,7 +22,6 @@ const MSA_PASSTHRU_TID = 'f8cdef31-a31e-4b4a-93e4-5f571e91255a'; export class MsalAuthProvider implements AuthenticationProvider { private readonly _disposables: { dispose(): void }[]; - private readonly _publicClientManager: CachedPublicClientApplicationManager; private readonly _eventBufferer = new EventBufferer(); /** @@ -43,15 +42,15 @@ export class MsalAuthProvider implements AuthenticationProvider { */ onDidChangeSessions = this._onDidChangeSessionsEmitter.event; - constructor( + private constructor( private readonly _context: ExtensionContext, private readonly _telemetryReporter: MicrosoftAuthenticationTelemetryReporter, private readonly _logger: LogOutputChannel, private readonly _uriHandler: UriEventHandler, + private readonly _publicClientManager: ICachedPublicClientApplicationManager, private readonly _env: Environment = Environment.AzureCloud ) { this._disposables = _context.subscriptions; - this._publicClientManager = new CachedPublicClientApplicationManager(_context.secrets, this._logger, this._env.name); const accountChangeEvent = this._eventBufferer.wrapEvent( this._publicClientManager.onDidAccountsChange, (last, newEvent) => { @@ -76,11 +75,24 @@ export class MsalAuthProvider implements AuthenticationProvider { )(e => this._handleAccountChange(e)); this._disposables.push( this._onDidChangeSessionsEmitter, - this._publicClientManager, accountChangeEvent ); } + static async create( + context: ExtensionContext, + telemetryReporter: MicrosoftAuthenticationTelemetryReporter, + logger: LogOutputChannel, + uriHandler: UriEventHandler, + env: Environment = Environment.AzureCloud + ): Promise { + const publicClientManager = await CachedPublicClientApplicationManager.create(context.secrets, logger, env.name); + context.subscriptions.push(publicClientManager); + const authProvider = new MsalAuthProvider(context, telemetryReporter, logger, uriHandler, publicClientManager, env); + await authProvider.initialize(); + return authProvider; + } + /** * Migrate sessions from the old secret storage to MSAL. * TODO: MSAL Migration. Remove this when we remove the old flow. @@ -109,9 +121,7 @@ export class MsalAuthProvider implements AuthenticationProvider { } } - async initialize(): Promise { - await this._eventBufferer.bufferEventsAsync(() => this._publicClientManager.initialize()); - + private async initialize(): Promise { if (!this._context.globalState.get('msalMigration', false)) { await this._migrateSessions(); } diff --git a/extensions/microsoft-authentication/src/node/cachedPublicClientApplication.ts b/extensions/microsoft-authentication/src/node/cachedPublicClientApplication.ts index 7eb74baac5f..f0f4eb7b9bc 100644 --- a/extensions/microsoft-authentication/src/node/cachedPublicClientApplication.ts +++ b/extensions/microsoft-authentication/src/node/cachedPublicClientApplication.ts @@ -39,7 +39,7 @@ export class CachedPublicClientApplication implements ICachedPublicClientApplica //#endregion - constructor( + private constructor( private readonly _clientId: string, private readonly _secretStorage: SecretStorage, private readonly _accountAccess: IAccountAccess, @@ -47,7 +47,7 @@ export class CachedPublicClientApplication implements ICachedPublicClientApplica ) { const loggerOptions = new MsalLoggerOptions(_logger); const nativeBrokerPlugin = new NativeBrokerPlugin(); - this._isBrokerAvailable = nativeBrokerPlugin.isBrokerAvailable ?? false; + this._isBrokerAvailable = nativeBrokerPlugin.isBrokerAvailable; this._pca = new PublicClientApplication({ auth: { clientId: _clientId }, system: { @@ -63,17 +63,26 @@ export class CachedPublicClientApplication implements ICachedPublicClientApplica this._disposable = Disposable.from( this._registerOnSecretStorageChanged(), this._onDidAccountsChangeEmitter, - this._onDidRemoveLastAccountEmitter + this._onDidRemoveLastAccountEmitter, + this._secretStorageCachePlugin ); } get accounts(): AccountInfo[] { return this._accounts; } get clientId(): string { return this._clientId; } - async initialize(): Promise { - if (this._isBrokerAvailable) { - await this._accountAccess.initialize(); - } + static async create( + clientId: string, + secretStorage: SecretStorage, + accountAccess: IAccountAccess, + logger: LogOutputChannel + ): Promise { + const app = new CachedPublicClientApplication(clientId, secretStorage, accountAccess, logger); + await app.initialize(); + return app; + } + + private async initialize(): Promise { await this._sequencer.queue(() => this._update()); } @@ -178,7 +187,15 @@ export class CachedPublicClientApplication implements ICachedPublicClientApplica */ async acquireTokenByRefreshToken(request: RefreshTokenRequest): Promise { this._logger.debug(`[acquireTokenByRefreshToken] [${this._clientId}] [${request.authority}] [${request.scopes.join(' ')}]`); - const result = await this._sequencer.queue(() => this._pca.acquireTokenByRefreshToken(request)); + const result = await this._sequencer.queue(async () => { + const result = await this._pca.acquireTokenByRefreshToken(request); + // Force an update so that the account cache is updated. + // TODO:@TylerLeonhardt The problem is, we use the sequencer for + // change events but we _don't_ use it for the accounts cache. + // We should probably use it for the accounts cache as well. + await this._update(); + return result; + }); if (result) { // this._setupRefresh(result); if (this._isBrokerAvailable && result.account) { diff --git a/extensions/microsoft-authentication/src/node/publicClientCache.ts b/extensions/microsoft-authentication/src/node/publicClientCache.ts index 7b46406c16e..16ccb80321f 100644 --- a/extensions/microsoft-authentication/src/node/publicClientCache.ts +++ b/extensions/microsoft-authentication/src/node/publicClientCache.ts @@ -20,37 +20,46 @@ export class CachedPublicClientApplicationManager implements ICachedPublicClient private readonly _pcaDisposables = new Map(); private _disposable: Disposable; - private _pcasSecretStorage: PublicClientApplicationsSecretStorage; - private _accountAccess: IAccountAccess | undefined; private readonly _onDidAccountsChangeEmitter = new EventEmitter<{ added: AccountInfo[]; changed: AccountInfo[]; deleted: AccountInfo[] }>(); readonly onDidAccountsChange = this._onDidAccountsChangeEmitter.event; - constructor( + private constructor( + private readonly _pcasSecretStorage: IPublicClientApplicationSecretStorage, + private readonly _accountAccess: IAccountAccess, private readonly _secretStorage: SecretStorage, private readonly _logger: LogOutputChannel, - private readonly _cloudName: string + disposables: Disposable[] ) { - this._pcasSecretStorage = new PublicClientApplicationsSecretStorage(_secretStorage, _cloudName); this._disposable = Disposable.from( - this._pcasSecretStorage, + ...disposables, this._registerSecretStorageHandler(), this._onDidAccountsChangeEmitter ); } + static async create( + secretStorage: SecretStorage, + logger: LogOutputChannel, + cloudName: string + ): Promise { + const pcasSecretStorage = await PublicClientApplicationsSecretStorage.create(secretStorage, cloudName); + // TODO: Remove the migrations in a version + const migrations = await pcasSecretStorage.getOldValue(); + const accountAccess = await ScopedAccountAccess.create(secretStorage, cloudName, logger, migrations); + const manager = new CachedPublicClientApplicationManager(pcasSecretStorage, accountAccess, secretStorage, logger, [pcasSecretStorage, accountAccess]); + await manager.initialize(); + return manager; + } + private _registerSecretStorageHandler() { return this._pcasSecretStorage.onDidChange(() => this._handleSecretStorageChange()); } - async initialize() { + private async initialize() { this._logger.debug('[initialize] Initializing PublicClientApplicationManager'); let clientIds: string[] | undefined; try { - await this._pcasSecretStorage.initialize(); - //TODO: Remove this in a version - const migrations = await this._pcasSecretStorage.getOldValue(); - this._accountAccess = new ScopedAccountAccess(this._secretStorage, this._cloudName, this._logger, migrations); clientIds = await this._pcasSecretStorage.get(); } catch (e) { // data is corrupted @@ -124,14 +133,12 @@ export class CachedPublicClientApplicationManager implements ICachedPublicClient this._logger.error(`[getOrCreate] [${clientId}] Error migrating refresh token:`, e); } } - // reinitialize the PCA so the account is properly cached - await pca.initialize(); } return pca; } private async _doCreatePublicClientApplication(clientId: string): Promise { - const pca = new CachedPublicClientApplication(clientId, this._secretStorage, this._accountAccess!, this._logger); + const pca = await CachedPublicClientApplication.create(clientId, this._secretStorage, this._accountAccess, this._logger); this._pcas.set(clientId, pca); const disposable = Disposable.from( pca, @@ -147,8 +154,10 @@ export class CachedPublicClientApplicationManager implements ICachedPublicClient }) ); this._pcaDisposables.set(clientId, disposable); - // Intialize the PCA after the `onDidAccountsChange` is set so we get initial state. - await pca.initialize(); + // Fire for the initial state and only if accounts exist + if (pca.accounts.length > 0) { + this._onDidAccountsChangeEmitter.fire({ added: pca.accounts, changed: [], deleted: [] }); + } return pca; } @@ -205,7 +214,15 @@ export class CachedPublicClientApplicationManager implements ICachedPublicClient } } -class PublicClientApplicationsSecretStorage { +interface IPublicClientApplicationSecretStorage { + get(): Promise; + getOldValue(): Promise<{ clientId: string; authority: string }[] | undefined>; + store(value: string[]): Thenable; + delete(): Thenable; + onDidChange: Event; +} + +class PublicClientApplicationsSecretStorage implements IPublicClientApplicationSecretStorage, Disposable { private _disposable: Disposable; private readonly _onDidChangeEmitter = new EventEmitter; @@ -214,7 +231,7 @@ class PublicClientApplicationsSecretStorage { private readonly _oldKey = `publicClientApplications-${this._cloudName}`; private readonly _key = `publicClients-${this._cloudName}`; - constructor(private readonly _secretStorage: SecretStorage, private readonly _cloudName: string) { + private constructor(private readonly _secretStorage: SecretStorage, private readonly _cloudName: string) { this._disposable = Disposable.from( this._onDidChangeEmitter, this._secretStorage.onDidChange(e => { @@ -225,11 +242,17 @@ class PublicClientApplicationsSecretStorage { ); } + static async create(secretStorage: SecretStorage, cloudName: string): Promise { + const storage = new PublicClientApplicationsSecretStorage(secretStorage, cloudName); + await storage.initialize(); + return storage; + } + /** * Runs the migration. * TODO: Remove this after a version. */ - async initialize() { + private async initialize() { const oldValue = await this.getOldValue(); if (!oldValue) { return;