[ReApply] Fixed the issue that context is not cleared when WAM is enabled

src/Accounts/Accounts/Account/DisconnectAzureRmAccount.cs

@@ -124,21 +124,20 @@ public override void ExecuteCmdlet()
 
                 if (ShouldProcess(string.Format("Log out principal '{0}'", azureAccount.Id), "log out"))
                 {
-                    if (GetContextModificationScope() == ContextModificationScope.CurrentUser)
-                    {
-                        AzureSession.Instance.AuthenticationFactory.RemoveUser(azureAccount, environment: null);
-                    }
-
                     if (AzureRmProfileProvider.Instance.Profile != null)
                     {
                         ModifyContext((localProfile, profileClient) =>
-                       {
-                           var matchingContexts = localProfile.Contexts?.Values?.Where((c) => c != null && c.Account != null && string.Equals(c.Account.Id, azureAccount.Id, StringComparison.CurrentCultureIgnoreCase));
-                           foreach (var context in matchingContexts)
-                           {
-                               profileClient.TryRemoveContext(context);
-                           }
-                       });
+                        {
+                            var matchingContexts = localProfile.Contexts?.Values?.Where((c) => c != null && c.Account != null && string.Equals(c.Account.Id, azureAccount.Id, StringComparison.CurrentCultureIgnoreCase));
+                            foreach (var context in matchingContexts)
+                            {
+                                if (GetContextModificationScope() == ContextModificationScope.CurrentUser)
+                                {
+                                    AzureSession.Instance.AuthenticationFactory.RemoveUser(azureAccount, context.Environment);
+                                }
+                                profileClient.TryRemoveContext(context);
+                            }
+                        });
                     }
 
                     WriteObject(new PSAzureRmAccount(azureAccount));

src/Accounts/Accounts/ChangeLog.md

@@ -19,6 +19,7 @@
 -->
 
 ## Upcoming Release
+* Fixed the issue that context is not cleared when WAM is enabled.
 
 ## Version 5.0.0
 * Changed the default output access token of `Get-AzAccessToken` from plain text to `SecureString`.

src/Accounts/Accounts/Context/ClearAzureRmContext.cs

@@ -65,29 +65,29 @@ void ClearContext(AzureRmProfile profile, RMProfileClient client)
             bool result = false;
             if (profile != null)
             {
+                PowerShellTokenCacheProvider tokenCacheProvider = null;
+                if (!AzureSession.Instance.TryGetComponent(PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey, out tokenCacheProvider))
+                {
+                    WriteWarning(Resources.ClientFactoryNotRegisteredClear);
+                }
+
                 var contexts = profile.Contexts.Values;
                 foreach (var context in contexts)
                 {
+                    tokenCacheProvider?.ClearCache(context?.Environment?.ActiveDirectoryAuthority);
                     client.TryRemoveContext(context);
                 }
 
-                PowerShellTokenCacheProvider tokenCacheProvider;
-                if (!AzureSession.Instance.TryGetComponent(PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey, out tokenCacheProvider))
+                if (tokenCacheProvider != null)
                 {
-                    WriteWarning(Resources.ClientFactoryNotRegisteredClear);
-                }
-                else
-                {
-                    tokenCacheProvider.ClearCache();
-                    var defaultContext = new AzureContext();
-                    profile.TrySetDefaultContext(defaultContext);
+                    profile.TrySetDefaultContext(new AzureContext());
                     result = true;
                 }
+
                 if (AzureSession.Instance.TryGetComponent(AzKeyStore.Name, out AzKeyStore keyStore))
                 {
                     keyStore?.Clear();
                 }
-
             }
 
             AzureSession.Instance.RaiseContextClearedEvent();

src/Accounts/Accounts/Context/GetAzureRMContext.cs

@@ -84,6 +84,10 @@ public override void ExecuteCmdlet()
                     var defaultProfile = DefaultProfile as AzureRmProfile;
                     if (defaultProfile != null && string.Equals(AzureSession.Instance?.ARMContextSaveMode, "CurrentUser"))
                     {
+                        if (!string.IsNullOrEmpty(DefaultContext?.Environment?.Name))
+                        {
+                            AzureSession.Instance.SetProperty(AzureSession.Property.Environment, DefaultContext.Environment.Name);
+                        }
                         defaultProfile.RefreshContextsFromCache(_cmdletContext);
                     }
                 }

src/Accounts/Accounts/Context/RemoveAzureRmContext.cs

