Protect WAM account picker deadlocks on STA thread apartments (Azure#36061)

Author: christothes

sdk/identity/Azure.Identity.BrokeredAuthentication/tests/ManualInteractiveBrowserCredentialBrokerTests.cs

@@ -2,7 +2,9 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Linq;
 using System.Runtime.InteropServices;
+using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core;
 using NUnit.Framework;
@@ -11,6 +13,8 @@ namespace Azure.Identity.BrokeredAuthentication.Tests
 {
     public class ManualInteractiveBrowserCredentialBrokerTests
     {
+        private static TokenRequestContext context = new TokenRequestContext(new string[] { "https://vault.azure.net/.default" });
+
         [DllImport("user32.dll")]
         private static extern IntPtr GetForegroundWindow();
 
@@ -28,5 +32,62 @@ public async Task AuthenticateWithBrokerAsync()
 
             Assert.NotNull(token.Token);
         }
+
+        [Test]
+        [Ignore("This test is an integration test which can only be run with user interaction")]
+        public void AuthenticateWithBrokerAsyncWithSTA([Values(true, false)] bool isAsync)
+        {
+            var assemblies = System.Reflection.Assembly.GetExecutingAssembly().GetReferencedAssemblies().FirstOrDefault(a => a.Name == "System.Windows.Forms");
+            if (assemblies != null)
+            {
+                throw new Exception("has winforms");
+            }
+            ManualResetEventSlim evt = new();
+            Thread thread = new Thread(async () =>
+            {
+                // do something with retVal
+
+                Console.WriteLine($"Thread {Thread.CurrentThread.GetApartmentState()}");
+                IntPtr parentWindowHandle = GetForegroundWindow();
+                var options = new InteractiveBrowserCredentialBrokerOptions(parentWindowHandle)
+                {
+                    TokenCachePersistenceOptions = new()
+                };
+
+                var cred = new InteractiveBrowserCredential(options);
+                var authRecord = isAsync ? await cred.AuthenticateAsync(context) : cred.Authenticate(context);
+                options.AuthenticationRecord = authRecord;
+                AccessToken token = isAsync ? await cred.GetTokenAsync(context).ConfigureAwait(false) : cred.GetToken(context);
+                Console.WriteLine("got token");
+                evt.Set();
+                // });
+            });
+
+            Thread thread2 = new Thread(async () =>
+            {
+                // do something with retVal
+
+                Console.WriteLine($"Thread {Thread.CurrentThread.GetApartmentState()}");
+                IntPtr parentWindowHandle = GetForegroundWindow();
+                var options = new InteractiveBrowserCredentialBrokerOptions(parentWindowHandle)
+                {
+                    TokenCachePersistenceOptions = new()
+                };
+
+                var cred = new InteractiveBrowserCredential(options);
+                var authRecord = isAsync ? await cred.AuthenticateAsync(context) : cred.Authenticate(context);
+                options.AuthenticationRecord = authRecord;
+                AccessToken token = isAsync ? await cred.GetTokenAsync(context).ConfigureAwait(false) : cred.GetToken(context);
+                Console.WriteLine("got token");
+                thread.Start();
+                // });
+            });
+#pragma warning disable CA1416 // Validate platform compatibility
+            thread.SetApartmentState(ApartmentState.STA);
+            thread2.SetApartmentState(ApartmentState.STA);
+#pragma warning restore CA1416 // Validate platform compatibility
+            thread2.Start();
+            evt.Wait(10000);
+        }
     }
 }

sdk/identity/Azure.Identity/src/MsalPublicClient.cs

@@ -126,11 +126,9 @@ protected virtual async ValueTask<AuthenticationResult> AcquireTokenSilentCoreAs
 
         public async ValueTask<AuthenticationResult> AcquireTokenInteractiveAsync(string[] scopes, string claims, Prompt prompt, string loginHint, string tenantId, bool async, CancellationToken cancellationToken)
         {
-#pragma warning disable AZC0109 // Misuse of 'async' parameter.
-            if (!async && !IdentityCompatSwitches.DisableInteractiveBrowserThreadpoolExecution)
-#pragma warning restore AZC0109 // Misuse of 'async' parameter.
+            if (Thread.CurrentThread.GetApartmentState() == ApartmentState.STA && !IdentityCompatSwitches.DisableInteractiveBrowserThreadpoolExecution)
             {
-                // In the synchronous case we need to use Task.Run to execute on the call to MSAL on the threadpool.
+                // In the case we are called in an STA apartment thread, we need to use Task.Run to execute on the call to MSAL on the threadpool.
                 // On certain platforms MSAL will use the embedded browser instead of launching the browser as a separate
                 // process. Executing with Task.Run prevents possibly deadlocking the UI thread in these cases.
                 // This workaround can be disabled by using the "Azure.Identity.DisableInteractiveBrowserThreadpoolExecution" app switch

sdk/identity/Azure.Identity/tests/InteractiveBrowserCredentialTests.cs

@@ -144,16 +144,16 @@ public async Task InteractiveBrowserValidateSyncWorkaroundCompatSwitch()
             // neither Environment variable or AppContext switch is set.
             // environment variable is set and AppContext switch is not set
             // AppContext switch is set
-            await ValidateSyncWorkaroundCompatSwitch(!IsAsync);
+            await ValidateSyncWorkaroundCompatSwitch(Thread.CurrentThread.GetApartmentState() == ApartmentState.STA);
 
             using (var envVar = new TestEnvVar("AZURE_IDENTITY_DISABLE_INTERACTIVEBROWSERTHREADPOOLEXECUTION", string.Empty))
             {
-                await ValidateSyncWorkaroundCompatSwitch(!IsAsync);
+                await ValidateSyncWorkaroundCompatSwitch(Thread.CurrentThread.GetApartmentState() == ApartmentState.STA);
             }
 
             using (var envVar = new TestEnvVar("AZURE_IDENTITY_DISABLE_INTERACTIVEBROWSERTHREADPOOLEXECUTION", "false"))
             {
-                await ValidateSyncWorkaroundCompatSwitch(!IsAsync);
+                await ValidateSyncWorkaroundCompatSwitch(Thread.CurrentThread.GetApartmentState() == ApartmentState.STA);
             }
 
             using (var envVar = new TestEnvVar("AZURE_IDENTITY_DISABLE_INTERACTIVEBROWSERTHREADPOOLEXECUTION", "true"))
@@ -172,7 +172,7 @@ public async Task InteractiveBrowserValidateSyncWorkaroundCompatSwitch()
 
             AppContext.SetSwitch("Azure.Identity.DisableInteractiveBrowserThreadpoolExecution", false);
 
-            await ValidateSyncWorkaroundCompatSwitch(!IsAsync);
+            await ValidateSyncWorkaroundCompatSwitch(Thread.CurrentThread.GetApartmentState() == ApartmentState.STA);
         }
 
         [Test]