From 5ea13d34de1b7f99b338a2fc5b60fc395603b214 Mon Sep 17 00:00:00 2001
From: Jin Lei <54836179+msJinLei@users.noreply.github.com>
Date: Mon, 1 Jul 2024 14:55:42 +0800
Subject: [PATCH] Disable WAM When UseDeviceAuthentication is true or
 paramterset is UserWithCredential (#25404)

* Disable WAM when UseDeviceAuthentication is true or paramterset is UserWithCredential

* Address review comments

* Address review comments

* Update src/Accounts/Accounts/ChangeLog.md

Co-authored-by: Yeming Liu <11371776+isra-fel@users.noreply.github.com>

---------

Co-authored-by: Yeming Liu <11371776+isra-fel@users.noreply.github.com>
---
 .../Accounts/Account/ConnectAzureRmAccount.cs |  7 ++++++
 src/Accounts/Accounts/ChangeLog.md            |  1 +
 .../Utilities/AzConfigReader.cs               | 23 ++++++++++++++++---
 3 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/src/Accounts/Accounts/Account/ConnectAzureRmAccount.cs b/src/Accounts/Accounts/Account/ConnectAzureRmAccount.cs
index b768dfdaea3d..9d2c1131bd0a 100644
--- a/src/Accounts/Accounts/Account/ConnectAzureRmAccount.cs
+++ b/src/Accounts/Accounts/Account/ConnectAzureRmAccount.cs
@@ -330,6 +330,13 @@ public override void ExecuteCmdlet()
             Guid subscriptionIdGuid;
             string subscriptionName = null;
             string subscriptionId = null;
+
+            //Disable WAM before the issue https://github.com/AzureAD/microsoft-authentication-library-for-dotnet/issues/4786 is fixed
+            if (ParameterSetName.Equals(UserParameterSet) && UseDeviceAuthentication == true || ParameterSetName.Equals(UserWithCredentialParameterSet))
+            {
+                AzConfigReader.Instance?.UpdateConfig(ConfigKeys.EnableLoginByWam, false, ConfigScope.CurrentUser);
+            }
+
             if (MyInvocation.BoundParameters.ContainsKey(nameof(Subscription)))
             {
                 if (Guid.TryParse(Subscription, out subscriptionIdGuid))
diff --git a/src/Accounts/Accounts/ChangeLog.md b/src/Accounts/Accounts/ChangeLog.md
index 56bf1295a90a..d0330e62ff10 100644
--- a/src/Accounts/Accounts/ChangeLog.md
+++ b/src/Accounts/Accounts/ChangeLog.md
@@ -19,6 +19,7 @@
 -->
 
 ## Upcoming Release
+* Disable WAM when the customers login with device code flow or username password (ROPC) flow to prevent a potential issue with token cache.
 * Fixed [CVE-2024-35255](https://github.com/advisories/GHSA-m5vv-6r4h-3vj9)
 * Updated `Microsoft.Identity.Client.NativeInterop` to fix the WAM pop window issue in elevated mode [#24967]
 * Updated the reference of Azure PowerShell Common to 1.3.98-preview.
diff --git a/src/Accounts/Authentication/Utilities/AzConfigReader.cs b/src/Accounts/Authentication/Utilities/AzConfigReader.cs
index da38112abc3a..5c32b7b9f152 100644
--- a/src/Accounts/Authentication/Utilities/AzConfigReader.cs
+++ b/src/Accounts/Authentication/Utilities/AzConfigReader.cs
@@ -31,14 +31,31 @@ private static IAzureSession Session
             }
         }
 
+        private static IConfigManager instance = null;
+
+        public static IConfigManager Instance
+        {
+            get
+            {
+                if (instance == null)
+                {
+                    if (!Session.TryGetComponent<IConfigManager>(nameof(IConfigManager), out instance))
+                    {
+                        instance = null;
+                    }
+                }
+                return instance;
+            }
+        }
+
         public static T GetAzConfig<T>(string key, T defaultValue = default(T))
         {
-            return Session.TryGetComponent<IConfigManager>(nameof(IConfigManager), out IConfigManager configManager) ? configManager.GetConfigValue<T>(key) : defaultValue;
+            return Instance != null ? Instance.GetConfigValue<T>(key) : defaultValue;
         }
 
         static public bool IsWamEnabled(string authority)
         {
-            if (!string.IsNullOrEmpty(authority) && Session.TryGetComponent<IConfigManager>(nameof(IConfigManager), out var config))
+            if (!string.IsNullOrEmpty(authority) && Instance != null)
             {
                 try
                 {
@@ -46,7 +63,7 @@ static public bool IsWamEnabled(string authority)
                     {
                         authority = authority + "/";
                     }
-                    return config.GetConfigValue<bool>(ConfigKeys.EnableLoginByWam) && 0 == string.Compare(authority, AzureAuthorityHosts.AzurePublicCloud.OriginalString, System.StringComparison.OrdinalIgnoreCase);
+                    return Instance.GetConfigValue<bool>(ConfigKeys.EnableLoginByWam) && 0 == string.Compare(authority, AzureAuthorityHosts.AzurePublicCloud.OriginalString, System.StringComparison.OrdinalIgnoreCase);
                 }
                 catch
                 {