Large cleanup of AzureCore - Introduction of getAccountSecurityToken and deprecation of getSecurityToken (#11446)

* do a large cleanup of azurecore

* Fix tests

* Rework Device Code

* Fix tests

* Fix AE scenario

* Fix firewall rule - clenaup logging

* Shorthand syntax

* Fix firewall tests

* Start on tests for azureAuth

* Add more tests

* Address comments

* Add a few more important tests

* Don't throw error on old code

* Fill in todo
This commit is contained in:
Amir Omidi
2020-07-22 15:03:42 -07:00
committed by GitHub
parent a61b85c9ff
commit 587abd43c2
40 changed files with 1045 additions and 895 deletions

View File

@@ -75,9 +75,13 @@ export class MainThreadAccountManagement extends Disposable implements MainThrea
clear(accountKey: azdata.AccountKey): Thenable<void> {
return self._proxy.$clear(handle, accountKey);
},
getSecurityToken(account: azdata.Account, resource: azdata.AzureResource): Thenable<{}> {
return self._proxy.$getSecurityToken(account, resource);
},
getAccountSecurityToken(account: azdata.Account, tenant: string, resource: azdata.AzureResource): Thenable<{ token: string }> {
return self._proxy.$getAccountSecurityToken(account, tenant, resource);
},
initialize(restoredAccounts: azdata.Account[]): Thenable<azdata.Account[]> {
return self._proxy.$initialize(handle, restoredAccounts);
},

View File

@@ -89,10 +89,7 @@ export class ExtHostAccountManagement extends ExtHostAccountManagementShape {
return Promise.all(promises).then(() => resultAccounts);
}
public $getSecurityToken(account: azdata.Account, resource?: azdata.AzureResource): Thenable<{}> {
if (resource === undefined) {
resource = AzureResource.ResourceManagement;
}
public $getSecurityToken(account: azdata.Account, resource: azdata.AzureResource = AzureResource.ResourceManagement): Thenable<{}> {
return this.$getAllAccounts().then(() => {
for (const handle in this._accounts) {
const providerHandle = parseInt(handle);
@@ -105,6 +102,20 @@ export class ExtHostAccountManagement extends ExtHostAccountManagementShape {
});
}
public $getAccountSecurityToken(account: azdata.Account, tenant: string, resource: azdata.AzureResource = AzureResource.ResourceManagement): Thenable<{ token: string }> {
return this.$getAllAccounts().then(() => {
for (const handle in this._accounts) {
const providerHandle = parseInt(handle);
if (firstIndex(this._accounts[handle], (acct) => acct.key.accountId === account.key.accountId) !== -1) {
return this._withProvider(providerHandle, (provider: azdata.AccountProvider) => provider.getAccountSecurityToken(account, tenant, resource));
}
}
throw new Error(`Account ${account.key.accountId} not found.`);
});
}
public get onDidChangeAccounts(): Event<azdata.DidChangeAccountsParams> {
return this._onDidChangeAccounts.event;
}

View File

@@ -160,6 +160,9 @@ export function createAdsApiFactory(accessor: ServicesAccessor): IAdsExtensionAp
getSecurityToken(account: azdata.Account, resource?: azdata.AzureResource): Thenable<{}> {
return extHostAccountManagement.$getSecurityToken(account, resource);
},
getAccountSecurityToken(account: azdata.Account, tenant: string, resource?: azdata.AzureResource): Thenable<{ token: string }> {
return extHostAccountManagement.$getAccountSecurityToken(account, tenant, resource);
},
onDidChangeAccounts(listener: (e: azdata.DidChangeAccountsParams) => void, thisArgs?: any, disposables?: extHostTypes.Disposable[]) {
return extHostAccountManagement.onDidChangeAccounts(listener, thisArgs, disposables);
}

View File

@@ -31,6 +31,7 @@ export abstract class ExtHostAccountManagementShape {
$autoOAuthCancelled(handle: number): Thenable<void> { throw ni(); }
$clear(handle: number, accountKey: azdata.AccountKey): Thenable<void> { throw ni(); }
$getSecurityToken(account: azdata.Account, resource?: azdata.AzureResource): Thenable<{}> { throw ni(); }
$getAccountSecurityToken(account: azdata.Account, tenant: string, resource?: azdata.AzureResource): Thenable<{ token: string }> { throw ni(); }
$initialize(handle: number, restoredAccounts: azdata.Account[]): Thenable<azdata.Account[]> { throw ni(); }
$prompt(handle: number): Thenable<azdata.Account | azdata.PromptFailedResult> { throw ni(); }
$refresh(handle: number, account: azdata.Account): Thenable<azdata.Account | azdata.PromptFailedResult> { throw ni(); }

View File

@@ -244,6 +244,19 @@ export class AccountManagementService implements IAccountManagementService {
});
}
/**
* Generates a security token by asking the account's provider
* @param account Account to generate security token for
* @param tenant Tenant to generate security token for
* @param resource The resource to get the security token for
* @return Promise to return the security token
*/
public getAccountSecurityToken(account: azdata.Account, tenant: string, resource: azdata.AzureResource): Thenable<{ token: string }> {
return this.doWithProvider(account.key.providerId, provider => {
return provider.provider.getAccountSecurityToken(account, tenant, resource);
});
}
/**
* Removes an account from the account store and clears sensitive data in the provider
* @param accountKey Key for the account to remove

View File

@@ -813,8 +813,8 @@ export class ConnectionManagementService extends Disposable implements IConnecti
const accounts = await this._accountManagementService.getAccounts();
const azureAccounts = accounts.filter(a => a.key.providerId.startsWith('azure'));
if (azureAccounts && azureAccounts.length > 0) {
let accountName = (connection.authenticationType === Constants.azureMFA || connection.authenticationType === Constants.azureMFAAndUser) ? connection.azureAccount : connection.userName;
let account = find(azureAccounts, account => account.key.accountId === accountName);
let accountId = (connection.authenticationType === Constants.azureMFA || connection.authenticationType === Constants.azureMFAAndUser) ? connection.azureAccount : connection.userName;
let account = find(azureAccounts, account => account.key.accountId === accountId);
if (account) {
this._logService.debug(`Getting security token for Azure account ${account.key.accountId}`);
if (account.isStale) {
@@ -827,26 +827,17 @@ export class ConnectionManagementService extends Disposable implements IConnecti
return false;
}
}
const tokensByTenant = await this._accountManagementService.getSecurityToken(account, azureResource);
this._logService.debug(`Got tokens for tenants [${Object.keys(tokensByTenant).join(',')}]`);
let token: string;
const tenantId = connection.azureTenantId;
if (tenantId && tokensByTenant[tenantId]) {
token = tokensByTenant[tenantId].token;
} else {
this._logService.debug(`No security token found for specific tenant ${tenantId} - falling back to first one`);
const tokens = values(tokensByTenant);
if (tokens.length === 0) {
this._logService.info(`No security tokens found for account`);
return false;
}
token = tokens[0].token;
const token = await this._accountManagementService.getAccountSecurityToken(account, tenantId, azureResource);
this._logService.debug(`Got token for tenant ${token}`);
if (!token) {
this._logService.info(`No security tokens found for account`);
}
connection.options['azureAccountToken'] = token;
connection.options['azureAccountToken'] = token.token;
connection.options['password'] = '';
return true;
} else {
this._logService.info(`Could not find Azure account with name ${accountName}`);
this._logService.info(`Could not find Azure account with name ${accountId}`);
}
} else {
this._logService.info(`Could not find any Azure accounts from accounts : [${accounts.map(a => `${a.key.accountId} (${a.key.providerId})`).join(',')}]`);

View File

@@ -6,7 +6,7 @@
import 'vs/css!./media/sqlConnection';
import { Button } from 'sql/base/browser/ui/button/button';
import { SelectBox } from 'sql/base/browser/ui/selectBox/selectBox';
import { SelectBox, SelectOptionItemSQL } from 'sql/base/browser/ui/selectBox/selectBox';
import { Checkbox } from 'sql/base/browser/ui/checkbox/checkbox';
import { InputBox } from 'sql/base/browser/ui/inputBox/inputBox';
import * as DialogHelper from 'sql/workbench/browser/modal/dialogHelper';
@@ -520,12 +520,18 @@ export class ConnectionWidget extends lifecycle.Disposable {
let oldSelection = this._azureAccountDropdown.value;
const accounts = await this._accountManagementService.getAccounts();
this._azureAccountList = accounts.filter(a => a.key.providerId.startsWith('azure'));
let accountDropdownOptions = this._azureAccountList.map(account => account.displayInfo.displayName);
let accountDropdownOptions: SelectOptionItemSQL[] = this._azureAccountList.map(account => {
return {
text: account.displayInfo.displayName,
value: account.key.accountId
} as SelectOptionItemSQL;
});
if (accountDropdownOptions.length === 0) {
// If there are no accounts add a blank option so that add account isn't automatically selected
accountDropdownOptions.unshift('');
accountDropdownOptions.unshift({ text: '', value: '' });
}
accountDropdownOptions.push(this._addAzureAccountMessage);
accountDropdownOptions.push({ text: this._addAzureAccountMessage, value: this._addAzureAccountMessage });
this._azureAccountDropdown.setOptions(accountDropdownOptions);
this._azureAccountDropdown.selectWithOptionName(oldSelection);
}

View File

@@ -1299,6 +1299,7 @@ suite('SQL ConnectionManagementService tests', () => {
let servername = 'test-database.database.windows.net';
azureConnectionProfile.serverName = servername;
let providerId = 'azure_PublicCloud';
azureConnectionProfile.azureTenantId = 'testTenant';
// Set up the account management service to return a token for the given user
accountManagementService.setup(x => x.getAccountsForProvider(TypeMoq.It.isAny())).returns(providerId => Promise.resolve<azdata.Account[]>([
@@ -1327,10 +1328,9 @@ suite('SQL ConnectionManagementService tests', () => {
]);
});
let testToken = 'testToken';
accountManagementService.setup(x => x.getSecurityToken(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve({
azure_publicCloud: {
token: testToken
}
accountManagementService.setup(x => x.getAccountSecurityToken(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve({
token: testToken,
tokenType: 'Bearer'
}));
connectionStore.setup(x => x.addSavedPassword(TypeMoq.It.is(profile => profile.authenticationType === 'AzureMFA'))).returns(profile => Promise.resolve({
profile: profile,
@@ -1384,11 +1384,8 @@ suite('SQL ConnectionManagementService tests', () => {
]);
});
let testToken = 'testToken';
let returnedTokens = {};
returnedTokens['azure_publicCloud'] = { token: 'badToken' };
returnedTokens[azureTenantId] = { token: testToken };
accountManagementService.setup(x => x.getSecurityToken(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(returnedTokens));
let returnedToken = { token: 'testToken', tokenType: 'Bearer' };
accountManagementService.setup(x => x.getAccountSecurityToken(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(returnedToken));
connectionStore.setup(x => x.addSavedPassword(TypeMoq.It.is(profile => profile.authenticationType === 'AzureMFA'))).returns(profile => Promise.resolve({
profile: profile,
savedCred: false
@@ -1399,7 +1396,7 @@ suite('SQL ConnectionManagementService tests', () => {
// Then the returned profile has the account token set corresponding to the requested tenant
assert.equal(profileWithCredentials.userName, azureConnectionProfile.userName);
assert.equal(profileWithCredentials.options['azureAccountToken'], testToken);
assert.equal(profileWithCredentials.options['azureAccountToken'], returnedToken.token);
});
test('getConnections test', () => {

View File

@@ -60,7 +60,12 @@ export class FirewallRuleDialogController {
private async handleOnCreateFirewallRule(): Promise<void> {
const resourceProviderId = this._resourceProviderId;
try {
const securityTokenMappings = await this._accountManagementService.getSecurityToken(this._firewallRuleDialog.viewModel.selectedAccount!, AzureResource.ResourceManagement);
const tenantId = this._connection.azureTenantId;
const token = await this._accountManagementService.getAccountSecurityToken(this._firewallRuleDialog.viewModel.selectedAccount!, tenantId, AzureResource.ResourceManagement);
const securityTokenMappings = {
[tenantId]: token
};
const firewallRuleInfo: azdata.FirewallRuleInfo = {
startIpAddress: this._firewallRuleDialog.viewModel.isIPAddressSelected ? this._firewallRuleDialog.viewModel.defaultIPAddress : this._firewallRuleDialog.viewModel.fromSubnetIPRange,
endIpAddress: this._firewallRuleDialog.viewModel.isIPAddressSelected ? this._firewallRuleDialog.viewModel.defaultIPAddress : this._firewallRuleDialog.viewModel.toSubnetIPRange,

View File

@@ -92,7 +92,8 @@ suite('Firewall rule dialog controller tests', () => {
providerName: mssqlProviderName,
options: {},
saveProfile: true,
id: ''
id: '',
azureTenantId: 'someTenant'
};
});
@@ -137,7 +138,7 @@ suite('Firewall rule dialog controller tests', () => {
// Then: it should get security token from account management service and call create firewall rule in resource provider
await deferredPromise;
mockAccountManagementService.verify(x => x.getSecurityToken(TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.once());
mockAccountManagementService.verify(x => x.getAccountSecurityToken(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.once());
mockResourceProvider.verify(x => x.createFirewallRule(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.once());
mockFirewallRuleDialog.verify(x => x.close(), TypeMoq.Times.once());
mockFirewallRuleDialog.verify(x => x.onServiceComplete(), TypeMoq.Times.once());
@@ -164,7 +165,7 @@ suite('Firewall rule dialog controller tests', () => {
// Then: it should get security token from account management service and an error dialog should have been opened
await deferredPromise;
mockAccountManagementService.verify(x => x.getSecurityToken(TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.once());
mockAccountManagementService.verify(x => x.getAccountSecurityToken(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.once());
mockErrorMessageService.verify(x => x.showDialog(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.once());
mockResourceProvider.verify(x => x.createFirewallRule(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.never());
});
@@ -191,7 +192,7 @@ suite('Firewall rule dialog controller tests', () => {
// Then: it should get security token from account management service and an error dialog should have been opened
await deferredPromise;
mockAccountManagementService.verify(x => x.getSecurityToken(TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.once());
mockAccountManagementService.verify(x => x.getAccountSecurityToken(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.once());
mockResourceProvider.verify(x => x.createFirewallRule(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.once());
mockErrorMessageService.verify(x => x.showDialog(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.once());
});
@@ -217,7 +218,7 @@ suite('Firewall rule dialog controller tests', () => {
// Then: it should get security token from account management service and an error dialog should have been opened
await deferredPromise;
mockAccountManagementService.verify(x => x.getSecurityToken(TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.once());
mockAccountManagementService.verify(x => x.getAccountSecurityToken(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.once());
mockResourceProvider.verify(x => x.createFirewallRule(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.once());
mockErrorMessageService.verify(x => x.showDialog(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.once());
});
@@ -226,8 +227,8 @@ suite('Firewall rule dialog controller tests', () => {
function getMockAccountManagementService(resolveSecurityToken: boolean): TypeMoq.Mock<TestAccountManagementService> {
let accountManagementTestService = new TestAccountManagementService();
let mockAccountManagementService = TypeMoq.Mock.ofInstance(accountManagementTestService);
mockAccountManagementService.setup(x => x.getSecurityToken(TypeMoq.It.isAny(), TypeMoq.It.isAny()))
.returns(() => resolveSecurityToken ? Promise.resolve({}) : Promise.reject(null));
mockAccountManagementService.setup(x => x.getAccountSecurityToken(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny()))
.returns(() => resolveSecurityToken ? Promise.resolve({ token: 'token' }) : Promise.reject(null));
return mockAccountManagementService;
}

View File

@@ -442,6 +442,8 @@ function getMockAccountManagementService(accounts: azdata.Account[]): TypeMoq.Mo
.returns(() => Promise.resolve(accounts));
mockAccountManagementService.setup(x => x.getSecurityToken(TypeMoq.It.isValue(accounts[0]), TypeMoq.It.isAny()))
.returns(() => Promise.resolve({}));
mockAccountManagementService.setup(x => x.getAccountSecurityToken(TypeMoq.It.isValue(accounts[0]), TypeMoq.It.isAny(), TypeMoq.It.isAny()))
.returns(() => Promise.resolve(undefined));
mockAccountManagementService.setup(x => x.updateAccountListEvent)
.returns(() => () => { return undefined; });