Skip to content

Commit fe49e2b

Browse files
author
Lei jin
committed
Fix that RemoveUser doesn't work when create Public Client with broker
1 parent 8b6fd9b commit fe49e2b

File tree

11 files changed

+101
-62
lines changed

11 files changed

+101
-62
lines changed

‎src/Accounts/Accounts/Account/DisconnectAzureRmAccount.cs‎

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -124,21 +124,20 @@ public override void ExecuteCmdlet()
124124

125125
if (ShouldProcess(string.Format("Log out principal '{0}'", azureAccount.Id), "log out"))
126126
{
127-
if (GetContextModificationScope() == ContextModificationScope.CurrentUser)
128-
{
129-
AzureSession.Instance.AuthenticationFactory.RemoveUser(azureAccount, null);
130-
}
131-
132127
if (AzureRmProfileProvider.Instance.Profile != null)
133128
{
134129
ModifyContext((localProfile, profileClient) =>
135-
{
136-
var matchingContexts = localProfile.Contexts?.Values?.Where((c) => c != null && c.Account != null && string.Equals(c.Account.Id, azureAccount.Id, StringComparison.CurrentCultureIgnoreCase));
137-
foreach (var context in matchingContexts)
138-
{
139-
profileClient.TryRemoveContext(context);
140-
}
141-
});
130+
{
131+
var matchingContexts = localProfile.Contexts?.Values?.Where((c) => c != null && c.Account != null && string.Equals(c.Account.Id, azureAccount.Id, StringComparison.CurrentCultureIgnoreCase));
132+
foreach (var context in matchingContexts)
133+
{
134+
if (GetContextModificationScope() == ContextModificationScope.CurrentUser)
135+
{
136+
AzureSession.Instance.AuthenticationFactory.RemoveUser(azureAccount, context.Environment.ActiveDirectoryAuthority);
137+
}
138+
profileClient.TryRemoveContext(context);
139+
}
140+
});
142141
}
143142

144143
WriteObject(new PSAzureRmAccount(azureAccount));

‎src/Accounts/Accounts/Context/ClearAzureRmContext.cs‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -65,29 +65,29 @@ void ClearContext(AzureRmProfile profile, RMProfileClient client)
6565
bool result = false;
6666
if (profile != null)
6767
{
68+
PowerShellTokenCacheProvider tokenCacheProvider = null;
69+
if (!AzureSession.Instance.TryGetComponent(PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey, out tokenCacheProvider))
70+
{
71+
WriteWarning(Resources.ClientFactoryNotRegisteredClear);
72+
}
73+
6874
var contexts = profile.Contexts.Values;
6975
foreach (var context in contexts)
7076
{
77+
tokenCacheProvider?.ClearCache(context.Environment.ActiveDirectoryAuthority);
7178
client.TryRemoveContext(context);
7279
}
7380

74-
PowerShellTokenCacheProvider tokenCacheProvider;
75-
if (!AzureSession.Instance.TryGetComponent(PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey, out tokenCacheProvider))
81+
if (tokenCacheProvider != null)
7682
{
77-
WriteWarning(Resources.ClientFactoryNotRegisteredClear);
78-
}
79-
else
80-
{
81-
tokenCacheProvider.ClearCache();
82-
var defaultContext = new AzureContext();
83-
profile.TrySetDefaultContext(defaultContext);
83+
profile.TrySetDefaultContext(new AzureContext());
8484
result = true;
8585
}
86+
8687
if (AzureSession.Instance.TryGetComponent(AzKeyStore.Name, out AzKeyStore keyStore))
8788
{
8889
keyStore?.Clear();
8990
}
90-
9191
}
9292

9393
AzureSession.Instance.RaiseContextClearedEvent();

‎src/Accounts/Accounts/Context/GetAzureRMContext.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,7 @@ public override void ExecuteCmdlet()
8484
var defaultProfile = DefaultProfile as AzureRmProfile;
8585
if (defaultProfile != null && string.Equals(AzureSession.Instance?.ARMContextSaveMode, "CurrentUser"))
8686
{
87+
AzureSession.Instance.SetProperty(AzureSession.Property.Environment, DefaultContext.Environment.Name);
8788
defaultProfile.RefreshContextsFromCache();
8889
}
8990
}

