diff --git a/PSReadLine/Cmdlets.cs b/PSReadLine/Cmdlets.cs index bc89cc1f9..7771fa920 100644 --- a/PSReadLine/Cmdlets.cs +++ b/PSReadLine/Cmdlets.cs @@ -502,6 +502,8 @@ public object ListPredictionTooltipColor set => _listPredictionTooltipColor = VTColorUtils.AsEscapeSequence(value); } + public bool TerminateOrphanedConsoleApps { get; set; } + internal string _defaultTokenColor; internal string _commentColor; internal string _keywordColor; @@ -808,6 +810,14 @@ public PredictionViewStyle PredictionViewStyle [Parameter] public Hashtable Colors { get; set; } + [Parameter] + public SwitchParameter TerminateOrphanedConsoleApps + { + get => _terminateOrphanedConsoleApps.GetValueOrDefault(); + set => _terminateOrphanedConsoleApps = value; + } + internal SwitchParameter? _terminateOrphanedConsoleApps; + [ExcludeFromCodeCoverage] protected override void EndProcessing() { diff --git a/PSReadLine/Options.cs b/PSReadLine/Options.cs index 19f366513..7485154b4 100644 --- a/PSReadLine/Options.cs +++ b/PSReadLine/Options.cs @@ -6,8 +6,10 @@ using System.Collections; using System.Collections.Generic; using System.Diagnostics.CodeAnalysis; +using System.Globalization; using System.Management.Automation; using System.Reflection; +using System.Runtime.InteropServices; using System.Threading; using Microsoft.PowerShell.PSReadLine; @@ -167,6 +169,22 @@ private void SetOptionsInternal(SetPSReadLineOption options) } } } + if (options._terminateOrphanedConsoleApps.HasValue) + { + if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows)) + { + Options.TerminateOrphanedConsoleApps = options.TerminateOrphanedConsoleApps; + PlatformWindows.SetTerminateOrphanedConsoleApps(Options.TerminateOrphanedConsoleApps); + } + else + { + throw new PlatformNotSupportedException( + string.Format( + CultureInfo.CurrentUICulture, + PSReadLineResources.OptionNotSupportedOnNonWindows, + nameof(Options.TerminateOrphanedConsoleApps))); + } + } } private void SetKeyHandlerInternal(string[] keys, Action handler, string briefDescription, string longDescription, ScriptBlock scriptBlock) diff --git a/PSReadLine/PSReadLine.format.ps1xml b/PSReadLine/PSReadLine.format.ps1xml index 5b20c58af..24f70799a 100644 --- a/PSReadLine/PSReadLine.format.ps1xml +++ b/PSReadLine/PSReadLine.format.ps1xml @@ -164,6 +164,9 @@ $d = [Microsoft.PowerShell.KeyHandler]::GetGroupingDescription($_.Group) PredictionViewStyle + + TerminateOrphanedConsoleApps + [Microsoft.PowerShell.VTColorUtils]::FormatColor($_.CommandColor) diff --git a/PSReadLine/PSReadLineResources.Designer.cs b/PSReadLine/PSReadLineResources.Designer.cs index 7d3435945..ac672d3f7 100644 --- a/PSReadLine/PSReadLineResources.Designer.cs +++ b/PSReadLine/PSReadLineResources.Designer.cs @@ -2257,5 +2257,16 @@ internal static string UpcaseWordDescription return ResourceManager.GetString("UpcaseWordDescription", resourceCulture); } } + + /// + /// Looks up a localized string similar to: The '{0}' option is not supported on non-Windows platforms. + /// + internal static string OptionNotSupportedOnNonWindows + { + get + { + return ResourceManager.GetString("OptionNotSupportedOnNonWindows", resourceCulture); + } + } } } diff --git a/PSReadLine/PSReadLineResources.resx b/PSReadLine/PSReadLineResources.resx index 618ef22b9..e619e7c06 100644 --- a/PSReadLine/PSReadLineResources.resx +++ b/PSReadLine/PSReadLineResources.resx @@ -873,4 +873,7 @@ Or not saving history with: Find the next word starting from the current position and then make it upper case. + + The '{0}' option is not supported on non-Windows platforms. + diff --git a/PSReadLine/PlatformWindows.cs b/PSReadLine/PlatformWindows.cs index b28077ada..86b4a73c5 100644 --- a/PSReadLine/PlatformWindows.cs +++ b/PSReadLine/PlatformWindows.cs @@ -6,6 +6,7 @@ using System.Collections.Generic; using System.ComponentModel; using System.Diagnostics; +using System.Linq; using System.Runtime.InteropServices; using Microsoft.PowerShell; using Microsoft.PowerShell.Internal; @@ -201,6 +202,15 @@ internal static void Init(ref ICharMap charMap) EnableAnsiInput(ref charMap); } + // Is the TerminateOrphanedConsoleApps feature enabled? + if (_allowedPids != null) + { + // We are about to disable Ctrl+C signals... so if there are still any + // console-attached children, the shell will be broken until they are + // gone, so we'll get rid of them: + TerminateStragglers(); + } + SetOurInputMode(); } } @@ -289,11 +299,11 @@ internal static void CallUsingOurInputMode(Action a) } } - private static readonly Lazy _inputHandle = new Lazy(() => + private static SafeFileHandle OpenConsoleHandle(string name) { // We use CreateFile here instead of GetStdWin32Handle, as GetStdWin32Handle will return redirected handles var handle = CreateFile( - "CONIN$", + name, (uint)(AccessQualifiers.GenericRead | AccessQualifiers.GenericWrite), (uint)ShareModes.ShareWrite, (IntPtr)0, @@ -305,33 +315,14 @@ internal static void CallUsingOurInputMode(Action a) { int err = Marshal.GetLastWin32Error(); Win32Exception innerException = new Win32Exception(err); - throw new Exception("Failed to retrieve the input console handle.", innerException); + throw new Exception($"Failed to retrieve the console handle ({name}).", innerException); } return new SafeFileHandle(handle, true); - }); - - private static readonly Lazy _outputHandle = new Lazy(() => - { - // We use CreateFile here instead of GetStdWin32Handle, as GetStdWin32Handle will return redirected handles - var handle = CreateFile( - "CONOUT$", - (uint)(AccessQualifiers.GenericRead | AccessQualifiers.GenericWrite), - (uint)ShareModes.ShareWrite, - (IntPtr)0, - (uint)CreationDisposition.OpenExisting, - 0, - (IntPtr)0); - - if (handle == INVALID_HANDLE_VALUE) - { - int err = Marshal.GetLastWin32Error(); - Win32Exception innerException = new Win32Exception(err); - throw new Exception("Failed to retrieve the input console handle.", innerException); - } + } - return new SafeFileHandle(handle, true); - }); + private static readonly Lazy _inputHandle = new Lazy(() => OpenConsoleHandle("CONIN$")); + private static readonly Lazy _outputHandle = new Lazy(() => OpenConsoleHandle("CONOUT$")); [DllImport("kernel32.dll", CharSet = CharSet.Unicode, SetLastError = true)] private static extern bool GetConsoleMode(IntPtr hConsole, out uint dwMode); @@ -343,6 +334,13 @@ private static uint GetConsoleInputMode() return result; } + private static uint GetConsoleOutputMode() + { + var handle = _outputHandle.Value.DangerousGetHandle(); + GetConsoleMode(handle, out var result); + return result; + } + [DllImport("kernel32.dll", CharSet = CharSet.Unicode, SetLastError = true)] private static extern bool SetConsoleMode(IntPtr hConsole, uint dwMode); @@ -614,4 +612,410 @@ public override void BlankRestOfLine() [DllImport("user32.dll", SetLastError = true)] [return: MarshalAs(UnmanagedType.Bool)] internal static extern bool SystemParametersInfo(uint uiAction, uint uiParam, ref bool pvParam, uint fWinIni); + + [StructLayout(LayoutKind.Sequential)] + internal struct PROCESS_BASIC_INFORMATION + { + public IntPtr ExitStatus; + public IntPtr PebBaseAddress; + public IntPtr AffinityMask; + public IntPtr BasePriority; + public IntPtr UniqueProcessId; + public IntPtr InheritedFromUniqueProcessId; + } + + [DllImport("ntdll.dll")] + internal static extern int NtQueryInformationProcess( + IntPtr processHandle, + int processInformationClass, + out PROCESS_BASIC_INFORMATION processInformation, + int processInformationLength, + out int returnLength); + + internal const int InvalidProcessId = -1; + + internal static int GetParentPid(Process process) + { + // (This is how ProcessCodeMethods in pwsh does it.) + PROCESS_BASIC_INFORMATION pbi; + int size; + var res = NtQueryInformationProcess(process.Handle, 0, out pbi, Marshal.SizeOf(), out size); + + return res != 0 ? InvalidProcessId : pbi.InheritedFromUniqueProcessId.ToInt32(); + } + + [DllImport("kernel32.dll", SetLastError = true, EntryPoint = "GetConsoleProcessList")] + private static extern uint native_GetConsoleProcessList([In, Out] uint[] lpdwProcessList, uint dwProcessCount); + + private static uint[] GetConsoleProcessList() + { + int size = 100; + uint[] pids = new uint[size]; + uint numPids = native_GetConsoleProcessList(pids, (uint) size); + + if (numPids > size) + { + size = (int) numPids + 10; // a bit extra, since we may be racing attaches. + pids = new uint[size]; + numPids = native_GetConsoleProcessList(pids, (uint) size); + } + + if (0 == numPids || numPids > size) + { + return null; // no TerminateOrphanedConsoleApps for you, sorry + } + + Array.Resize(ref pids, (int) numPids); + return pids; + } + + // If the TerminateOrphanedConsoleApps option is enabled, this is the list of PIDs + // that are allowed to stay attached to the console (effectively the current process + // plus ancestors). + private static uint[] _allowedPids; + + internal static void SetTerminateOrphanedConsoleApps(bool enabled) + { + if (enabled) + { + _allowedPids = GetConsoleProcessList(); + } + else + { + _allowedPids = null; + } + } + + private static bool ItLooksLikeWeAreInTerminal() + { + return !String.IsNullOrEmpty(Environment.GetEnvironmentVariable("WT_SESSION")); + } + + [DllImport("kernel32.dll", SetLastError = true)] + private static extern IntPtr GetConsoleWindow(); + + internal enum TaskbarStates + { + NoProgress = 0, + Indeterminate = 0x1, + Normal = 0x2, + Error = 0x4, + Paused = 0x8, + } + + internal static class TaskbarProgress + { + [ComImport()] + [Guid("ea1afb91-9e28-4b86-90e9-9e9f8a5eefaf")] + [InterfaceType(ComInterfaceType.InterfaceIsIUnknown)] + private interface ITaskbarList3 + { + // ITaskbarList + [PreserveSig] + int HrInit(); + + [PreserveSig] + int AddTab(IntPtr hwnd); + + [PreserveSig] + int DeleteTab(IntPtr hwnd); + + [PreserveSig] + int ActivateTab(IntPtr hwnd); + + [PreserveSig] + int SetActiveAlt(IntPtr hwnd); + + // ITaskbarList2 + [PreserveSig] + int MarkFullscreenWindow(IntPtr hwnd, [MarshalAs(UnmanagedType.Bool)] bool fFullscreen); + + // ITaskbarList3 + [PreserveSig] + int SetProgressValue(IntPtr hwnd, UInt64 ullCompleted, UInt64 ullTotal); + + [PreserveSig] + int SetProgressState(IntPtr hwnd, TaskbarStates state); + + // N.B. for copy/pasters: we've left out the rest of the ITaskbarList3 methods... + } + + [ComImport()] + [Guid("56fdf344-fd6d-11d0-958a-006097c9a090")] + [ClassInterface(ClassInterfaceType.None)] + private class TaskbarInstance + { + } + + private static Lazy _taskbarInstance = new Lazy(() => (ITaskbarList3) new TaskbarInstance()); + + public static int SetProgressState(IntPtr windowHandle, TaskbarStates taskbarState) + { + return _taskbarInstance.Value.SetProgressState(windowHandle, taskbarState); + } + + public static int SetProgressValue(IntPtr windowHandle, int progressValue, int progressMax) + { + return _taskbarInstance.Value.SetProgressValue(windowHandle, (ulong) progressValue, (ulong) progressMax); + } + } + + private static readonly Lazy _myPid = new(() => + { + using var me = Process.GetCurrentProcess(); + return (uint)me.Id; + }); + + // Calculates what processes need to be terminated (populated into procsToTerminate), + // and returns the count. A "straggler" is a console-attached process (so GUI + // processes don't count) that is not in the _allowedPids list. + private static int GatherStragglers(List procsToTerminate) + { + procsToTerminate.Clear(); + + // These are the processes currently attached to this console. Note that GUI + // processes will not be attached to the console. + uint[] currentPids = GetConsoleProcessList(); + + foreach (var pid in currentPids) + { + if (!_allowedPids.Contains(pid)) + { + Process proc = null; + try + { + proc = Process.GetProcessById((int) pid); + } + catch (ArgumentException) + { + // Ignore it: process could be gone, or something else that we + // likely can't do anything about it. + } + + if (proc != null) + { + // Q: Why the check against the parent pid (below)? + // + // A: The idea is that a user could do something like this: + // + // $p = Start-Process pwsh -ArgumentList '-c Write-Host start $pid; sleep -seconds 30; Write-Host stop' -NoNewWindow -passThru + // + // Such a process *is* indeed _capable_ of wrecking the interactive prompt (all it has to do is to attempt to read input; and any output + // will be interleaved with your interactive session)... but MAYBE it won't. So the idea with letting such processes live is that perhaps + // the user did this on purpose, to do some sort of "background work" (even though it may not seem like the best way to do that); and we + // only want to kill *actually-orphaned* processes: processes whose parent is gone, so they should be gone, too. + // + // We only check the immediate children processes here for simplicity. However, an immediate child process may have children that accidentally + // derive the standard input (which technically is a wrong thing to do), so ideally we should check if the parent of a console-attached process + // is still alive -- the parent process id points to an alive process that was created earlier. + // We will wait for feedback to see if this check needs to be updated. + + if (GetParentPid(proc) != _myPid.Value) + { + procsToTerminate.Add(proc); + } + else + { + proc.Dispose(); + } + } + } + } + return procsToTerminate.Count; + } + + [DllImport("kernel32.dll")] + internal static extern ulong GetTickCount64(); + + private static int MillisLeftUntilDeadline(ulong deadline) + { + long diff = (long) (deadline - GetTickCount64()); + + if (diff < 0) + { + diff = 0; + } + else if (diff >= (long) Int32.MaxValue) + { + // Should not ever actually happen... + diff = DefaultGraceMillis; + } + + return (int) diff; + } + + private const int DefaultGraceMillis = 1000; + private const int MaxRounds = 2; + + // + // TerminateOrphanedConsoleApps + // + // This feature works around a bad interaction on Windows between: + // * a race condition between ctrl+c and console attachment, and + // * poor behavior when multiple processes want console input. + // + // This bad interaction is most likely to happen when the user has launched a process + // that is launching many, MANY more child processes (imagine a build system, for + // example): if the user types ctrl+c to cancel, all processes *currently attached* to + // the console will receive the ctrl+c signal (and presumably exit). However, there + // *may* have been some processes that had been created, but are not yet attached to + // the console--these grandchildren will have missed the ctrl+c signal (that's the + // race condition). If those grandchildren do not somehow figure out on their own that + // they should exit, the console enters a highly problematic state ("the borked + // state"): because pwsh's immediate child has exited, the shell will return to the + // prompt and wait for input. But those straggler granchildren are ALSO attached to + // the console... so when the user starts typing, who gets the input? + // + // It turns out that the console will just sort of randomly distribute pieces of input + // between all processes who want input--a straggler grandchild process might get a + // "key down" record, and then PSReadLine might get the corresponding "key up". This + // is obviously untenable; it makes the shell totally unusable. (The console team has + // been made aware, and there are several ideas of how to Do Better, but who knows + // when any of those will come to fruition.) + // + // To make matters worse: when returning to the prompt, PSReadLine disables ctrl+c + // signals (we prefer to handle those keys specially ourselves). So if you hit this + // situation with cmd.exe as your shell, you can just mash on ctrl+c for a while and + // kill all the stragglers manually; but if you have PSReadLine loaded, your shell is + // borked, and you are stuck. You CAN recover, IF you can track down and kill all the + // straggler processes manually. + // + // So when enabled, this feature does that for you: it kills all those straggler + // processes, right before we disable ctrl+c signals and wait for user input, ensuring + // that the user has a usable shell. + // + // Note that GUI processes do not attach to the console, so if you have launched + // notepad, for example, TerminateOrphanedConsoleApps will never even "see" it; they + // are immune from getting terminated. + // + // Q: But isn't terminating processes that we know nothing about kind of risky and + // extreme? + // + // A: Perhaps so... but consider the alternative: by definition, if you get into a + // situation where the TerminateOrphanedConsoleApps feature would actually kill + // anything, your shell will be Completely Broken. It's "them or us": allow the + // stragglers to live, but leave the user without their shell; or kill the + // stragglers and give the user their shell back. There is no middle ground. So + // when the TerminateOrphanedConsoleApps feature is enabled, that means the user + // has opted for "give me back my shell". + // + // Note that we do give stragglers a small grace period before terminating them, in + // case they are somehow just slow shutting down. But if you're wondering "should + // we make that grace period longer?", remember that another way to think of that + // period is "how long do I want the shell to potentially be unusable after + // displaying the prompt?" + // + // Q: What if the user *didn't* type ctrl+c? + // + // A: We don't care. When TerminateOrphanedConsoleApps is called, all we know is that + // the shell has displayed the prompt and believes it is time to wait for user + // input. Whether this situation came about because of a ctrl+c, or some other + // situation (for example, if the shell's immediate child crashed or was manually + // killed), if there are leftover straggler processes (console-attached + // grandchildren), the shell will be broken until they are gone, and thus we must + // take action (if the feature is enabled). + // + // Q: Should this really be baked into PSReadLine, or could we leave it to some other + // module to implement? (See: https://github.com/jazzdelightsme/ConsoleBouncer) + // + // A: We should have the option in PSReadLine. An external module can do something + // very *similar* to what we do here in PSReadLine, but not quite the same, and is + // strictly inferior. An external module would have to rely on receiving a ctrl+c + // signal, but "there was a ctrl+c signal" is NOT equivalent to "the shell is about + // to wait for input". For example, some child processes may depend on handling + // ctrl+c signals, *without* exiting (kd.exe / cdb.exe, for example). In such a + // case, control would not return to the shell, but an external module would have + // no way to know that (hence it is inferior). That could be worked around, but + // only clumsily--the user would have to have a way to tell the module "hey BTW + // please don't kill these ones, even though they will *look* like stragglers". + // + // And in fact, an external module solution may still be attractive to some users + // (and could safely be used with TerminateOrphanedConsoleApps enabled). Because + // the (external solution) ConsoleBouncer module reacts to ctrl+c signals, that + // makes it a bit more aggressive than what we do here: + // TerminateOrphanedConsoleApps only comes into play when control has returned to + // the shell, which might not be right away after the user types ctrl+c--there + // might be "Terminate batch job (Y/N)?" messages, etc. So if the user understands + // the limitations of the ConsoleBouncer module and has an environment where it + // would be suitable, they could still opt to use it to get much more responsive + // ctrl+c behavior. (A metaphor with a club: the PSReadLine built-in feature + // patiently waits for the host of a private party to leave before kicking the rest + // of the guests out; whereas the ConsoleBouncer, upon receipt of a ctrl+c signal, + // just clears the whole place out right away (which *might* not be the right thing + // to do, but you're paying them to be tough, not smart).) + // + private static void TerminateStragglers() + { + var procsToTerminate = new List(); + + // The theory for why more than one round might be needed is that the same race + // between process creation and console attachment that could cause lingering + // processes in the first place could cause us to need a second round of + // cleanup... but I've never actually seen more than one round be needed. Probably + // because in my specific scenario the process that was spawning processes got + // taken out with the original ctrl+c signal. + // + // If it takes more than a few rounds of cleanup, we may be in some kind of + // pathological situation, and we'll bow out. + int round = 0; + int killAttempts = 0; + + while (round++ < MaxRounds && + GatherStragglers(procsToTerminate) > 0) + { + // We'll give them up to GracePeriodMillis for them to exit on their + // own, in case they actually did receive the original ctrl+c, and are + // just a tad slow shutting down. + ulong deadline = GetTickCount64() + (ulong) DefaultGraceMillis; + + var notDeadYet = procsToTerminate.Where( + (p) => !p.WaitForExit(MillisLeftUntilDeadline(deadline))); + + foreach (var process in notDeadYet) + { + try + { + killAttempts++; + process.Kill(); + } + // Ignore problems; maybe it's gone already, maybe something else; + // whatever. + catch (InvalidOperationException) { } + catch (Win32Exception) { } + } + + foreach (var process in procsToTerminate) + { + process.Dispose(); + } + } // end retry loop + + // In forcible termination scenarios, if there was a child updating the terminal's + // progress state, it may be left stuck that way... we can clear that out. + // + // The preferred way to do that is with a VT sequence, but there's no way to know + // if the console we are attached to supports that sequence. If we are in Windows + // Terminal, we know we can use the VT sequence; else we'll fall back to the old + // (Win7-era?) COM API (which does the same thing). + uint consoleMode = GetConsoleOutputMode(); + if (ItLooksLikeWeAreInTerminal()) + { + // We can use the [semi-]standard OSC sequence: + // https://conemu.github.io/en/AnsiEscapeCodes.html#ConEmu_specific_OSC + if (0 != (consoleMode & (uint) ENABLE_VIRTUAL_TERMINAL_PROCESSING)) + { + // Use "bell" if we actually tried to whack anything. + string final = (killAttempts > 0) ? "\a" : "\x001b\\"; + Console.Write("\x001b]9;4;0;0" + final); + } + } + else + { + IntPtr hwnd = GetConsoleWindow(); + if (hwnd != IntPtr.Zero) + { + int ret = TaskbarProgress.SetProgressState(hwnd, TaskbarStates.NoProgress); + } + } + } }