@@ -12,17 +12,16 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
-using System;
-using System.Linq;
-using System.Management.Automation;
-
 using Microsoft.Azure.Commands.Common.Authentication;
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
 using Microsoft.Azure.Commands.Common.Authentication.Models;
 using Microsoft.Azure.Commands.Profile.Common;
 using Microsoft.Azure.Commands.Profile.Models.Core;
 using Microsoft.Azure.Commands.Profile.Properties;
-using Microsoft.WindowsAzure.Commands.Utilities.Common;
+
+using System;
+using System.Linq;
+using System.Management.Automation;
 
 namespace Microsoft.Azure.Commands.Profile.Context
 {
@@ -91,7 +90,7 @@ public override void ExecuteCmdlet()
                                         }
                                         else
                                         {
-                                            if (!tokenCacheProvider.TryRemoveAccount(removedContext.Account.Id))
+                                            if (!tokenCacheProvider.TryRemoveAccount(removedContext.Account.Id, removedContext?.Environment?.ActiveDirectoryAuthority))
                                             {
                                                 WriteWarning(string.Format(Resources.NoContextsRemain, removedContext.Account.Id));
                                             }

src/Accounts/Authentication.ResourceManager/AzureRmProfile.cs

@@ -819,9 +819,12 @@ public void RefreshContextsFromCache(ICmdletContext cmdletContext)
                 out PowerShellTokenCacheProvider tokenCacheProvider);
 
             string authority = null;
+            //If the function is called from "public virtual IAzureContext DefaultContext", authroity is empty and then ListAccounts will return empty.
+            //But as "ShouldRefreshContextsFromCache" is always false, the only call path is from GetAzureRMContext for now.
             if (TryGetEnvironment(AzureSession.Instance.GetProperty(AzureSession.Property.Environment), out IAzureEnvironment sessionEnvironment))
             {
-                authority = $"{sessionEnvironment.ActiveDirectoryAuthority}organizations";
+                authority = new Uri(new Uri(sessionEnvironment.ActiveDirectoryAuthority), "organizations").AbsoluteUri;
+
             }
             var accounts = tokenCacheProvider.ListAccounts(authority);
             if (!accounts.Any())

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

@@ -14,6 +14,7 @@
 
 using Azure.Identity;
 
+using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
 using Microsoft.Identity.Client;
 
 namespace Microsoft.Azure.Commands.Common.Authentication
@@ -46,7 +47,7 @@ public override void FlushTokenData()
             }
         }
 
-        public override void ClearCache()
+        public override void ClearCache(string authority)
         {
             InMemoryTokenCacheOptions = new InMemoryTokenCacheOptions();
         }

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

@@ -12,10 +12,6 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
-using System;
-using System.Collections.Generic;
-using System.Linq;
-
 using Azure.Identity;
 
 using Hyak.Common;
@@ -24,13 +20,15 @@
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions.Extensions;
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions.Interfaces;
 using Microsoft.Azure.Commands.Common.Authentication.Utilities;
-using Microsoft.Azure.Commands.Shared.Config;
 using Microsoft.Azure.Internal.Subscriptions;
 using Microsoft.Azure.Internal.Subscriptions.Models;
-using Microsoft.Azure.PowerShell.Common.Config;
 using Microsoft.Identity.Client;
 using Microsoft.Identity.Client.Broker;
 
+using System;
+using System.Collections.Generic;
+using System.Linq;
+
 namespace Microsoft.Azure.Commands.Common.Authentication
 {
     public abstract class PowerShellTokenCacheProvider
@@ -55,14 +53,14 @@ public virtual void FlushTokenData()
             _tokenCacheDataToFlush = null;
         }
 
-        public virtual void ClearCache()
+        public virtual void ClearCache(string authority = null)
         {
         }
 