‎src/Accounts/Accounts/Context/RemoveAzureRmContext.cs‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,17 +12,16 @@
1212
// limitations under the License.
1313
// ----------------------------------------------------------------------------------
1414

15-
using System;
16-
using System.Linq;
17-
using System.Management.Automation;
18-
1915
using Microsoft.Azure.Commands.Common.Authentication;
2016
using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
2117
using Microsoft.Azure.Commands.Common.Authentication.Models;
2218
using Microsoft.Azure.Commands.Profile.Common;
2319
using Microsoft.Azure.Commands.Profile.Models.Core;
2420
using Microsoft.Azure.Commands.Profile.Properties;
25-
using Microsoft.WindowsAzure.Commands.Utilities.Common;
21+
22+
using System;
23+
using System.Linq;
24+
using System.Management.Automation;
2625

2726
namespace Microsoft.Azure.Commands.Profile.Context
2827
{
@@ -91,7 +90,7 @@ public override void ExecuteCmdlet()
9190
}
9291
else
9392
{
94-
if (!tokenCacheProvider.TryRemoveAccount(removedContext.Account.Id))
93+
if (!tokenCacheProvider.TryRemoveAccount(removedContext.Account.Id, removedContext.Environment.ActiveDirectoryAuthority))
9594
{
9695
WriteWarning(string.Format(Resources.NoContextsRemain, removedContext.Account.Id));
9796
}

‎src/Accounts/Authentication.ResourceManager/AzureRmProfile.cs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -822,8 +822,9 @@ public void RefreshContextsFromCache()
822822
string authority = null;
823823
if (TryGetEnvironment(AzureSession.Instance.GetProperty(AzureSession.Property.Environment), out IAzureEnvironment sessionEnvironment))
824824
{
825-
authority = $"{sessionEnvironment.ActiveDirectoryAuthority}organizations";
825+
authority = $"{sessionEnvironment.ActiveDirectoryAuthority}/organizations";
826826
}
827+
//fixme, if Connect-AzAccount not run, when to get authority
827828
var accounts = tokenCacheProvider.ListAccounts(authority);
828829
if (!accounts.Any())
829830
{

‎src/Accounts/Authentication/Authentication/TokenCache/InMemoryTokenCacheProvider.cs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
using Azure.Identity;
1616

17+
using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
1718
using Microsoft.Identity.Client;
1819

1920
namespace Microsoft.Azure.Commands.Common.Authentication
@@ -46,7 +47,7 @@ public override void FlushTokenData()
4647
}
4748
}
4849

49-
public override void ClearCache()
50+
public override void ClearCache(string authority)
5051
{
5152
InMemoryTokenCacheOptions = new InMemoryTokenCacheOptions();
5253
}

‎src/Accounts/Authentication/Authentication/TokenCache/PowerShellTokenCacheProvider.cs‎

Lines changed: 9 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -12,23 +12,21 @@
1212
// limitations under the License.
1313
// ----------------------------------------------------------------------------------
1414

15-
using System;
16-
using System.Collections.Generic;
17-
using System.Linq;
18-
1915
using Azure.Identity;
2016

2117
using Hyak.Common;
2218

2319
using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
2420
using Microsoft.Azure.Commands.Common.Authentication.Utilities;
25-
using Microsoft.Azure.Commands.Shared.Config;
2621
using Microsoft.Azure.Internal.Subscriptions;
2722
using Microsoft.Azure.Internal.Subscriptions.Models;
28-
using Microsoft.Azure.PowerShell.Common.Config;
2923
using Microsoft.Identity.Client;
3024
using Microsoft.Identity.Client.Broker;
3125

26+
using System;
27+
using System.Collections.Generic;
28+
using System.Linq;
29+
3230
namespace Microsoft.Azure.Commands.Common.Authentication
3331
{
3432
public abstract class PowerShellTokenCacheProvider
@@ -53,14 +51,14 @@ public virtual void FlushTokenData()
5351
_tokenCacheDataToFlush = null;
5452
}
5553

