Fix deadlock in AzureCliCredential and other process based TokenCredential implementations

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

@@ -0,0 +1,24 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System.Diagnostics;
+
+namespace Azure.Identity
+{
+    internal class DataReceivedEventArgsWrapper
+    {
+        public DataReceivedEventArgsWrapper(string data)
+        {
+            Data = data;
+        }
+
+        public DataReceivedEventArgsWrapper(DataReceivedEventArgs args)
+        {
+            Data = args?.Data;
+        }
+
+        public string Data { get; }
+    }
+
+    internal delegate void DataReceivedEventWrapperHandler(object sender, DataReceivedEventArgsWrapper e);
+}

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

@@ -16,6 +16,8 @@ internal interface IProcess : IDisposable
         ProcessStartInfo StartInfo { get; set; }
 
         event EventHandler Exited;
+        event DataReceivedEventWrapperHandler OutputDataReceived;
+        event DataReceivedEventWrapperHandler ErrorDataReceived;
 
         void Start();
         void Kill();

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

@@ -2,6 +2,8 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Diagnostics;
+using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -13,16 +15,24 @@ internal sealed class ProcessRunner
     {
         private readonly IProcess _process;
         private readonly TimeSpan _timeout;
-        private readonly TaskCompletionSource<string> _tcs;
+        private readonly TaskCompletionSource<bool> _processExitedCompletionSource;
         private readonly CancellationToken _cancellationToken;
         private readonly CancellationTokenSource _timeoutCts;
         private CancellationTokenRegistration _ctRegistration;
+        private readonly StringBuilder _stdOutBuilder;
+        private readonly StringBuilder _stdErrBuilder;
+        private readonly TaskCompletionSource<string> _stdOutCompletionSource;
+        private readonly TaskCompletionSource<string> _stdErrCompletionSource;
 
         public ProcessRunner(IProcess process, TimeSpan timeout, CancellationToken cancellationToken)
         {
             _process = process;
             _timeout = timeout;
-            _tcs = new TaskCompletionSource<string>(TaskCreationOptions.RunContinuationsAsynchronously);
+            _processExitedCompletionSource = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+            _stdOutBuilder = new StringBuilder();
+            _stdErrBuilder = new StringBuilder();
+            _stdOutCompletionSource = new TaskCompletionSource<string>(TaskCreationOptions.RunContinuationsAsynchronously);
+            _stdErrCompletionSource = new TaskCompletionSource<string>(TaskCreationOptions.RunContinuationsAsynchronously);
 
             if (timeout.TotalMilliseconds >= 0)
             {
@@ -35,29 +45,53 @@ public ProcessRunner(IProcess process, TimeSpan timeout, CancellationToken cance
             }
         }
 
-        public Task<string> RunAsync()
+        public async Task<string> RunAsync()
         {
             StartProcess();
-            return _tcs.Task;
+
+            try
+            {
+                await _processExitedCompletionSource.Task.ConfigureAwait(false);
+                var output = await _stdOutCompletionSource.Task.ConfigureAwait(false);
+                await _stdErrCompletionSource.Task.ConfigureAwait(false);
+                return output;
+            }
+            finally
+            {
+                DisposeProcess();
+            }
         }
 
         public string Run()
         {
             StartProcess();
 #pragma warning disable AZC0102 // Do not use GetAwaiter().GetResult().
-            return _tcs.Task.GetAwaiter().GetResult();
+            try
+            {
+                _processExitedCompletionSource.Task.GetAwaiter().GetResult();
+                var output = _stdOutCompletionSource.Task.GetAwaiter().GetResult();
+                _stdErrCompletionSource.Task.GetAwaiter().GetResult();
+                return output;
+            }
+            finally
+            {
+                DisposeProcess();
+            }
 #pragma warning restore AZC0102 // Do not use GetAwaiter().GetResult().
         }
 
         private void StartProcess()
         {
-            if (TrySetCanceled() || _tcs.Task.IsCompleted)
+            if (TrySetCanceled() || _processExitedCompletionSource.Task.IsCompleted)
             {
                 return;
             }
 
             _process.Exited += (o, e) => HandleExit();
 
+            _process.OutputDataReceived += HandleStdOutDataReceived;
+            _process.ErrorDataReceived += HandleStdErrDataReceived;
+
             _process.StartInfo.UseShellExecute = false;
             _process.StartInfo.RedirectStandardOutput = true;
             _process.StartInfo.RedirectStandardError = true;
@@ -72,17 +106,48 @@ private void HandleExit()
         {
             if (_process.ExitCode == 0)
             {
-                TrySetResult(_process.StandardOutput.ReadToEnd());
+                _processExitedCompletionSource.TrySetResult(true);
             }
             else
             {
-                TrySetException(new InvalidOperationException(_process.StandardError.ReadToEnd()));
+                _processExitedCompletionSource.TrySetResult(false);
+            }
+        }
+        private void HandleStdErrDataReceived(object sender, DataReceivedEventArgsWrapper e)
+        {
+            if (e.Data is null)
+            {
+                var result = _stdErrBuilder.ToString();
+                if (result.Length == 0)
+                {
+                    _stdErrCompletionSource.TrySetResult(result);
+                }
+                else
+                {
+                    _stdErrCompletionSource.TrySetException(new InvalidOperationException(result));
+                }
+            }
+            else
+            {
+                _stdErrBuilder.Append(e.Data);
+            }
+        }
+
+        private void HandleStdOutDataReceived(object sender, DataReceivedEventArgsWrapper e)
+        {
+            if (e.Data is null)
+            {
+                _stdOutCompletionSource.TrySetResult(_stdOutBuilder.ToString());
+            }
+            else
+            {
+                _stdOutBuilder.Append(e.Data);
             }
         }
 
         private void HandleCancel()
         {
-            if (_tcs.Task.IsCompleted)
+            if (_processExitedCompletionSource.Task.IsCompleted)
             {
                 return;
             }
@@ -95,39 +160,31 @@ private void HandleCancel()
                 }
                 catch (Exception ex)
                 {
-                    TrySetException(ex);
+                    _processExitedCompletionSource.TrySetException(ex);
                     return;
                 }
             }
 
             TrySetCanceled();
         }
 
-        private void TrySetResult(string result)
-        {
-            DisposeProcess();
-            _tcs.TrySetResult(result);
-        }
-
         private bool TrySetCanceled()
         {
             if (_cancellationToken.IsCancellationRequested)
             {
-                DisposeProcess();
-                _tcs.TrySetCanceled(_cancellationToken);
+                _stdOutCompletionSource.TrySetCanceled(_cancellationToken);
+                _stdErrCompletionSource.TrySetCanceled(_cancellationToken);
+                _processExitedCompletionSource.TrySetCanceled(_cancellationToken);
             }
 
             return _cancellationToken.IsCancellationRequested;
         }
 
-        private void TrySetException(Exception exception)
-        {
-            DisposeProcess();
-            _tcs.TrySetException(exception);
-        }
-
         private void DisposeProcess()
         {
+            _process.OutputDataReceived -= HandleStdOutDataReceived;
+            _process.ErrorDataReceived -= HandleStdErrDataReceived;
+
             _process.Dispose();
             _ctRegistration.Dispose();
             _timeoutCts?.Dispose();

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

@@ -26,6 +26,9 @@ public ProcessWrapper(ProcessStartInfo processStartInfo)
                     StartInfo = processStartInfo,
                     EnableRaisingEvents = true
                 };
+
+                _process.OutputDataReceived += HandleOutputDataReceived;
+                _process.ErrorDataReceived += HandleErrorDataReceived;
             }
 
             public bool HasExited => _process.HasExited;
