Skip to content

Commit b517d9f

Browse files
author
Lei jin
committed
Fix that RemoveUser doesn't work when create Public Client with broker
1 parent 2bca398 commit b517d9f

File tree

11 files changed

+99
-59
lines changed

11 files changed

+99
-59
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(_cmdletContext);
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
@@ -823,8 +823,9 @@ public void RefreshContextsFromCache(ICmdletContext cmdletContext)
823823
string authority = null;
824824
if (TryGetEnvironment(AzureSession.Instance.GetProperty(AzureSession.Property.Environment), out IAzureEnvironment sessionEnvironment))
825825
{
826-
authority = $"{sessionEnvironment.ActiveDirectoryAuthority}organizations";
826+
authority = $"{sessionEnvironment.ActiveDirectoryAuthority}/organizations";
827827
}
828+
//fixme, if Connect-AzAccount not run, when to get authority
828829
var accounts = tokenCacheProvider.ListAccounts(authority);
829830
if (!accounts.Any())
830831
{

‎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,10 +12,6 @@
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;
@@ -24,13 +20,15 @@
2420
using Microsoft.Azure.Commands.Common.Authentication.Abstractions.Extensions;
2521
using Microsoft.Azure.Commands.Common.Authentication.Abstractions.Interfaces;
2622
using Microsoft.Azure.Commands.Common.Authentication.Utilities;
27-
using Microsoft.Azure.Commands.Shared.Config;
2823
using Microsoft.Azure.Internal.Subscriptions;
2924
using Microsoft.Azure.Internal.Subscriptions.Models;
30-
using Microsoft.Azure.PowerShell.Common.Config;
3125
using Microsoft.Identity.Client;
3226
using Microsoft.Identity.Client.Broker;
3327

28+
using System;
29+
using System.Collections.Generic;
30+
using System.Linq;
31+
3432
namespace Microsoft.Azure.Commands.Common.Authentication
3533
{
3634
public abstract class PowerShellTokenCacheProvider
@@ -55,14 +53,14 @@ public virtual void FlushTokenData()
5553
_tokenCacheDataToFlush = null;
5654
}
5755

58-
public virtual void ClearCache()
56+
public virtual void ClearCache(string authority = null)
5957
{
6058
}
6159

62-
public bool TryRemoveAccount(string accountId)
60+
public bool TryRemoveAccount(string accountId, string authority = null)
6361
{
6462
TracingAdapter.Information(string.Format("[AuthenticationClientFactory] Calling GetAccountsAsync"));
65-
var client = CreatePublicClient();
63+
var client = CreatePublicClient(authority);
6664
var account = client.GetAccountsAsync()
6765
.ConfigureAwait(false).GetAwaiter().GetResult()
6866
.FirstOrDefault(a => string.Equals(a.Username, accountId, StringComparison.OrdinalIgnoreCase));
@@ -89,7 +87,7 @@ public IEnumerable<IAccount> ListAccounts(string authority = null)
8987
{
9088
TracingAdapter.Information(string.Format("[PowerShellTokenCacheProvider] Calling GetAccountsAsync on {0}", authority ?? "AzureCloud"));
9189

92-
return CreatePublicClient(authority: authority)
90+
return CreatePublicClient(authority)
9391
.GetAccountsAsync()
9492
.ConfigureAwait(false).GetAwaiter().GetResult();
9593
}
@@ -195,18 +193,7 @@ public virtual IPublicClientApplication CreatePublicClient(string authority, str
195193
/// </summary>
196194
public virtual IPublicClientApplication CreatePublicClient(string authority = null)
197195
{
198-
var builder = PublicClientApplicationBuilder.Create(Constants.PowerShellClientId);
199-
if (AzConfigReader.IsWamEnabled(authority))
200-
{
201-
builder = builder.WithBroker(new BrokerOptions(BrokerOptions.OperatingSystems.Windows));
202-
}
203-
if (!string.IsNullOrEmpty(authority))
204-
{
205-
builder.WithAuthority(authority);
206-
}
207-
var client = builder.Build();
208-
RegisterCache(client);
209-
return client;
196+
return CreatePublicClient(authority, organizationTenant);
210197
}
211198

212199
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: 44 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -522,6 +522,47 @@ public void RemoveUser(IAzureAccount account, IAzureTokenCache tokenCache)
522522
}
523523
}
524524

525+
/// <summary>
526+
/// Remove any stored credentials for the given user and the Azure environment used.
527+
/// </summary>
528+
/// <param name="account">The account to remove credentials for</param>
529+
/// <param name="authority">The Microsoft Entra authority</param>
530+
public void RemoveUser(IAzureAccount account, string authority)
531+
{
532+
if (account != null && !string.IsNullOrEmpty(account.Id) && !string.IsNullOrWhiteSpace(account.Type))
533+
{
534+
switch (account.Type)
535+
{
536+
case AzureAccount.AccountType.AccessToken:
537+
account.SetProperty(AzureAccount.Property.AccessToken, null);
538+
account.SetProperty(AzureAccount.Property.GraphAccessToken, null);
539+
account.SetProperty(AzureAccount.Property.KeyVaultAccessToken, null);
540+
break;
541+
case AzureAccount.AccountType.ManagedService:
542+
account.SetProperty(AzureAccount.Property.MSILoginUri, null);
543+
break;
544+
case AzureAccount.AccountType.ServicePrincipal:
545+
try
546+
{
547+
KeyStore.RemoveSecureString(new ServicePrincipalKey(AzureAccount.Property.ServicePrincipalSecret,
548+
account.Id, account.GetTenants().FirstOrDefault()));
549+
KeyStore.RemoveSecureString(new ServicePrincipalKey(AzureAccount.Property.CertificatePassword,
550+
account.Id, account.GetTenants().FirstOrDefault()));
551+
}
552+
catch
553+
{
554+
// make best effort to remove credentials
555+
}
556+
557+
RemoveFromTokenCache(account, authority);
558+
break;
559+
case AzureAccount.AccountType.User:
560+
RemoveFromTokenCache(account, authority);
561+
break;
562+
}
563+
}
564+
}
565+
525566
private string GetResourceId(string resourceIdorEndpointName, IAzureEnvironment environment)
526567
{
527568
return environment.GetEndpoint(resourceIdorEndpointName) ?? resourceIdorEndpointName;
@@ -558,20 +599,20 @@ private string GetEndpointToken(IAzureAccount account, string targetEndpoint)
558599
return account.GetProperty(tokenKey);
559600
}
560601

561-
private void RemoveFromTokenCache(IAzureAccount account)
602+
private void RemoveFromTokenCache(IAzureAccount account, string authority = null)
562603
{
563604
PowerShellTokenCacheProvider tokenCacheProvider;
564605
if (!AzureSession.Instance.TryGetComponent(PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey, out tokenCacheProvider))
565606
{
566607
throw new NullReferenceException(Resources.AuthenticationClientFactoryNotRegistered);
567608
}
568609

569-
var publicClient = tokenCacheProvider.CreatePublicClient();
610+
var publicClient = tokenCacheProvider.CreatePublicClient(authority);
570611
var accounts = publicClient.GetAccountsAsync()
571612
.ConfigureAwait(false).GetAwaiter().GetResult();
572613
var tokenAccounts = accounts.Where(a => MatchCacheItem(account, a));
573614
foreach (var tokenAccount in tokenAccounts)
574-
{
615+
{
575616
publicClient.RemoveAsync(tokenAccount)
576617
.ConfigureAwait(false).GetAwaiter().GetResult();
577618
}

‎tools/TestFx/Mocks/MockCertificateAuthenticationFactory.cs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,5 +124,10 @@ public IAccessToken Authenticate(IAzureAccount account, IAzureEnvironment enviro
124124
{
125125
throw new NotImplementedException();
126126
}
127+
128+
public void RemoveUser(IAzureAccount account, string authority)
129+
{
130+
throw new NotImplementedException();
131+
}
127132
}
128133
}

‎tools/TestFx/Mocks/MockTokenAuthenticationFactory.cs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,5 +148,10 @@ public ServiceClientCredentials GetServiceClientCredentials(IAzureContext contex
148148
{
149149
return GetServiceClientCredentials(context, targetEndpoint, AzureCmdletContext.CmdletNone);
150150
}
151+
152+
public void RemoveUser(IAzureAccount account, string authority)
153+
{
154+
throw new NotImplementedException();
155+
}
151156
}
152157
}

0 commit comments

Comments
 (0)