Last active
December 11, 2024 21:23
-
-
Save alx9r/8b785d0085b4caa0d8ba4f63ae6e6c5f to your computer and use it in GitHub Desktop.
Proof-of-concept of command combining New-Object and Set-Variable to achieve construction and assignment that can't be interrupted stopping runspace.
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.Management.Automation.Host; | |
using Automation = System.Management.Automation; | |
using System.Management.Automation.Runspaces; | |
using System; | |
namespace Assuage.PowerShell { | |
/* This command combines New-Object and Set-Variable to achieve construction and assignment in a manner that cannot be interrupted by the runspace stopping. This is required to ensure reliable cleanup using the assigned-to variable in clean{} because otherwise it is possible for construction, but not, assignment can succeed. | |
See also: | |
https://stackoverflow.com/questions/79265633/are-assignments-from-net-methods-atomic-when-stopping-powershell | |
https://github.com/PowerShell/PowerShell/issues/24658 | |
*/ | |
[Cmdlet(VerbsCommon.Set,"CleanupVariable")] | |
public class SetCleanupVariableCommand : Cmdlet { | |
// mirror the Name parameter of Set-Variable | |
// https://github.com/PowerShell/PowerShell/blob/294adb47f163124b0b0355fd995f7c8f1656c54c/src/Microsoft.PowerShell.Commands.Utility/commands/utility/Var.cs#L570-L571 | |
[Parameter(Mandatory = true,Position = 0)] | |
public string Name {get;set;} | |
// mirror the TypeName parameter of New-Object | |
// https://github.com/PowerShell/PowerShell/blob/294adb47f163124b0b0355fd995f7c8f1656c54c/src/Microsoft.PowerShell.Commands.Utility/commands/utility/New-Object.cs#L33-L35 | |
[Parameter(Mandatory = true, Position = 1)] | |
[ValidateTrustedData] | |
[Alias("NewObjectType")] | |
public string TypeName {get;set;} | |
// mirror the ArgumentList parameter of New-Object | |
// https://github.com/PowerShell/PowerShell/blob/294adb47f163124b0b0355fd995f7c8f1656c54c/src/Microsoft.PowerShell.Commands.Utility/commands/utility/New-Object.cs#L51-L53 | |
[Parameter(Mandatory = false, Position = 3)] | |
[ValidateTrustedData] | |
[Alias("Args")] | |
public object[] ArgumentList {get;set;} = new object[] {}; | |
// Gain access to the variables in the current runspace's session state. | |
// Assuming this command is invoked from PowerShell, those variables should be the ones at the call site. | |
PSVariableIntrinsics getPSVariableIntrinsics() { | |
using (var ps = Automation.PowerShell.Create(RunspaceMode.CurrentRunspace)) | |
{ | |
return ps | |
.AddScript(". {[CmdletBinding()]param()$PSCmdlet.SessionState.PSVariable}") | |
.Invoke<PSVariableIntrinsics>()[0]; | |
} | |
} | |
// Obtain a reference to the host at the call site. | |
PSHost getPSHost() { | |
using (var ps = Automation.PowerShell.Create(RunspaceMode.CurrentRunspace)) | |
{ | |
return ps | |
.AddScript("$Host") | |
.Invoke<PSHost>()[0]; | |
} | |
} | |
protected override void BeginProcessing() { | |
/* | |
Get a reference to the variables that will later be set. | |
Note that this is done in advance of the call to New-Object because getPSVariableIntrinsics() uses the current runspace which could be stopping. Doing after New-Object risks the possibility that getPSVariableIntrinsics() will fail because the runspace is stopping which would mean this isn't atomic. | |
*/ | |
PSVariableIntrinsics psvariable = getPSVariableIntrinsics(); | |
// Get a reference to the host used at the call site. | |
PSHost host = getPSHost(); | |
if (Stopping) { | |
ThrowTerminatingError( | |
new ErrorRecord ( | |
new PipelineStoppedException(), | |
null, | |
ErrorCategory.OperationStopped, | |
null | |
)); | |
return; | |
} | |
if (psvariable == null) { | |
// if somehow getPSVariableIntrinsics() silently failed, we end up here | |
ThrowTerminatingError( | |
new ErrorRecord ( | |
new CmdletInvocationException("Failed to retrieve PSVariable for the current session state."), | |
null, | |
ErrorCategory.ObjectNotFound, | |
null | |
)); | |
return; | |
} | |
// the variable to hold the newly-constructed object output by New-Object | |
object newObject; | |
Func<Runspace> createRunspace; | |
// Create separate runspace in which to invoke New-Object | |
// Using the default runspace here would defeat the purpose of this command being atomic since that runspace might be stopping. | |
// If a host was obtained from the call site, direct streams to that, otherwise use no host. | |
if (host == null) { | |
createRunspace = RunspaceFactory.CreateRunspace; | |
} | |
else { | |
createRunspace = ()=>RunspaceFactory.CreateRunspace(host); | |
} | |
// Prepare the environment in which to invoke New-Object. | |
using (var rs = createRunspace()) { | |
using (var ps = Automation.PowerShell.Create(rs)) { | |
rs.Open(); | |
// Redirect the error stream from the New-Object invocation be output from this Cmdlet. | |
ps.Streams.Error.DataAdding += (s,e) => WriteError((ErrorRecord)e.ItemAdded); | |
// create and invoke the call to New-Object | |
newObject = | |
ps.AddCommand("New-Object") | |
.AddParameter("TypeName",TypeName) | |
.AddParameter("ArgumentList",ArgumentList) | |
.Invoke()[0]; | |
if (ps.Streams.Error.Count > 0 && | |
newObject == null ) { | |
// Something went wrong with New-Object that we can't recover from, so don't assign the Null value to the variable. | |
// The error should already have surfaced by way of the earlier redirecting the error stream. | |
return; | |
} | |
}} | |
// Set the variable to the new object. | |
psvariable.Set(Name,newObject); | |
} | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment