From a9240f38f720cc3a7c95f7c45a1de9f05f5d77cc Mon Sep 17 00:00:00 2001 From: Amir Omidi Date: Fri, 27 Mar 2020 13:18:27 -0700 Subject: [PATCH] Fix a few problems with new Azure auth (#9760) * Fix a few problems * Fix bug * One resource per line * Dispose before cleaning * Dispose the event handler * Dispose webserver --- .../src/account-provider/auths/azureAuth.ts | 11 ++++++--- .../auths/azureAuthCodeGrant.ts | 24 +++++++++---------- .../account-provider/auths/azureDeviceCode.ts | 9 ++++--- .../account-provider/azureAccountProvider.ts | 18 +++++++++----- .../azureAccountProviderService.ts | 22 +++++++++++++++-- .../src/account-provider/providerSettings.ts | 10 ++++---- .../account-provider/auths/azureAuth.test.ts | 2 +- 7 files changed, 64 insertions(+), 32 deletions(-) diff --git a/extensions/azurecore/src/account-provider/auths/azureAuth.ts b/extensions/azurecore/src/account-provider/auths/azureAuth.ts index 6c9c46a541..4e6dfa7e14 100644 --- a/extensions/azurecore/src/account-provider/auths/azureAuth.ts +++ b/extensions/azurecore/src/account-provider/auths/azureAuth.ts @@ -91,7 +91,7 @@ export interface TokenClaims { // https://docs.microsoft.com/en-us/azure/active- export type TokenRefreshResponse = { accessToken: AccessToken, refreshToken: RefreshToken, tokenClaims: TokenClaims }; -export abstract class AzureAuth { +export abstract class AzureAuth implements vscode.Disposable { protected readonly memdb = new MemoryDatabase(); protected readonly WorkSchoolAccountType: string = 'work_school'; @@ -110,6 +110,7 @@ export abstract class AzureAuth { protected readonly metadata: AzureAccountProviderMetadata, protected readonly tokenCache: SimpleTokenCache, protected readonly context: vscode.ExtensionContext, + protected readonly uriEventEmitter: vscode.EventEmitter, protected readonly authType: AzureAuthType, public readonly userFriendlyName: string ) { @@ -119,8 +120,11 @@ export abstract class AzureAuth { this.clientId = this.metadata.settings.clientId; this.resources = [ - this.metadata.settings.armResource, this.metadata.settings.sqlResource, - this.metadata.settings.graphResource, this.metadata.settings.ossRdbmsResource + this.metadata.settings.armResource, + this.metadata.settings.sqlResource, + this.metadata.settings.graphResource, + this.metadata.settings.ossRdbmsResource, + this.metadata.settings.azureKeyVaultResource ]; this.scopes = [...this.metadata.settings.scopes]; @@ -131,6 +135,7 @@ export abstract class AzureAuth { public abstract async autoOAuthCancelled(): Promise; + public dispose() { } public async refreshAccess(account: azdata.Account): Promise { const response = await this.getCachedToken(account.key); diff --git a/extensions/azurecore/src/account-provider/auths/azureAuthCodeGrant.ts b/extensions/azurecore/src/account-provider/auths/azureAuthCodeGrant.ts index 2eabce3cd8..ff26b261b3 100644 --- a/extensions/azurecore/src/account-provider/auths/azureAuthCodeGrant.ts +++ b/extensions/azurecore/src/account-provider/auths/azureAuthCodeGrant.ts @@ -30,12 +30,6 @@ import { SimpleWebServer } from '../utils/simpleWebServer'; import { SimpleTokenCache } from '../simpleTokenCache'; const localize = nls.loadMessageBundle(); -class UriEventHandler extends vscode.EventEmitter implements vscode.UriHandler { - public handleUri(uri: vscode.Uri) { - this.fire(uri); - } -} - function parseQuery(uri: vscode.Uri) { return uri.query.split('&').reduce((prev: any, current) => { const queryString = current.split('='); @@ -52,14 +46,14 @@ interface AuthCodeResponse { export class AzureAuthCodeGrant extends AzureAuth { private static readonly USER_FRIENDLY_NAME: string = localize('azure.azureAuthCodeGrantName', "Azure Auth Code Grant"); private server: SimpleWebServer; - private readonly _uriHandler: UriEventHandler; - constructor(metadata: AzureAccountProviderMetadata, + constructor( + metadata: AzureAccountProviderMetadata, tokenCache: SimpleTokenCache, - context: vscode.ExtensionContext) { - super(metadata, tokenCache, context, AzureAuthType.AuthCodeGrant, AzureAuthCodeGrant.USER_FRIENDLY_NAME); - this._uriHandler = new UriEventHandler(); - vscode.window.registerUriHandler(this._uriHandler); + context: vscode.ExtensionContext, + uriEventEmitter: vscode.EventEmitter, + ) { + super(metadata, tokenCache, context, uriEventEmitter, AzureAuthType.AuthCodeGrant, AzureAuthCodeGrant.USER_FRIENDLY_NAME); } public async autoOAuthCancelled(): Promise { @@ -146,7 +140,7 @@ export class AzureAuthCodeGrant extends AzureAuth { public async handleCodeResponse(state: string): Promise { let uriEventListener: vscode.Disposable; return new Promise((resolve: (value: any) => void, reject) => { - uriEventListener = this._uriHandler.event(async (uri: vscode.Uri) => { + uriEventListener = this.uriEventEmitter.event(async (uri: vscode.Uri) => { try { const query = parseQuery(uri); const code = query.code; @@ -316,4 +310,8 @@ export class AzureAuthCodeGrant extends AzureAuth { return this.getToken(postData); } + + public dispose() { + this.server?.shutdown().catch(console.error); + } } diff --git a/extensions/azurecore/src/account-provider/auths/azureDeviceCode.ts b/extensions/azurecore/src/account-provider/auths/azureDeviceCode.ts index d42ac1b225..8ab6d77d34 100644 --- a/extensions/azurecore/src/account-provider/auths/azureDeviceCode.ts +++ b/extensions/azurecore/src/account-provider/auths/azureDeviceCode.ts @@ -46,10 +46,13 @@ export class AzureDeviceCode extends AzureAuth { private static readonly USER_FRIENDLY_NAME: string = localize('azure.azureDeviceCodeAuth', "Azure Device Code"); private readonly pageTitle: string; - constructor(metadata: AzureAccountProviderMetadata, + constructor( + metadata: AzureAccountProviderMetadata, tokenCache: SimpleTokenCache, - context: vscode.ExtensionContext) { - super(metadata, tokenCache, context, AzureAuthType.AuthCodeGrant, AzureDeviceCode.USER_FRIENDLY_NAME); + context: vscode.ExtensionContext, + uriEventEmitter: vscode.EventEmitter, + ) { + super(metadata, tokenCache, context, uriEventEmitter, AzureAuthType.AuthCodeGrant, AzureDeviceCode.USER_FRIENDLY_NAME); this.pageTitle = localize('addAccount', "Add {0} account", this.metadata.displayName); } diff --git a/extensions/azurecore/src/account-provider/azureAccountProvider.ts b/extensions/azurecore/src/account-provider/azureAccountProvider.ts index 99db70ab72..4a8f997a57 100644 --- a/extensions/azurecore/src/account-provider/azureAccountProvider.ts +++ b/extensions/azurecore/src/account-provider/azureAccountProvider.ts @@ -20,7 +20,7 @@ import { AzureDeviceCode } from './auths/azureDeviceCode'; const localize = nls.loadMessageBundle(); -export class AzureAccountProvider implements azdata.AccountProvider { +export class AzureAccountProvider implements azdata.AccountProvider, vscode.Disposable { private static readonly CONFIGURATION_SECTION = 'accounts.azure.auth'; private readonly authMappings = new Map(); private initComplete: Deferred; @@ -30,23 +30,29 @@ export class AzureAccountProvider implements azdata.AccountProvider { metadata: AzureAccountProviderMetadata, tokenCache: SimpleTokenCache, context: vscode.ExtensionContext, + uriEventHandler: vscode.EventEmitter, private readonly forceDeviceCode: boolean = false ) { vscode.workspace.onDidChangeConfiguration((changeEvent) => { const impact = changeEvent.affectsConfiguration(AzureAccountProvider.CONFIGURATION_SECTION); if (impact === true) { - this.handleAuthMapping(metadata, tokenCache, context); + this.handleAuthMapping(metadata, tokenCache, context, uriEventHandler); } }); - this.handleAuthMapping(metadata, tokenCache, context); + this.handleAuthMapping(metadata, tokenCache, context, uriEventHandler); + } + + dispose() { + this.authMappings.forEach(x => x.dispose()); } clearTokenCache(): Thenable { return this.getAuthMethod().deleteAllCache(); } - private handleAuthMapping(metadata: AzureAccountProviderMetadata, tokenCache: SimpleTokenCache, context: vscode.ExtensionContext) { + private handleAuthMapping(metadata: AzureAccountProviderMetadata, tokenCache: SimpleTokenCache, context: vscode.ExtensionContext, uriEventHandler: vscode.EventEmitter) { + this.authMappings.forEach(m => m.dispose()); this.authMappings.clear(); const configuration = vscode.workspace.getConfiguration(AzureAccountProvider.CONFIGURATION_SECTION); @@ -54,10 +60,10 @@ export class AzureAccountProvider implements azdata.AccountProvider { const deviceCodeMethod: boolean = configuration.get('deviceCode'); if (codeGrantMethod === true && !this.forceDeviceCode) { - this.authMappings.set(AzureAuthType.AuthCodeGrant, new AzureAuthCodeGrant(metadata, tokenCache, context)); + this.authMappings.set(AzureAuthType.AuthCodeGrant, new AzureAuthCodeGrant(metadata, tokenCache, context, uriEventHandler)); } if (deviceCodeMethod === true || this.forceDeviceCode) { - this.authMappings.set(AzureAuthType.DeviceCode, new AzureDeviceCode(metadata, tokenCache, context)); + this.authMappings.set(AzureAuthType.DeviceCode, new AzureDeviceCode(metadata, tokenCache, context, uriEventHandler)); } } diff --git a/extensions/azurecore/src/account-provider/azureAccountProviderService.ts b/extensions/azurecore/src/account-provider/azureAccountProviderService.ts index 7b56e7d510..9307dde7af 100644 --- a/extensions/azurecore/src/account-provider/azureAccountProviderService.ts +++ b/extensions/azurecore/src/account-provider/azureAccountProviderService.ts @@ -15,6 +15,12 @@ import * as loc from '../localizedConstants'; let localize = nls.loadMessageBundle(); +class UriEventHandler extends vscode.EventEmitter implements vscode.UriHandler { + public handleUri(uri: vscode.Uri) { + this.fire(uri); + } +} + export class AzureAccountProviderService implements vscode.Disposable { // CONSTANTS /////////////////////////////////////////////////////////////// private static CommandClearTokenCache = 'accounts.clearTokenCache'; @@ -22,12 +28,14 @@ export class AzureAccountProviderService implements vscode.Disposable { private static CredentialNamespace = 'azureAccountProviderCredentials'; // MEMBER VARIABLES //////////////////////////////////////////////////////// + private _disposables: vscode.Disposable[] = []; private _accountDisposals: { [accountProviderId: string]: vscode.Disposable }; private _accountProviders: { [accountProviderId: string]: azdata.AccountProvider }; private _credentialProvider: azdata.CredentialProvider; private _configChangePromiseChain: Thenable; private _currentConfig: vscode.WorkspaceConfiguration; private _event: events.EventEmitter; + private readonly _uriEventHandler: UriEventHandler; constructor(private _context: vscode.ExtensionContext, private _userStoragePath: string) { this._accountDisposals = {}; @@ -35,6 +43,9 @@ export class AzureAccountProviderService implements vscode.Disposable { this._configChangePromiseChain = Promise.resolve(); this._currentConfig = null; this._event = new events.EventEmitter(); + + this._uriEventHandler = new UriEventHandler(); + this._disposables.push(vscode.window.registerUriHandler(this._uriEventHandler)); } // PUBLIC METHODS ////////////////////////////////////////////////////// @@ -65,7 +76,14 @@ export class AzureAccountProviderService implements vscode.Disposable { }); } - public dispose() { } + public dispose() { + while (this._disposables.length) { + const item = this._disposables.pop(); + if (item) { + item.dispose(); + } + } + } // PRIVATE HELPERS ///////////////////////////////////////////////////// private onClearTokenCache(): Thenable { @@ -133,7 +151,7 @@ export class AzureAccountProviderService implements vscode.Disposable { await simpleTokenCache.init(); const isSaw: boolean = vscode.env.appName.toLowerCase().indexOf('saw') > 0; - let accountProvider = new AzureAccountProvider(provider.metadata as AzureAccountProviderMetadata, simpleTokenCache, this._context, isSaw); + let accountProvider = new AzureAccountProvider(provider.metadata as AzureAccountProviderMetadata, simpleTokenCache, this._context, this._uriEventHandler, isSaw); this._accountProviders[provider.metadata.id] = accountProvider; this._accountDisposals[provider.metadata.id] = azdata.accounts.registerAccountProvider(provider.metadata, accountProvider); diff --git a/extensions/azurecore/src/account-provider/providerSettings.ts b/extensions/azurecore/src/account-provider/providerSettings.ts index cd8e28615f..90c4c48e36 100644 --- a/extensions/azurecore/src/account-provider/providerSettings.ts +++ b/extensions/azurecore/src/account-provider/providerSettings.ts @@ -39,8 +39,9 @@ const publicAzureSettings: ProviderSettings = { azureResourceId: AzureResource.OssRdbms }, azureKeyVaultResource: { - id: 'https://vault.azure.net', - endpoint: 'https://vault.azure.net' + id: 'vault', + endpoint: 'https://vault.azure.net', + azureResourceId: AzureResource.AzureKeyVault }, redirectUri: 'https://vscode-redirect.azurewebsites.net/', scopes: [ @@ -82,8 +83,9 @@ const usGovAzureSettings: ProviderSettings = { azureResourceId: AzureResource.OssRdbms }, azureKeyVaultResource: { - id: 'https://vault.usgovcloudapi.net', - endpoint: 'https://vault.usgovcloudapi.net' + id: 'vault', + endpoint: 'https://vault.usgovcloudapi.net', + azureResourceId: AzureResource.AzureKeyVault }, redirectUri: 'https://vscode-redirect.azurewebsites.net/', scopes: [ diff --git a/extensions/azurecore/src/test/account-provider/auths/azureAuth.test.ts b/extensions/azurecore/src/test/account-provider/auths/azureAuth.test.ts index db47317606..604c70fc50 100644 --- a/extensions/azurecore/src/test/account-provider/auths/azureAuth.test.ts +++ b/extensions/azurecore/src/test/account-provider/auths/azureAuth.test.ts @@ -48,7 +48,7 @@ describe('AccountProvider.AzureAuth', function (): void { beforeEach(async function (): Promise { const tokenCache = new SimpleTokenCache('testTokenService', os.tmpdir(), true, new CredentialsTestProvider()); await tokenCache.init(); - baseAuth = new BasicAzureAuth(providerSettings[0].metadata, tokenCache, undefined, AzureAuthType.AuthCodeGrant, 'Auth Code Grant'); + baseAuth = new BasicAzureAuth(providerSettings[0].metadata, tokenCache, undefined, undefined, AzureAuthType.AuthCodeGrant, 'Auth Code Grant'); }); it('Basic token set and get', async function (): Promise {