@@ -45,7 +48,20 @@ public event EventHandler Exited
                 remove => _process.Exited -= value;
             }
 
-            public void Start() => _process.Start();
+            public event DataReceivedEventWrapperHandler OutputDataReceived;
+            public event DataReceivedEventWrapperHandler ErrorDataReceived;
+
+            private void HandleErrorDataReceived(object sender, DataReceivedEventArgs e) => ErrorDataReceived?.Invoke(this, new DataReceivedEventArgsWrapper(e));
+            private void HandleOutputDataReceived(object sender, DataReceivedEventArgs e) => OutputDataReceived?.Invoke(this, new DataReceivedEventArgsWrapper(e));
+
+            public void Start()
+            {
+                _process.Start();
+
+                _process.BeginOutputReadLine();
+                _process.BeginErrorReadLine();
+            }
+
             public void Kill() => _process.Kill();
             public void Dispose() => _process.Dispose();
         }

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

@@ -4,6 +4,7 @@
 using System;
 using System.ComponentModel;
 using System.Diagnostics;
+using System.Linq;
 using System.Runtime.InteropServices;
 using System.Text;
 using System.Threading;
@@ -45,6 +46,17 @@ public async Task ProcessRunnerRealProcessSucceeded()
             Assert.AreEqual(output, result.TrimEnd());
         }
 
+        [Test]
+        public async Task ProcessRunnerRealProcessLargeOutputSucceeded()
+        {
+            var output = string.Concat(Enumerable.Repeat("ab", 3000));
+            var process = CreateRealProcess($"echo {output}", $"echo {output}");
+            var runner = new ProcessRunner(process, TimeSpan.FromSeconds(30), default);
+            var result = await Run(runner);
+
+            Assert.AreEqual(output, result.TrimEnd());
+        }
+
         [Test]
         public void ProcessRunnerCanceledByTimeout()
         {

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

@@ -70,6 +70,8 @@ public int ExitCode
 
         public event EventHandler Exited;
         public event EventHandler Started;
+        public event DataReceivedEventWrapperHandler OutputDataReceived;
+        public event DataReceivedEventWrapperHandler ErrorDataReceived;
 
         public void Start()
         {
@@ -101,8 +103,23 @@ private void FinishProcessRun(Task delayTask)
         private void FinishProcessRun()
         {
             WriteToStream(Output, _outputStreamWriter);
+
+            if (Output != default)
+            {
+                OutputDataReceived?.Invoke(this, new DataReceivedEventArgsWrapper(Output));
+            }
+
             WriteToStream(Error, _errorStreamWriter);
 
+            if (Error != default)
+            {
+                ErrorDataReceived?.Invoke(this, new DataReceivedEventArgsWrapper(Error));
+            }
+
+            // signal completion
+            OutputDataReceived?.Invoke(this, new DataReceivedEventArgsWrapper((string)null));
+            ErrorDataReceived?.Invoke(this, new DataReceivedEventArgsWrapper((string)null));
+
             _hasExited = true;
             _exitCode = CodeOnExit ?? (Error != default ? 1 : 0);
             Exited?.Invoke(this, EventArgs.Empty);