Update PATH environment variable for package manager executable on Windows (PowerShell#25847)

Implementation for PowerShell/PowerShell-RFC#391. A few implementation decisions:

1. Check the package manager list only once per process (**same as CMD**). So, changing the package manager list in registry doesn't affect a PowerShell session that is already running.

2. When detecting user/system PATH value changes, only check new items PREPENDED or APPENDED to the original string (**same as CMD**). The detection is made simple, which is not intended to detect complex changes to the user/system PATH.

3. The PATH update feature is disabled when the `Environment` provider is not available, or when the current session is restricted.
    - Both `Environment` provider and `IsSessionRestricted` are checked only once per PowerShell session.
    - Use `ConditionalWeakTable` to cache whether the feature is enabled or not for a given session.
    - We cannot use `thread-static` variable for the caching because it's possible for Runspace to reuse threads.

Author: daxian-dbw

src/System.Management.Automation/engine/NativeCommandProcessor.cs

@@ -11,15 +11,17 @@
 using System.Diagnostics.CodeAnalysis;
 using System.Globalization;
 using System.IO;
+using System.Linq;
 using System.Management.Automation.Internal;
-using System.Management.Automation.Runspaces;
-using System.Runtime.InteropServices;
+using System.Runtime.CompilerServices;
 using System.Runtime.Serialization;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using System.Xml;
+using Microsoft.PowerShell.Commands;
 using Microsoft.PowerShell.Telemetry;
+using Microsoft.Win32;
 using Dbg = System.Management.Automation.Diagnostics;
 
 namespace System.Management.Automation
@@ -194,27 +196,211 @@ internal NativeCommandExitException(string path, int exitCode, int processId, st
     /// </summary>
     internal class NativeCommandProcessor : CommandProcessorBase
     {
-        // This is the list of files which will trigger Legacy behavior if
-        // PSNativeCommandArgumentPassing is set to "Windows".
-        private static readonly IReadOnlySet<string> s_legacyFileExtensions = new HashSet<string>(StringComparer.OrdinalIgnoreCase)
+        /// <summary>
+        /// This is the list of files which will trigger Legacy behavior if 'PSNativeCommandArgumentPassing' is set to "Windows".
+        /// </summary>
+        private static readonly HashSet<string> s_legacyFileExtensions = new(StringComparer.OrdinalIgnoreCase)
+            {
+                ".js",
+                ".wsf",
+                ".cmd",
+                ".bat",
+                ".vbs",
+            };
+
+        /// <summary>
+        /// This is the list of native commands that have non-standard behavior with regard to argument passing.
+        /// We use Legacy argument parsing for them when 'PSNativeCommandArgumentPassing' is set to "Windows".
+        /// </summary>
+        private static readonly HashSet<string> s_legacyCommands = new(StringComparer.OrdinalIgnoreCase)
+            {
+                "cmd",
+                "cscript",
+                "find",
+                "sqlcmd",
+                "wscript",
+            };
+
+#if !UNIX
+        /// <summary>
+        /// List of known package managers pulled from the registry.
+        /// </summary>
+        private static readonly HashSet<string> s_knownPackageManagers = GetPackageManagerListFromRegistry();
+
+        /// <summary>
+        /// Indicates whether the Path Update feature is enabled in a given session.
+        /// PowerShell sessions could reuse the same thread, so we cannot cache the value with a thread static variable.
+        /// </summary>
+        private static readonly ConditionalWeakTable<ExecutionContext, string> s_pathUpdateFeatureEnabled = new();
+
+        private readonly bool _isPackageManager;
+        private string _originalUserEnvPath;
+        private string _originalSystemEnvPath;
+
+        /// <summary>
+        /// Gets the known package managers from the registry.
+        /// </summary>
+        private static HashSet<string> GetPackageManagerListFromRegistry()
+        {
+            // We only account for the first 8 package managers. This is the same behavior as in CMD.
+            const int MaxPackageManagerCount = 8;
+            const string RegKeyPath = @"Software\Microsoft\Command Processor\KnownPackageManagers";
+
+            string[] subKeyNames = null;
+            HashSet<string> retSet = null;
+
+            try
+            {
+                using RegistryKey key = Registry.LocalMachine.OpenSubKey(RegKeyPath);
+                subKeyNames = key?.GetSubKeyNames();
+            }
+            catch
+            {
+                return null;
+            }
+
+            if (subKeyNames is { Length: > 0 })
+            {
+                IEnumerable<string> names = subKeyNames.Length <= MaxPackageManagerCount
+                    ? subKeyNames
+                    : subKeyNames.Take(MaxPackageManagerCount);
+
+                retSet = new(names, StringComparer.OrdinalIgnoreCase);
+            }
+
+            return retSet;
+        }
+
+        /// <summary>
+        /// Check if the given name is a known package manager from the registry list.
+        /// </summary>
+        private static bool IsKnownPackageManager(string name)
         {
-            ".js",
-            ".wsf",
-            ".cmd",
-            ".bat",
-            ".vbs",
-        };
-
-        // The following native commands have non-standard behavior with regard to argument passing,
-        // so we use Legacy argument parsing for them when PSNativeCommandArgumentPassing is set to Windows.
-        private static readonly IReadOnlySet<string> s_legacyCommands = new HashSet<string>(StringComparer.OrdinalIgnoreCase)
+            if (s_knownPackageManagers is null)
+            {
+                return false;
+            }
+
+            if (s_knownPackageManagers.Contains(name))
+            {
+                return true;
+            }
+
+            int lastDotIndex = name.LastIndexOf('.');
+            if (lastDotIndex > 0)
+            {
+                string nameWithoutExt = name[..lastDotIndex];
+                if (s_knownPackageManagers.Contains(nameWithoutExt))
+                {
+                    return true;
+                }
+            }
+
+            return false;
+        }
+
+        /// <summary>
+        /// Check if the Path Update feature is enabled for the given session.
+        /// </summary>
+        private static bool IsPathUpdateFeatureEnabled(ExecutionContext context)
         {
-            "cmd",
-            "cscript",
-            "find",
-            "sqlcmd",
-            "wscript",
-        };
+            // We check only once per session.
+            if (s_pathUpdateFeatureEnabled.TryGetValue(context, out string value))
+            {
+                // The feature is enabled if the value is not null.
+                return value is { };
+            }
+
+            // Disable Path Update if 'EnvironmentProvider' is disabled in the current session, or the current session is restricted.
+            bool enabled = context.EngineSessionState.Providers.ContainsKey(EnvironmentProvider.ProviderName)
+                && !Utils.IsSessionRestricted(context);
+
+            // - Use the static empty string instance to indicate that the feature is enabled.
+            // - Use the null value to indicate that the feature is disabled.
+            s_pathUpdateFeatureEnabled.TryAdd(context, enabled ? string.Empty : null);
+            return enabled;
+        }
+
+        /// <summary>
+        /// Gets the added part of the new string compared to the old string.
+        /// </summary>
+        private static ReadOnlySpan<char> GetAddedPartOfString(string oldString, string newString)
+        {
+            if (oldString.Length >= newString.Length)
+            {
+                // Nothing added or something removed.
+                return ReadOnlySpan<char>.Empty;
+            }
+
+            int index = newString.IndexOf(oldString);
+            if (index is -1)
+            {
+                // The new and old strings are drastically different. Stop trying in this case.
+                return ReadOnlySpan<char>.Empty;
+            }
+
+            if (index > 0)
+            {
+                // Found the old string at non-zero offset, so something was prepended to the old string.
+                return newString.AsSpan(0, index);
+            }
+            else
+            {
+                // Found the old string at the beginning of the new string, so something was appended to the old string.
+                return newString.AsSpan(oldString.Length);
+            }
+        }
+
+        /// <summary>
+        /// Update the process-scope environment variable Path based on the changes in the user-scope and system-scope Path.
+        /// </summary>
+        /// <param name="oldUserPath">The old value of the user-scope Path retrieved from registry.</param>
+        /// <param name="oldSystemPath">The old value of the system-scope Path retrieved from registry.</param>
+        private static void UpdateProcessEnvPath(string oldUserPath, string oldSystemPath)
+        {
+            string newUserEnvPath = Environment.GetEnvironmentVariable("Path", EnvironmentVariableTarget.User);
+            string newSystemEnvPath = Environment.GetEnvironmentVariable("Path", EnvironmentVariableTarget.Machine);
+            string procEnvPath = Environment.GetEnvironmentVariable("Path");
+
+            ReadOnlySpan<char> userPathChange = GetAddedPartOfString(oldUserPath, newUserEnvPath).Trim(';');
+            ReadOnlySpan<char> systemPathChange = GetAddedPartOfString(oldSystemPath, newSystemEnvPath).Trim(';');
+
+            // Add 2 to account for the path separators we may need to add.
+            int maxLength = procEnvPath.Length + userPathChange.Length + systemPathChange.Length + 2;
+            StringBuilder newPath = null;
+
+            if (userPathChange.Length > 0)
+            {
+                CreateNewProcEnvPath(userPathChange);
+            }
+
+            if (systemPathChange.Length > 0)
+            {
+                CreateNewProcEnvPath(systemPathChange);
+            }
+
+            if (newPath is { Length: > 0 })
+            {
+                // Update the process env Path.
+                Environment.SetEnvironmentVariable("Path", newPath.ToString());
+            }
+
+            // Helper method to create a new env Path string.
+            void CreateNewProcEnvPath(ReadOnlySpan<char> newChange)
+            {
+                newPath ??= new StringBuilder(procEnvPath, capacity: maxLength);
+
+                if (newPath.Length is 0 || newPath[^1] is ';')
+                {
+                    newPath.Append(newChange);
+                }
+                else
+                {
+                    newPath.Append(';').Append(newChange);
+                }
+            }
+        }
+#endif
 
         #region ctor/native command properties
 
@@ -262,7 +448,11 @@ internal NativeCommandProcessor(ApplicationInfo applicationInfo, ExecutionContex
             // Create input writer for providing input to the process.
             _inputWriter = new ProcessInputWriter(Command);
 
-            _isTranscribing = this.Command.Context.EngineHostInterface.UI.IsTranscribing;
+            _isTranscribing = context.EngineHostInterface.UI.IsTranscribing;
+
+#if !UNIX
+            _isPackageManager = IsKnownPackageManager(_applicationInfo.Name) && IsPathUpdateFeatureEnabled(context);
+#endif
         }
 
         /// <summary>
@@ -418,7 +608,7 @@ internal override void ProcessRecord()
         /// <summary>
         /// Process object for the invoked application.
         /// </summary>
-        private System.Diagnostics.Process _nativeProcess;
+        private Process _nativeProcess;
 
         /// <summary>
         /// This is used for writing input to the process.
@@ -560,6 +750,12 @@ private void InitNativeProcess()
                 // must set UseShellExecute to false if we modify the environment block
                 startInfo.UseShellExecute = false;
             }
