Skip to content

Instantly share code, notes, and snippets.

@alx9r
Created December 20, 2024 18:37
Show Gist options
  • Save alx9r/6438a06c93b62c500f21e9a35e969ea8 to your computer and use it in GitHub Desktop.
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)
# 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'
}}
}
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