Can clear token cache when wam is enabled (#28289)

Co-authored-by: Yeming Liu <[email protected]>

Author: isra-fel

src/Accounts/Accounts/ChangeLog.md

@@ -19,6 +19,7 @@
 -->
 
 ## Upcoming Release
+* Fixed an issue where `Clear-AzContext` does not clear the token cache when broker is enabled.
 * Added new parameter `-ClaimsChallenge` to `Connect-AzAccount` to support claims challenge authentication for MFA.
 * Refined the error message when a cmdlet fails because of policy violations about Multi-Factor Authentication (MFA) to provide more actionable guidance.
 

src/Accounts/Accounts/Context/ClearAzureRmContext.cs

@@ -65,6 +65,9 @@ void ClearContext(AzureRmProfile profile, RMProfileClient client)
             bool result = false;
             if (profile != null)
             {
+                string authorityHost = profile.DefaultContext?.Environment?.ActiveDirectoryAuthority;
+                string tenantId = profile.DefaultContext?.Tenant?.Id;
+
                 var contexts = profile.Contexts.Values;
                 foreach (var context in contexts)
                 {
@@ -78,11 +81,19 @@ void ClearContext(AzureRmProfile profile, RMProfileClient client)
                 }
                 else
                 {
-                    tokenCacheProvider.ClearCache();
+                    string authority = null;
+                    if (authorityHost != null)
+                    {
+                        authority = $"{authorityHost}/{tenantId}";
+                    }
+                    WriteDebug($"Clearing token cache for authority: {authority}");
+                    tokenCacheProvider.ClearCache(authority);
                     var defaultContext = new AzureContext();
                     profile.TrySetDefaultContext(defaultContext);
                     result = true;
                 }
+
+
                 if (AzureSession.Instance.TryGetComponent(AzKeyStore.Name, out AzKeyStore keyStore))
                 {
                     keyStore?.Clear();

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

@@ -24,10 +24,8 @@
 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;
 
@@ -59,9 +57,14 @@ public virtual void ClearCache()
         {
         }
 
+        public virtual void ClearCache(string authority)
+        {
+            ClearCache();
+        }
+
         public bool TryRemoveAccount(string accountId)
         {
-            TracingAdapter.Information(string.Format("[AuthenticationClientFactory] Calling GetAccountsAsync"));
+            TracingAdapter.Information(string.Format("[AuthenticationClientFactory] Calling TryRemoveAccount"));
             var client = CreatePublicClient();
             var account = client.GetAccountsAsync()
                             .ConfigureAwait(false).GetAwaiter().GetResult()
@@ -73,7 +76,7 @@ public bool TryRemoveAccount(string accountId)
 
             try
             {
-                TracingAdapter.Information(string.Format("[AuthenticationClientFactory] Calling RemoveAsync - Account: '{0}'", account.Username));
+                TracingAdapter.Information(string.Format("[AuthenticationClientFactory] Calling TryRemoveAccount - Account: '{0}'", account.Username));
                 client.RemoveAsync(account)
                     .ConfigureAwait(false).GetAwaiter().GetResult();
             }

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

@@ -12,12 +12,11 @@
 // 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
 {
@@ -105,7 +104,17 @@ protected override void RegisterCache(IPublicClientApplication client)
 
         public override void ClearCache()
         {
-            var client = CreatePublicClient();
+            ClearCacheInternal(null);
+        }
+
+        public override void ClearCache(string authority)
+        {
+            ClearCacheInternal(authority);
+        }
+
+        private void ClearCacheInternal(string authority)
+        {
+            var client = CreatePublicClient(authority);
             var accounts = client.GetAccountsAsync().GetAwaiter().GetResult();
             foreach (var account in accounts)
             {

src/Accounts/Authentication/Utilities/AzConfigReader.cs

@@ -63,7 +63,7 @@ static public bool IsWamEnabled(string authority)
                     {
                         authority = authority + "/";
                     }
-                    return Instance.GetConfigValue<bool>(ConfigKeys.EnableLoginByWam) && 0 == string.Compare(authority, AzureAuthorityHosts.AzurePublicCloud.OriginalString, System.StringComparison.OrdinalIgnoreCase);
+                    return Instance.GetConfigValue<bool>(ConfigKeys.EnableLoginByWam) && authority.StartsWith(AzureAuthorityHosts.AzurePublicCloud.OriginalString, System.StringComparison.OrdinalIgnoreCase);
                 }
                 catch
                 {