-        public bool TryRemoveAccount(string accountId)
+        public bool TryRemoveAccount(string accountId, string authority = null)
         {
             TracingAdapter.Information(string.Format("[AuthenticationClientFactory] Calling GetAccountsAsync"));
-            var client = CreatePublicClient();
+            var client = CreatePublicClient(authority);
             var account = client.GetAccountsAsync()
                             .ConfigureAwait(false).GetAwaiter().GetResult()
                             .FirstOrDefault(a => string.Equals(a.Username, accountId, StringComparison.OrdinalIgnoreCase));
@@ -89,7 +87,7 @@ public IEnumerable<IAccount> ListAccounts(string authority = null)
         {
             TracingAdapter.Information(string.Format("[PowerShellTokenCacheProvider] Calling GetAccountsAsync on {0}", authority ?? "AzureCloud"));
 
-            return CreatePublicClient(authority: authority)
+            return CreatePublicClient(authority)
                     .GetAccountsAsync()
                     .ConfigureAwait(false).GetAwaiter().GetResult();
         }
@@ -195,18 +193,7 @@ public virtual IPublicClientApplication CreatePublicClient(string authority, str
         /// </summary>
         public virtual IPublicClientApplication CreatePublicClient(string authority = null)
         {
-            var builder = PublicClientApplicationBuilder.Create(Constants.PowerShellClientId);
-            if (AzConfigReader.IsWamEnabled(authority))
-            {
-                builder = builder.WithBroker(new BrokerOptions(BrokerOptions.OperatingSystems.Windows));
-            }
-            if (!string.IsNullOrEmpty(authority))
-            {
-                builder.WithAuthority(authority);
-            }
-            var client = builder.Build();
-            RegisterCache(client);
-            return client;
+            return CreatePublicClient(authority, organizationTenant);
         }
 
         public abstract TokenCachePersistenceOptions GetTokenCachePersistenceOptions();

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

@@ -12,13 +12,14 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
-using System;
-
 using Azure.Identity;
 
+using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
 using Microsoft.Identity.Client;
 using Microsoft.Identity.Client.Extensions.Msal;
 
+using System;
+
 namespace Microsoft.Azure.Commands.Common.Authentication
 {
     public class SharedTokenCacheProvider : PowerShellTokenCacheProvider
@@ -103,9 +104,9 @@ protected override void RegisterCache(IPublicClientApplication client)
             }
         }
 
-        public override void ClearCache()
+        public override void ClearCache(string authority)
         {
-            var client = CreatePublicClient();
+            var client = CreatePublicClient(authority);
             var accounts = client.GetAccountsAsync().GetAwaiter().GetResult();
             foreach (var account in accounts)
             {

src/Accounts/Authentication/Factories/AuthenticationFactory.cs

@@ -482,10 +482,10 @@ public ServiceClientCredentials GetServiceClientCredentials(string accessToken,
         }
 
         /// <summary>
-        /// Remove a user from token cache.
+        /// Remove any stored credentials for the given user and the Azure environment used.
         /// </summary>
-        /// <param name="account"></param>
-        /// <param name="environment"></param>
+        /// <param name="account">The account to remove credentials for</param>
+        /// <param name="environment">The environment which account belongs to</param>
         public void RemoveUser(IAzureAccount account, IAzureEnvironment environment)
         {
             if (account != null && !string.IsNullOrEmpty(account.Id) && !string.IsNullOrWhiteSpace(account.Type))
@@ -513,10 +513,10 @@ public void RemoveUser(IAzureAccount account, IAzureEnvironment environment)
                             // make best effort to remove credentials
                         }
 
-                        RemoveFromTokenCache(account);
+                        RemoveFromTokenCache(account, environment?.ActiveDirectoryAuthority);
                         break;
                     case AzureAccount.AccountType.User:
-                        RemoveFromTokenCache(account);
+                        RemoveFromTokenCache(account, environment?.ActiveDirectoryAuthority);
                         break;
                 }
             }
@@ -558,20 +558,20 @@ private string GetEndpointToken(IAzureAccount account, string targetEndpoint)
             return account.GetProperty(tokenKey);
         }
 
-        private void RemoveFromTokenCache(IAzureAccount account)
+        private void RemoveFromTokenCache(IAzureAccount account, string authority = null)
         {
             PowerShellTokenCacheProvider tokenCacheProvider;
             if (!AzureSession.Instance.TryGetComponent(PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey, out tokenCacheProvider))
             {
                 throw new NullReferenceException(Resources.AuthenticationClientFactoryNotRegistered);
             }
 
-            var publicClient = tokenCacheProvider.CreatePublicClient();
+            var publicClient = tokenCacheProvider.CreatePublicClient(authority);
             var accounts = publicClient.GetAccountsAsync()
                             .ConfigureAwait(false).GetAwaiter().GetResult();
             var tokenAccounts = accounts.Where(a => MatchCacheItem(account, a));
             foreach (var tokenAccount in tokenAccounts)
-                {
+            {
                 publicClient.RemoveAsync(tokenAccount)
                                 .ConfigureAwait(false).GetAwaiter().GetResult();
             }