From f76b2fcbafb92c784aa8a200e412353f7e864824 Mon Sep 17 00:00:00 2001
From: PetSerAl <PetSerAl@gmail.com>
Date: Sun, 2 Apr 2017 02:36:04 +0300
Subject: [PATCH] Make `Out-Default -Transcript` more robust in how it handles
 `TranscribeOnly` state (#3436)

---
 .../FormatAndOutput/out-console/OutConsole.cs | 23 ++++++++---
 .../engine/hostifaces/MshHostUserInterface.cs | 25 +++++++++++-
 .../Out-Default.Tests.ps1                     | 40 +++++++++++++++++++
 3 files changed, 82 insertions(+), 6 deletions(-)
 create mode 100644 test/powershell/Modules/Microsoft.PowerShell.Core/Out-Default.Tests.ps1

diff --git a/src/System.Management.Automation/commands/utility/FormatAndOutput/out-console/OutConsole.cs b/src/System.Management.Automation/commands/utility/FormatAndOutput/out-console/OutConsole.cs
index 484e99674..f685c43f3 100644
--- a/src/System.Management.Automation/commands/utility/FormatAndOutput/out-console/OutConsole.cs
+++ b/src/System.Management.Automation/commands/utility/FormatAndOutput/out-console/OutConsole.cs
@@ -79,10 +79,9 @@ namespace Microsoft.PowerShell.Commands
                 mrt.MergeUnclaimedPreviousErrorResults = true;
             }
 
-            _savedTranscribeOnly = Host.UI.TranscribeOnly;
             if (Transcript)
             {
-                Host.UI.TranscribeOnly = true;
+                _transcribeOnlyCookie = Host.UI.SetTranscribeOnly();
             }
 
             // This needs to be done directly through the command runtime, as Out-Default
@@ -143,15 +142,29 @@ namespace Microsoft.PowerShell.Commands
             }
 
             base.EndProcessing();
+        }
 
-            if (Transcript)
+        /// <summary>
+        /// Revert transcription state on Dispose
+        /// </summary>
+        protected override void InternalDispose()
+        {
+            try
             {
-                Host.UI.TranscribeOnly = _savedTranscribeOnly;
+                base.InternalDispose();
+            }
+            finally
+            {
+                if (_transcribeOnlyCookie != null)
+                {
+                    _transcribeOnlyCookie.Dispose();
+                    _transcribeOnlyCookie = null;
+                }
             }
         }
 
         private ArrayList _outVarResults = null;
-        private bool _savedTranscribeOnly = false;
+        private IDisposable _transcribeOnlyCookie = null;
     }
 
     /// <summary>
diff --git a/src/System.Management.Automation/engine/hostifaces/MshHostUserInterface.cs b/src/System.Management.Automation/engine/hostifaces/MshHostUserInterface.cs
index 8249eb33d..fc3bc4f97 100644
--- a/src/System.Management.Automation/engine/hostifaces/MshHostUserInterface.cs
+++ b/src/System.Management.Automation/engine/hostifaces/MshHostUserInterface.cs
@@ -11,6 +11,7 @@ using System.Security;
 using System.Globalization;
 using System.Management.Automation.Runspaces;
 using Microsoft.PowerShell.Commands;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace System.Management.Automation.Host
@@ -390,7 +391,29 @@ namespace System.Management.Automation.Host
         /// so that when content is sent through Out-Default it doesn't
         /// make it to the actual host.
         /// </summary>
-        internal bool TranscribeOnly { get; set; }
+        internal bool TranscribeOnly => Interlocked.CompareExchange(ref _transcribeOnlyCount, 0, 0) != 0;
+        private int _transcribeOnlyCount = 0;
+        internal IDisposable SetTranscribeOnly() => new TranscribeOnlyCookie(this);
+        private sealed class TranscribeOnlyCookie : IDisposable
+        {
+            private PSHostUserInterface _ui;
+            private bool _disposed = false;
+            public TranscribeOnlyCookie(PSHostUserInterface ui)
+            {
+                _ui=ui;
+                Interlocked.Increment(ref _ui._transcribeOnlyCount);
+            }
+            public void Dispose()
+            {
+                if (!_disposed)
+                {
+                    Interlocked.Decrement(ref _ui._transcribeOnlyCount);
+                    _disposed = true;
+                    GC.SuppressFinalize(this);
+                }
+            }
+            ~TranscribeOnlyCookie() => Dispose();
+        }
 
         /// <summary>
         /// Flag to determine whether the host is transcribing.
diff --git a/test/powershell/Modules/Microsoft.PowerShell.Core/Out-Default.Tests.ps1 b/test/powershell/Modules/Microsoft.PowerShell.Core/Out-Default.Tests.ps1
new file mode 100644
index 000000000..fa8ddc3dd
--- /dev/null
+++ b/test/powershell/Modules/Microsoft.PowerShell.Core/Out-Default.Tests.ps1
@@ -0,0 +1,40 @@
+Describe "Out-Default Tests" -tag CI {
+    BeforeAll {
+        # due to https://github.com/PowerShell/PowerShell/issues/3405, `Out-Default -Transcript` emits output to pipeline
+        # as running in Pester effectively wraps everything in parenthesis, workaround is to use another powershell
+        # to run the test script passed as a string
+        $powershell = "$PSHOME/powershell"
+    }
+
+    It "'Out-Default -Transcript' shows up in transcript, but not host" {
+        $script = @"
+            `$null = Start-Transcript -Path "$testdrive\transcript.txt";
+            'hello' | Microsoft.PowerShell.Core\Out-Default -Transcript;
+            'bye';
+            `$null = Stop-Transcript
+"@
+
+        & $powershell -c $script | Should BeExactly 'bye'
+        "TestDrive:\transcript.txt" | Should Contain 'hello'
+    }
+
+    It "Out-Default reverts transcription state when used more than once in a pipeline" {
+        & $powershell -c "Out-Default -Transcript | Out-Default -Transcript; 'Hello'" | Should BeExactly "Hello"
+    }
+
+    It "Out-Default reverts transcription state when exception occurs in pipeline" {
+        & $powershell -c "try { & { throw } | Out-Default -Transcript } catch {}; 'Hello'" | Should BeExactly "Hello"
+    }
+
+    It "Out-Default reverts transcription state even if Dispose() isn't called" {
+        $script = @"
+            `$sp = {Out-Default -Transcript}.GetSteppablePipeline();
+            `$sp.Begin(`$false);
+            `$sp = `$null;
+            [GC]::Collect();
+            [GC]::WaitForPendingFinalizers();
+            'hello'
+"@
+        & $powershell -c $script | Should BeExactly 'hello'
+    }
+}