Skip to content

Instantly share code, notes, and snippets.

@alx9r
Last active December 11, 2024 21:23
Show Gist options
  • Save alx9r/8b785d0085b4caa0d8ba4f63ae6e6c5f to your computer and use it in GitHub Desktop.
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.
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