56-
public virtual void ClearCache()
54+
public virtual void ClearCache(string authority = null)
5755
{
5856
}
5957

60-
public bool TryRemoveAccount(string accountId)
58+
public bool TryRemoveAccount(string accountId, string authority = null)
6159
{
6260
TracingAdapter.Information(string.Format("[AuthenticationClientFactory] Calling GetAccountsAsync"));
63-
var client = CreatePublicClient();
61+
var client = CreatePublicClient(authority);
6462
var account = client.GetAccountsAsync()
6563
.ConfigureAwait(false).GetAwaiter().GetResult()
6664
.FirstOrDefault(a => string.Equals(a.Username, accountId, StringComparison.OrdinalIgnoreCase));
@@ -87,7 +85,7 @@ public IEnumerable<IAccount> ListAccounts(string authority = null)
8785
{
8886
TracingAdapter.Information(string.Format("[PowerShellTokenCacheProvider] Calling GetAccountsAsync on {0}", authority ?? "AzureCloud"));
8987

90-
return CreatePublicClient(authority: authority)
88+
return CreatePublicClient(authority)
9189
.GetAccountsAsync()
9290
.ConfigureAwait(false).GetAwaiter().GetResult();
9391
}
@@ -192,18 +190,7 @@ public virtual IPublicClientApplication CreatePublicClient(string authority, str
192190
/// </summary>
193191
public virtual IPublicClientApplication CreatePublicClient(string authority = null)
194192
{
195-
var builder = PublicClientApplicationBuilder.Create(Constants.PowerShellClientId);
196-
if (AzConfigReader.IsWamEnabled(authority))
197-
{
198-
builder = builder.WithBroker(new BrokerOptions(BrokerOptions.OperatingSystems.Windows));
199-
}
200-
if (!string.IsNullOrEmpty(authority))
201-
{
202-
builder.WithAuthority(authority);
203-
}
204-
var client = builder.Build();
205-
RegisterCache(client);
206-
return client;
193+
return CreatePublicClient(authority, organizationTenant);
207194
}
208195

209196
public abstract TokenCachePersistenceOptions GetTokenCachePersistenceOptions();

‎src/Accounts/Authentication/Authentication/TokenCache/SharedTokenCacheProvider.cs‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,13 +12,14 @@
1212
// limitations under the License.
1313
// ----------------------------------------------------------------------------------
1414

15-
using System;
16-
1715
using Azure.Identity;
1816

17+
using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
1918
using Microsoft.Identity.Client;
2019
using Microsoft.Identity.Client.Extensions.Msal;
2120

21+
using System;
22+
2223
namespace Microsoft.Azure.Commands.Common.Authentication
2324
{
2425
public class SharedTokenCacheProvider : PowerShellTokenCacheProvider
@@ -103,9 +104,9 @@ protected override void RegisterCache(IPublicClientApplication client)
103104
}
104105
}
105106

106-
public override void ClearCache()
107+
public override void ClearCache(string authority)
107108
{
108-
var client = CreatePublicClient();
109+
var client = CreatePublicClient(authority);
109110
var accounts = client.GetAccountsAsync().GetAwaiter().GetResult();
110111
foreach (var account in accounts)
111112
{

‎src/Accounts/Authentication/Factories/AuthenticationFactory.cs‎

Lines changed: 46 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,19 +13,18 @@
1313
// ----------------------------------------------------------------------------------
1414

1515
using Hyak.Common;
16+
1617
using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
1718
using Microsoft.Azure.Commands.Common.Authentication.Authentication;
1819
using Microsoft.Azure.Commands.Common.Authentication.Properties;
1920
using Microsoft.Azure.Commands.Common.Authentication.Utilities;
2021
using Microsoft.Azure.Commands.Common.Exceptions;
2122
using Microsoft.Azure.Commands.ResourceManager.Common;
22-
using Microsoft.Azure.Commands.Shared.Config;
23-
using Microsoft.Azure.PowerShell.Common.Config;
2423
using Microsoft.Identity.Client;
2524
using Microsoft.Rest;
2625
using Microsoft.WindowsAzure.Commands.Common;
26+
2727
using System;
28-
using System.Diagnostics;
2928
using System.Linq;
3029
using System.Security;
3130
using System.Threading.Tasks;
@@ -449,6 +448,47 @@ public void RemoveUser(IAzureAccount account, IAzureTokenCache tokenCache)
449448
}
450449
}
451450

451+
/// <summary>
452+
/// Remove any stored credentials for the given user and the Azure environment used.
453+
/// </summary>
454+
/// <param name="account">The account to remove credentials for</param>
455+
/// <param name="authority">The Microsoft Entra authority</param>
456+
public void RemoveUser(IAzureAccount account, string authority)
457+
{
458+
if (account != null && !string.IsNullOrEmpty(account.Id) && !string.IsNullOrWhiteSpace(account.Type))
459+
{
460+
switch (account.Type)
461+
{
462+
case AzureAccount.AccountType.AccessToken:
463+
account.SetProperty(AzureAccount.Property.AccessToken, null);
464+
account.SetProperty(AzureAccount.Property.GraphAccessToken, null);
465+
account.SetProperty(AzureAccount.Property.KeyVaultAccessToken, null);
466+
break;
467+
case AzureAccount.AccountType.ManagedService:
468+
account.SetProperty(AzureAccount.Property.MSILoginUri, null);
469+
break;
470+
case AzureAccount.AccountType.ServicePrincipal:
471+
try
472+
{
473+
KeyStore.RemoveSecureString(new ServicePrincipalKey(AzureAccount.Property.ServicePrincipalSecret,
474+
account.Id, account.GetTenants().FirstOrDefault()));
475+
KeyStore.RemoveSecureString(new ServicePrincipalKey(AzureAccount.Property.CertificatePassword,
476+
account.Id, account.GetTenants().FirstOrDefault()));
477+
}
478+
catch
479+
{
480+
// make best effort to remove credentials
481+
}
482+
483+
RemoveFromTokenCache(account, authority);
484+
break;
485+
case AzureAccount.AccountType.User:
486+
RemoveFromTokenCache(account, authority);
487+
break;
488+
}
489+
}
490+
}
491+
452492
private string GetResourceId(string resourceIdorEndpointName, IAzureEnvironment environment)
453493
{
454494
return environment.GetEndpoint(resourceIdorEndpointName) ?? resourceIdorEndpointName;
@@ -485,20 +525,20 @@ private string GetEndpointToken(IAzureAccount account, string targetEndpoint)
485525
return account.GetProperty(tokenKey);
486526
}
487527

488-
private void RemoveFromTokenCache(IAzureAccount account)
528+
private void RemoveFromTokenCache(IAzureAccount account, string authority = null)
489529
{
490530
PowerShellTokenCacheProvider tokenCacheProvider;
491531
if (!AzureSession.Instance.TryGetComponent(PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey, out tokenCacheProvider))
492532
{
493533
throw new NullReferenceException(Resources.AuthenticationClientFactoryNotRegistered);
494534
}
495535

496-
var publicClient = tokenCacheProvider.CreatePublicClient();
536+
var publicClient = tokenCacheProvider.CreatePublicClient(authority);
497537
var accounts = publicClient.GetAccountsAsync()
498538
.ConfigureAwait(false).GetAwaiter().GetResult();
499539
var tokenAccounts = accounts.Where(a => MatchCacheItem(account, a));
500540
foreach (var tokenAccount in tokenAccounts)
501-
{
541+
{
502542
publicClient.RemoveAsync(tokenAccount)
503543
.ConfigureAwait(false).GetAwaiter().GetResult();
504544
}

‎tools/TestFx/Mocks/MockCertificateAuthenticationFactory.cs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,5 +102,10 @@ public void RemoveUser(IAzureAccount account, IAzureTokenCache tokenCache)
102102
{
103103
throw new NotImplementedException();
104104
}
105+
106+
public void RemoveUser(IAzureAccount account, string authority)
107+
{
108+
throw new NotImplementedException();
109+
}
105110
}
106111
}

0 commit comments

Comments
 (0)