ckaczor/azuredatastudio
1
0
Fork 0
mirror of https://github.com/ckaczor/azuredatastudio.git synced 2026-02-16 10:58:30 -05:00
Code Issues Packages Projects Releases Wiki Activity

Fix ADAL to MSAL transition of connections and account list (#21425)

Browse Source
This commit is contained in:
Cheena Malhotra
2022-12-15 16:38:45 -08:00
committed by GitHub
parent a5e2e77042
commit 112de46723
4 changed files with 72 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

47
extensions/azurecore/src/account-provider/auths/azureAuth.ts
View File

@@ -318,23 +318,12 @@ export abstract class AzureAuth implements vscode.Disposable {
* @returns The authentication result, including the access token
*/
public async getTokenMsal(accountId: string, azureResource: azdata.AzureResource, tenantId: string): Promise<AuthenticationResult | null> {
const cache = this.clientApplication.getTokenCache();
if (!cache) {
Logger.error('Error: Could not fetch token cache.');
return null;
}
const resource = this.resources.find(s => s.azureResourceId === azureResource);
if (!resource) {
Logger.error(`Error: Could not fetch the azure resource ${azureResource} `);
return null;
}
let account: AccountInfo | null;
// if the accountId is a home ID, it will include a "." character
if (accountId.includes(".")) {
account = await cache.getAccountByHomeId(accountId);
} else {
account = await cache.getAccountByLocalId(accountId);
}
let account: AccountInfo | null = await this.getAccountFromMsalCache(accountId);
if (!account) {
Logger.error('Error: Could not fetch account when acquiring token');
return null;
@@ -374,6 +363,23 @@ export abstract class AzureAuth implements vscode.Disposable {
}
}
public async getAccountFromMsalCache(accountId: string): Promise<AccountInfo | null> {
const cache = this.clientApplication.getTokenCache();
if (!cache) {
Logger.error('Error: Could not fetch token cache.');
return null;
}
let account: AccountInfo | null = null;
// if the accountId is a home ID, it will include a "." character
if (accountId.includes(".")) {
account = await cache.getAccountByHomeId(accountId);
} else {
account = await cache.getAccountByLocalId(accountId);
}
return account;
}
public async getTokenAdal(tenant: Tenant, resource: Resource, postData: AuthorizationCodePostData | TokenPostData | RefreshTokenPostData): Promise<OAuthTokenResponse | undefined> {
Logger.verbose('Fetching token for tenant {0}', tenant.id);
const tokenUrl = `${this.loginEndpointUrl}${tenant.id}/oauth2/token`;
@@ -817,26 +823,19 @@ export abstract class AzureAuth implements vscode.Disposable {
// remove account based on authLibrary field, accounts added before this field was present will default to
// ADAL method of account removal
if (account.authLibrary === Constants.AuthLibrary.MSAL) {
return this.deleteAccountCacheMsal(account);
return await this.deleteAccountCacheMsal(account);
} else { // fallback to ADAL by default
return this.deleteAccountCacheAdal(account);
return await this.deleteAccountCacheAdal(account);
}
} catch (ex) {
const msg = localize('azure.cacheErrrorRemove', "Error when removing your account from the cache.");
void vscode.window.showErrorMessage(msg);
Logger.error('Error when removing tokens.', ex);
// We need not prompt user for error if token could not be removed from cache.
Logger.error('Error when removing token from cache: ', ex);
}
}
public async deleteAccountCacheMsal(account: azdata.AccountKey): Promise<void> {
const tokenCache = this.clientApplication.getTokenCache();
let msalAccount: AccountInfo | null;
// if the accountId is a home ID, it will include a "." character
if (account.accountId.includes(".")) {
msalAccount = await tokenCache.getAccountByHomeId(account.accountId);
} else {
msalAccount = await tokenCache.getAccountByLocalId(account.accountId);
}
let msalAccount: AccountInfo | null = await this.getAccountFromMsalCache(account.accountId);
if (!msalAccount) {
Logger.error(`MSAL: Unable to find account ${account.accountId} for removal`);
throw Error(`Unable to find account ${account.accountId}`);

45
extensions/azurecore/src/account-provider/azureAccountProvider.ts
View File

@@ -115,7 +115,13 @@ export class AzureAccountProvider implements azdata.AccountProvider, vscode.Disp
} else {
account.isStale = false;
if (this.authLibrary === Constants.AuthLibrary.MSAL) {
// Check MSAL Cache before adding account, to mark it as stale if it is not present in cache
const accountInCache = await azureAuth.getAccountFromMsalCache(account.key.accountId);
if (!accountInCache) {
account.isStale = true;
}
accounts.push(account);
} else { // fallback to ADAL as default
accounts.push(await azureAuth.refreshAccessAdal(account));
}
@@ -137,23 +143,30 @@ export class AzureAccountProvider implements azdata.AccountProvider, vscode.Disp
private async _getAccountSecurityToken(account: AzureAccount, tenantId: string, resource: azdata.AzureResource): Promise<Token | undefined> {
await this.initCompletePromise;
const azureAuth = this.getAuthMethod(account);
Logger.pii(`Getting account security token for ${JSON.stringify(account.key)} (tenant ${tenantId}). Auth Method = ${azureAuth.userFriendlyName}`, [], []);
if (this.authLibrary === Constants.AuthLibrary.MSAL) {
let authResult = await azureAuth?.getTokenMsal(account.key.accountId, resource, tenantId);
if (!authResult || !authResult.account || !authResult.account.idTokenClaims) {
Logger.error(`MSAL: getToken call failed`);
throw Error('Failed to get token');
} else {
const token: Token = {
key: authResult.account.homeAccountId,
token: authResult.accessToken,
tokenType: authResult.tokenType,
expiresOn: authResult.account.idTokenClaims.exp
};
return token;
if (azureAuth) {
Logger.pii(`Getting account security token for ${JSON.stringify(account.key)} (tenant ${tenantId}). Auth Method = ${azureAuth.userFriendlyName}`, [], []);
if (this.authLibrary === Constants.AuthLibrary.MSAL) {
let authResult = await azureAuth.getTokenMsal(account.key.accountId, resource, tenantId);
if (!authResult || !authResult.account || !authResult.account.idTokenClaims) {
Logger.error(`MSAL: getToken call failed`);
throw Error('Failed to get token');
} else {
const token: Token = {
key: authResult.account.homeAccountId,
token: authResult.accessToken,
tokenType: authResult.tokenType,
expiresOn: authResult.account.idTokenClaims.exp
};
return token;
}
} else { // fallback to ADAL as default
return azureAuth.getAccountSecurityTokenAdal(account, tenantId, resource);
}
} else { // fallback to ADAL as default
return azureAuth?.getAccountSecurityTokenAdal(account, tenantId, resource);
} else {
account.isStale = true;
Logger.error(`_getAccountSecurityToken: Authentication method not found for account ${account.displayInfo.displayName}`);
throw Error('Failed to get authentication method, please remove and re-add the account');
}
}