From e1495e1509227d9a03399d202665029094f89fbe Mon Sep 17 00:00:00 2001
From: Matej Kafka <git@a-n.cc>
Date: Mon, 8 Apr 2024 20:30:12 +0200
Subject: [PATCH] InitialSessionState: Add support for configuring initial
 working directory (#17603)

---
 .../engine/InitialSessionState.cs             | 33 +++++++++++++++----
 1 file changed, 26 insertions(+), 7 deletions(-)

diff --git a/src/System.Management.Automation/engine/InitialSessionState.cs b/src/System.Management.Automation/engine/InitialSessionState.cs
index c30e199b2b1b..0c6e373c5fef 100644
--- a/src/System.Management.Automation/engine/InitialSessionState.cs
+++ b/src/System.Management.Automation/engine/InitialSessionState.cs
@@ -1678,6 +1678,7 @@ public InitialSessionState Clone()
             ss.ThreadOptions = this.ThreadOptions;
             ss.ThrowOnRunspaceOpenError = this.ThrowOnRunspaceOpenError;
             ss.ApartmentState = this.ApartmentState;
+            ss.WorkingDirectory = this.WorkingDirectory;
 
             foreach (ModuleSpecification modSpec in this.ModuleSpecificationsToImport)
             {
@@ -1831,6 +1832,12 @@ public Microsoft.PowerShell.ExecutionPolicy ExecutionPolicy
         /// </summary>
         public bool ThrowOnRunspaceOpenError { get; set; } = false;
 
+        /// <summary>
+        /// If not null, the working directory of the runspace is set to this path. If null,
+        /// the <see cref="Environment.CurrentDirectory">process working directory</see> is used as a default.
+        /// </summary>
+        public string WorkingDirectory { get; set; } = null;
+
         /// <summary>
         /// This property will be set only if we are refreshing the Type/Format settings by calling UpdateTypes/UpdateFormats directly.
         /// In this case, we should wait until all type/format entries get processed. After that, if there were errors
@@ -2128,6 +2135,7 @@ public virtual HashSet<string> StartupScripts
 
         private readonly object _syncObject = new object();
 
+        /// <exception cref="ItemNotFoundException">User configured a custom WorkingDirectory, which does not exist in the configured providers.</exception>
         internal void Bind(ExecutionContext context, bool updateOnly, PSModuleInfo module, bool noClobber, bool local, bool setLocation)
         {
             Host = context.EngineHostInterface;
@@ -2231,7 +2239,7 @@ internal void Bind(ExecutionContext context, bool updateOnly, PSModuleInfo modul
                 }
             }
 
-            SetSessionStateDrive(context, setLocation: setLocation);
+            SetSessionStateDrive(context, setLocation, WorkingDirectory);
         }
 
         private void Bind_SetVariables(SessionStateInternal ss)
@@ -3331,6 +3339,7 @@ private RunspaceOpenModuleLoadException ValidateAndReturnRunspaceOpenModuleLoadE
         /// This allows for runspace reuse with minimal chance for contamination.
         /// </summary>
         /// <param name="context"></param>
+        /// <exception cref="ItemNotFoundException">User configured a custom WorkingDirectory, which does not exist in the configured providers.</exception>
         internal void ResetRunspaceState(ExecutionContext context)
         {
             lock (_syncObject)
@@ -3370,7 +3379,7 @@ internal void ResetRunspaceState(ExecutionContext context)
                 InitialSessionState.CreateQuestionVariable(context);
 
                 // Reset the path for this runspace.
-                SetSessionStateDrive(context, true);
+                SetSessionStateDrive(context, true, WorkingDirectory);
 
                 // Reset the event, transaction and debug managers.
                 context.ResetManagers();
@@ -3381,7 +3390,7 @@ internal void ResetRunspaceState(ExecutionContext context)
             }
         }
 
-        private static bool TryInitSessionStateCurrentDrive(ExecutionContext context)
+        private static bool TryInitSessionStateCurrentDrive(ExecutionContext context, string location)
         {
             try
             {
@@ -3408,13 +3417,14 @@ private static bool TryInitSessionStateCurrentDrive(ExecutionContext context)
             }
             else
             {
-                ItemNotFoundException itemNotFound = new(Environment.CurrentDirectory, "PathNotFound", SessionStateStrings.PathNotFound);
+                ItemNotFoundException itemNotFound = new(location ?? Environment.CurrentDirectory, "PathNotFound", SessionStateStrings.PathNotFound);
                 context.ReportEngineStartupError(itemNotFound);
                 return false;
             }
         }
 
-        internal static void SetSessionStateDrive(ExecutionContext context, bool setLocation)
+        /// <exception cref="ItemNotFoundException">User configured a custom WorkingDirectory, which does not exist in the configured providers.</exception>
+        internal static void SetSessionStateDrive(ExecutionContext context, bool setLocation, string location)
         {
             if (context.EngineSessionState.ProviderCount == 0)
             {
@@ -3424,7 +3434,7 @@ internal static void SetSessionStateDrive(ExecutionContext context, bool setLoca
             // NTRAID#Windows Out Of Band Releases-908481-2005/07/01-JeffJon
             // Make sure we have a CurrentDrive set so that we can deal with
             // UNC paths
-            if (context.EngineSessionState.CurrentDrive == null && !TryInitSessionStateCurrentDrive(context))
+            if (context.EngineSessionState.CurrentDrive == null && !TryInitSessionStateCurrentDrive(context, location))
             {
                 return;
             }
@@ -3436,7 +3446,16 @@ internal static void SetSessionStateDrive(ExecutionContext context, bool setLoca
 
             var providerContext = new CmdletProviderContext(context) { SuppressWildcardExpansion = true };
 
-            // Set the starting location to the current process working directory
+            // User set a custom initial working directory.
+            if (location != null)
+            {
+                // If the location is invalid or does not exist, let the exception bubble up; since the user explicitly
+                // configured the working directory, he probably wants to get notified on failure.
+                context.EngineSessionState.SetLocation(location, providerContext);
+                return;
+            }
+
+            // As a fallback, set the starting location to the current process working directory.
             // Ignore any errors as the file system provider may not be loaded or
             // a drive with the same name as the real file system drive may not have
             // been mounted.