From 82341b0e3bddfab4e82a1a7635a9e56329564601 Mon Sep 17 00:00:00 2001
From: Jordan Borean <jborean93@gmail.com>
Date: Mon, 2 Sep 2024 12:28:36 +1000
Subject: [PATCH] Add remote runspace check for PushRunspace

Add explicit exception when attempting to call PushRunspace on a
runspace that is not remote. Instead of failing with a hard to
understand NullReferenceException an ArgumentException is raised with
information about what went wrong.
---
 .../host/msh/ConsoleHost.cs                       |  9 +++++++--
 .../resources/ConsoleHostStrings.resx             |  3 +++
 .../engine/hostifaces/MshHost.cs                  |  5 +++++
 test/powershell/Host/HostUtilities.Tests.ps1      | 15 +++++++++++++++
 4 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/Microsoft.PowerShell.ConsoleHost/host/msh/ConsoleHost.cs b/src/Microsoft.PowerShell.ConsoleHost/host/msh/ConsoleHost.cs
index 8cc7ee00f5..747d8415ce 100644
--- a/src/Microsoft.PowerShell.ConsoleHost/host/msh/ConsoleHost.cs
+++ b/src/Microsoft.PowerShell.ConsoleHost/host/msh/ConsoleHost.cs
@@ -622,14 +622,19 @@ namespace Microsoft.PowerShell
         /// <summary>
         /// See base class.
         /// </summary>
-        public void PushRunspace(Runspace newRunspace)
+        public void PushRunspace(Runspace runspace)
         {
             if (_runspaceRef == null)
             {
                 return;
             }
 
-            RemoteRunspace remoteRunspace = newRunspace as RemoteRunspace;
+            RemoteRunspace remoteRunspace = runspace as RemoteRunspace;
+            if (remoteRunspace is null)
+            {
+                throw new ArgumentException(ConsoleHostStrings.PushRunspaceNotRemote, nameof(runspace));
+            }
+
             Dbg.Assert(remoteRunspace != null, "Expected remoteRunspace != null");
             remoteRunspace.StateChanged += HandleRemoteRunspaceStateChanged;
 
diff --git a/src/Microsoft.PowerShell.ConsoleHost/resources/ConsoleHostStrings.resx b/src/Microsoft.PowerShell.ConsoleHost/resources/ConsoleHostStrings.resx
index ce124ec084..80b3d4aafe 100644
--- a/src/Microsoft.PowerShell.ConsoleHost/resources/ConsoleHostStrings.resx
+++ b/src/Microsoft.PowerShell.ConsoleHost/resources/ConsoleHostStrings.resx
@@ -182,4 +182,7 @@ The current session does not support debugging; execution will continue.
   <data name="RunAsAdministrator" xml:space="preserve">
     <value>Run as Administrator</value>
   </data>
+  <data name="PushRunspaceNotRemote" xml:space="preserve">
+    <value>PushRunspace can only push a remote runspace.</value>
+  </data>
 </root>
diff --git a/src/System.Management.Automation/engine/hostifaces/MshHost.cs b/src/System.Management.Automation/engine/hostifaces/MshHost.cs
index 32e07f48f3..bfa79e89f1 100644
--- a/src/System.Management.Automation/engine/hostifaces/MshHost.cs
+++ b/src/System.Management.Automation/engine/hostifaces/MshHost.cs
@@ -296,6 +296,11 @@ namespace System.Management.Automation.Host
         /// <summary>
         /// Called by the engine to notify the host that a runspace push has been requested.
         /// </summary>
+        /// <param name="runspace">
+        /// The runspace to push. This runspace must be a remote runspace and
+        /// not a locally created runspace.
+        /// </param>
+        /// <exception cref="ArgumentException">The specified runspace is not a remote runspace.</exception>
         /// <seealso cref="System.Management.Automation.Host.IHostSupportsInteractiveSession.PushRunspace"/>
         [SuppressMessage("Microsoft.Naming", "CA1704:IdentifiersShouldBeSpelledCorrectly", MessageId = "Runspace")]
         [SuppressMessage("Microsoft.Naming", "CA1704:IdentifiersShouldBeSpelledCorrectly", MessageId = "runspace")]
diff --git a/test/powershell/Host/HostUtilities.Tests.ps1 b/test/powershell/Host/HostUtilities.Tests.ps1
index da65dc8162..e151a2cfc2 100644
--- a/test/powershell/Host/HostUtilities.Tests.ps1
+++ b/test/powershell/Host/HostUtilities.Tests.ps1
@@ -81,3 +81,18 @@ Describe 'PromptForCredential' -Tags "CI" {
         $out.UserName | Should -BeExactly 'myDomain\myUser'
     }
 }
+
+Describe 'PushRunspaceLocalFailure' -Tags 'CI' {
+    It 'Should throw an exception when pushing a local runspace' {
+        $runspace = [RunspaceFactory]::CreateRunspace()
+        try {
+            $runspace.Open()
+            $exc = { $Host.PushRunspace($runspace) } | Should -Throw -PassThru
+            $exc.Exception.InnerException | Should -BeOfType ([System.ArgumentException])
+            [string]$exc | Should -BeLike "*PushRunspace can only push a remote runspace. (Parameter 'runspace')*"
+        }
+        finally {
+            $runspace.Dispose()
+        }
+    }
+}