+
+            if (_isPackageManager)
+            {
+                _originalUserEnvPath = Environment.GetEnvironmentVariable("Path", EnvironmentVariableTarget.User);
+                _originalSystemEnvPath = Environment.GetEnvironmentVariable("Path", EnvironmentVariableTarget.Machine);
+            }
 #endif
 
             if (this.Command.Context.CurrentPipelineStopping)
@@ -898,6 +1094,13 @@ internal override void Complete()
                     ConsumeAvailableNativeProcessOutput(blocking: true);
                     _nativeProcess.WaitForExit();
 
+#if !UNIX
+                    if (_isPackageManager)
+                    {
+                        UpdateProcessEnvPath(_originalUserEnvPath, _originalSystemEnvPath);
+                    }
+#endif
+
                     // Capture screen output if we are transcribing and running stand alone
                     if (_isTranscribing && (s_supportScreenScrape == true) && _runStandAlone)
                     {
@@ -1717,6 +1920,7 @@ private bool IsExecutable(string path)
         #region Minishell Interop
 
         private bool _isMiniShell = false;
+
         /// <summary>
         /// Returns true if native command being invoked is mini-shell.
         /// </summary>

src/System.Management.Automation/engine/Utils.cs

@@ -1539,14 +1539,14 @@ internal static string DisplayHumanReadableFileSize(long bytes)
         /// <returns>True if the session is restricted.</returns>
         internal static bool IsSessionRestricted(ExecutionContext context)
         {
-                CmdletInfo cmdletInfo = context.SessionState.InvokeCommand.GetCmdlet("Microsoft.PowerShell.Core\\Import-Module");
-                // if import-module is visible, then the session is not restricted,
-                // because the user can load arbitrary code.
-                if (cmdletInfo != null && cmdletInfo.Visibility == SessionStateEntryVisibility.Public)
-                {
-                   return false;
-                }
-                return true;
+            CmdletInfo cmdletInfo = context.SessionState.InvokeCommand.GetCmdlet("Microsoft.PowerShell.Core\\Import-Module");
+            // if import-module is visible, then the session is not restricted,
+            // because the user can load arbitrary code.
+            if (cmdletInfo != null && cmdletInfo.Visibility == SessionStateEntryVisibility.Public)
+            {
+                return false;
+            }
+            return true;
         }
     }
 }