Created
December 20, 2024 18:37
-
-
Save alx9r/6438a06c93b62c500f21e9a35e969ea8 to your computer and use it in GitHub Desktop.
Proof-of-Concept of Cmdlet to protect cleanup from being curtailed by console. (PowerShell/PowerShell#23786)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# add and import the Protect-Cleanup command | |
Add-Type ` | |
-Path .\protectCleanupCommand.cs ` | |
-PassThru | | |
% Assembly | | |
Import-Module | |
function DivideByZero { [CmdletBinding()]param() 1/0 } | |
Protect-Cleanup -ErrorStreamProtection ActionPreference { | |
. { | |
try { | |
# $ErrorActionPreference here is unaltered | |
Write-Verbose "try{} ErrorActionPreference: $ErrorActionPreference" -Verbose | |
Write-Host 'Press Ctrl+C' | |
Start-Sleep 1 | |
Write-Host 'Ctrl+C not pressed in time' | |
} | |
finally { | |
<# | |
If Ctrl+C has been pressed and ErrorStreamProtection has the ActionPreference flag set, | |
the $ErrorActionPreference should have been set to 'SilentlyContinue' by Ctrl+C. | |
#> | |
Write-Verbose "finally{} ErrorActionPreference: $ErrorActionPreference" -Verbose | |
# Trigger invocation of the downstream script block | |
'output' | |
} | |
} | | |
. { | |
process { | |
# These following lines don't cause problems and are output regardless after Ctrl+C. | |
Write-Host 'Write-Host' | |
Write-Information 'Write-Information' | |
Write-Verbose 'Write-Verbose' -Verbose | |
Write-Warning 'Write-Warning' | |
<# | |
This error record is | |
- redirected to the success stream | |
- output on Protect-Cleanup's success stream before Ctrl+C | |
- blocked from being output from Protect-Cleanup's success stream after Ctrl+C | |
#> | |
DivideByZero | |
Write-Error 'Write-Error' | |
<# | |
The 'Write-Output' and 'output' objects are | |
- output on Protect-Cleanup's success stream before Ctrl+C | |
- blocked from being output from Protect-Cleanup's success stream after Ctrl+C | |
#> | |
Write-Output 'Write-Output' | |
'output' | |
# Reaching this after Ctrl+C indicates cleanup was protected from curtailment by the console. | |
Write-Host 'execution completed' | |
}} | |
} |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
using System.Management.Automation; | |
using System.Runtime.InteropServices; | |
using System.Threading; | |
using System; | |
[Flags] | |
public enum CleanupProtection { | |
None = 0b_0000_0000, | |
ActionPreference = 0b_0000_0001, | |
StreamRedirection = 0b_0000_0010 | |
} | |
// Proof-of-Concept command to protect cleanup from curtailment by the console. | |
// See also https://github.com/PowerShell/PowerShell/issues/23786 | |
[Cmdlet(VerbsSecurity.Protect,"Cleanup")] | |
public class ProtectCleanupCommand : Cmdlet { | |
[Parameter(Position=1)] | |
// the user script block | |
public ScriptBlock ScriptBlock {get; set;} | |
[Parameter()] | |
/* | |
The protection measures applied to the error stream: | |
- ActionPreference : Sets ErrorActionPreference to 'SilentlyContinue' in ScriptBlock's after Ctrl+C. | |
- StreamRedirection : Redirects all errors from ScriptBlock the information stream. | |
Use of ActionPreference is always recommended. Use of StreamRedirection is recommended only if ActionPreference is no sufficient to prevent cleanup curtailment. StreamRedirection has the drawback that error formatting performed by the console is lost for errors including those encountered when Ctrl+C has not been pressed. | |
*/ | |
public CleanupProtection ErrorStreamProtection {get; set;} | |
[Parameter()] | |
/* | |
The protection measures applied to the success stream: | |
- StreamRedirection : Redirects success stream output from ScriptBlock after Ctrl+C to the information stream. | |
*/ | |
public CleanupProtection SuccessStreamProtection {get; set;} | |
// pipeline machinery for invoking the user script block | |
SteppablePipeline steppable; | |
/* | |
The following two class/property pairs rely on the following from the language specification for thread safety: | |
>Reads and writes of the following data types shall be atomic: bool, char, byte, sbyte, short, ushort, uint, int, float, and reference types. In addition, reads and writes of enum types with an underlying type in the previous list shall also be atomic. | |
*/ | |
// Object to capture whether Ctrl+C has been pressed. | |
CtrlCFlag ctrlCFlag; | |
class CtrlCFlag { | |
bool flag = false; | |
public bool IsSet { get { return flag; } } | |
public void Set() { | |
flag = true; | |
} | |
} | |
// Object to capture the $ErrorActionPreference seen by ScriptBlock. | |
// The value of the PSVariable is altered when Ctrl+C is detected. | |
ActionPreferenceVariableReference errorActionPreference; | |
class ActionPreferenceVariableReference { | |
public PSVariable ActionPreferenceVariable = null; | |
public void Set(ActionPreference value) { | |
if (ActionPreferenceVariable == null) { | |
return; | |
} | |
ActionPreferenceVariable.Value = value; | |
} | |
} | |
// registration of the Ctrl+C signal | |
PosixSignalRegistration registration; | |
public ProtectCleanupCommand() { | |
// initialize capture properties | |
ctrlCFlag = new CtrlCFlag(); | |
errorActionPreference = new ActionPreferenceVariableReference(); | |
// default parameter values | |
ErrorStreamProtection = CleanupProtection.ActionPreference; | |
SuccessStreamProtection = CleanupProtection.StreamRedirection; | |
} | |
protected override void BeginProcessing() { | |
/* | |
PosixSignal.SIGINT doesn't match exactly the event the console state hinges on. | |
That event seems to be resolved starting at | |
https://github.com/PowerShell/PowerShell/blob/617dbda8f47cf06d115947f0db282e1994294604/src/Microsoft.PowerShell.ConsoleHost/host/msh/ConsoleHost.cs#L1231-L1236. | |
Mimicking that event reliably is quite a bit more complicated. | |
*/ | |
// register the SIGINT handler | |
registration = | |
PosixSignalRegistration.Create( | |
PosixSignal.SIGINT, | |
ctx=>{ | |
errorActionPreference.Set(ActionPreference.SilentlyContinue); | |
ctrlCFlag.Set(); // set the Ctrl+C flag on SIGINT | |
} | |
); | |
// create the steppable pipeline to invoke the user script block | |
steppable = | |
ScriptBlock | |
.Create(@" | |
# the requisite single pipeline script block | |
# these params are hidden from $__sb by the steppable pipeline machinery | |
param($sb,$flag,$eap,$errorMode,$successMode,$cmdlet) . { | |
# the utility script block | |
# These params are unavoidably visible to $__sb because it is dot-sourced. | |
param($__sb,$__flag,$__eap,$__errorMode,$__successMode,$__cmdlet) | |
if ($__errorMode.HasFlag([CleanupProtection]'ActionPreference')) { | |
<# | |
Capture the reference to $ErrorActionPreference. The value of this reference | |
will be seen by the commands in $sb. | |
#> | |
$__eap.ActionPreferenceVariable = Get-Variable ErrorActionPreference | |
} | |
. { | |
# invoke the user script block | |
if ($__errorMode.HasFlag([CleanupProtection]'StreamRedirection')) { | |
# Wrap the success stream so it can be distinguished from the error stream. | |
# Merge the error stream with the success stream. | |
. { . $__sb | . { process { [pscustomobject]@{ Success = $_ } } } } 2>&1 | |
} | |
else { | |
. $__sb | |
} | |
} | | |
& { | |
param( | |
[Parameter(ValueFromPipeline)] | |
# Objects from the success stream output by the user block. | |
# The success stream has the error stream merged into it when | |
# error stream StreamRedirection is selected. | |
$InputObject | |
) | |
begin { | |
# count used in verbose output | |
$count = 0 | |
# script block for verbose output | |
$verbose = { | |
$count+=1 | |
$msg = 'Stopped '+$count+' objects from reaching the console after Ctrl+C' | |
$__cmdlet.WriteVerbose($msg) | |
} | |
} | |
process { | |
if ($__errorMode.HasFlag([CleanupProtection]'StreamRedirection') -and | |
($InputObject -is [System.Management.Automation.ErrorRecord])) { | |
# Any error records encountered here are directed to the information | |
# stream. They will lose their formatting, but at least they will be | |
# seen. | |
& { | |
$InformationPreference = 'Continue' | |
$InputObject | Out-String | Write-Information | |
} | |
. $verbose | |
return | |
} | |
& { | |
if ($__errorMode.HasFlag([CleanupProtection]'StreamRedirection')) { | |
# Unwrap the success stream objects that were previously wrapped. | |
,$InputObject.Success | |
} | |
else { | |
$InputObject | |
} | |
} | | |
. { | |
process { | |
if ($__flag.IsSet -and | |
$__successMode.HasFlag([CleanupProtection]'StreamRedirection')) { | |
# Any success stream objects encountered here are directed to the information | |
# stream. Sending them to the success stream will curtail execution. | |
& { | |
$InformationPreference = 'Continue' | |
$_ | Out-String | Write-Information | |
} | |
. $verbose | |
} | |
else { | |
,$_ | |
} | |
} | |
} | |
} | |
end { | |
# this block is not reliably reached after Ctrl+C | |
} | |
} | |
} $sb $flag $eap $errorMode $successMode $cmdlet") | |
.GetSteppablePipeline( | |
CommandOrigin.Internal, | |
new object [] { | |
// pass the arguments required by the script block | |
ScriptBlock , // the user script block | |
ctrlCFlag , // the flag that indicates whether Ctrl+C has been pressed | |
errorActionPreference , // a reference to $ErrorActionPreference that is altered when Ctrl+C is pressed | |
ErrorStreamProtection , // the flags indicating which measures to apply | |
SuccessStreamProtection, // " | |
this // a reference for using WriteVerbose() | |
} | |
); | |
// start invoking user script block | |
steppable.Begin(this); | |
} | |
protected override void EndProcessing() | |
{ | |
// by the time this is reached after Ctrl+C tearing down has already disabled some facilities | |
// complete invoking the user script block | |
steppable.End(); | |
} | |
public void Dispose() | |
{ | |
steppable.Dispose(); | |
registration.Dispose(); | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment