Skip to content

Commit b9c70b1

Browse files
jborean93jborean93
authored
Add cancellation for -Wait (#14)
Adds support for cancelling Start-ProcessEx and Start-ProcessWith when using the -Wait parameter.
1 parent 363a262 commit b9c70b1

File tree

7 files changed

+145
-17
lines changed

7 files changed

+145
-17
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
* Added support for `ProcessIntString` parameters accepting `uint` values as returned by WMI/CIM calls
66
* Added `-ArgumentEscaping` to `Start-ProcessWith` and `Start-ProcessEx`
77
* Added `Credential` attribute transformer to transform username to PSCredential prompt on `-Credential` parameters
8+
* Add support for cancelling `Start-ProcessEx` and `Start-ProcessWith` with `-Wait`
89

910
## v0.4.0 - 2024-11-07
1011

‎src/ProcessEx/Commands/ProcessEx.cs‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,10 @@
22
using ProcessEx.Security;
33
using System;
44
using System.Collections;
5-
using System.Collections.Generic;
65
using System.Linq;
76
using System.Management.Automation;
87
using System.Runtime.InteropServices;
8+
using System.Threading;
99

1010
namespace ProcessEx.Commands;
1111

@@ -60,8 +60,10 @@ protected override void ProcessRecord()
6060
DefaultParameterSetName = "FilePath"
6161
)]
6262
[OutputType(null, typeof(ProcessInfo))]
63-
public class StartProcessEx : PSCmdlet
63+
public class StartProcessEx : PSCmdlet, IDisposable
6464
{
65+
private CancellationTokenSource? _cancelSource;
66+
6567
[Parameter(
6668
Mandatory = true,
6769
Position = 0,
@@ -203,8 +205,11 @@ protected override void ProcessRecord()
203205
WriteVerbose($"Process created with PID {info.ProcessId} and TID {info.ThreadId}");
204206
if (Wait)
205207
{
206-
WriteVerbose("Resuming process and waiting for it to complete");
207-
ProcessRunner.ResumeAndWait(info);
208+
using (_cancelSource = new())
209+
{
210+
WriteVerbose("Resuming process and waiting for it to complete");
211+
ProcessRunner.ResumeAndWait(info, _cancelSource.Token);
212+
}
208213
}
209214
else if (!suspended)
210215
{
@@ -215,4 +220,14 @@ protected override void ProcessRecord()
215220
if (PassThru)
216221
WriteObject(info);
217222
}
223+
224+
protected override void StopProcessing()
225+
{
226+
_cancelSource?.Cancel();
227+
}
228+
229+
public void Dispose()
230+
{
231+
_cancelSource?.Dispose();
232+
}
218233
}

‎src/ProcessEx/Commands/ProcessWith.cs‎

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
using ProcessEx.Native;
22
using System;
33
using System.Collections;
4-
using System.Collections.Generic;
54
using System.Linq;
65
using System.Management.Automation;
76
using System.Runtime.InteropServices;
87
using System.Security;
8+
using System.Threading;
99

1010
namespace ProcessEx.Commands
1111
{
@@ -14,8 +14,10 @@ namespace ProcessEx.Commands
1414
DefaultParameterSetName = "FilePathCredential"
1515
)]
1616
[OutputType(null, typeof(ProcessInfo))]
17-
public class StartProcessWith : PSCmdlet
17+
public class StartProcessWith : PSCmdlet, IDisposable
1818
{
19+
private CancellationTokenSource? _cancelSource;
20+
1921
[Parameter(
2022
Mandatory = true,
2123
Position = 0,
@@ -201,8 +203,11 @@ protected override void ProcessRecord()
201203
WriteVerbose($"Process created with PID {info.ProcessId} and TID {info.ThreadId}");
202204
if (Wait)
203205
{
204-
WriteVerbose("Resuming process and waiting for it to complete");
205-
ProcessRunner.ResumeAndWait(info);
206+
using (_cancelSource = new())
207+
{
208+
WriteVerbose("Resuming process and waiting for it to complete");
209+
ProcessRunner.ResumeAndWait(info, _cancelSource.Token);
210+
}
206211
}
207212
else if (!suspended)
208213
{
@@ -213,5 +218,17 @@ protected override void ProcessRecord()
213218
if (PassThru)
214219
WriteObject(info);
215220
}
221+
222+
protected override void StopProcessing()
223+
{
224+
_cancelSource?.Cancel();
225+
}
226+
227+
228+
public void Dispose()
229+
{
230+
_cancelSource?.Dispose();
231+
}
232+
216233
}
217234
}

‎src/ProcessEx/Native/Kernel32.cs‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -596,6 +596,28 @@ public static SafeNativeHandle OpenProcess(int processId, ProcessAccessRights ac
596596
return handle;
597597
}
598598

599+
[DllImport("Kernel32.dll", EntryPoint = "PostQueuedCompletionStatus", SetLastError = true)]
600+
private static extern bool NativePostQueuedCompletionStatus(
601+
SafeHandle CompletionPort,
602+
uint dwNumberOfBytesTransferred,
603+
nint dwCompletionKey,
604+
nint lpOverlapped);
605+
606+
public static void PostQueuedCompletionStatus(
607+
SafeHandle completionPort,
608+
uint data,
609+
nint completionKey)
610+
{
611+
if (!NativePostQueuedCompletionStatus(
612+
completionPort,
613+
data,
614+
completionKey,
615+
IntPtr.Zero))
616+
{
617+
throw new NativeException("PostQueuedCompletionStatus");
618+
}
619+
}
620+
599621
[DllImport("Kernel32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
600622
private static extern bool QueryFullProcessImageNameW(
601623
SafeHandle hProcess,

‎src/ProcessEx/Process.cs‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
using System.Security.AccessControl;
1414
using System.Security.Principal;
1515
using System.Text;
16+
using System.Threading;
1617

1718
namespace ProcessEx
1819
{
@@ -821,7 +822,9 @@ private static ProcessInfo CreateProcessInternal(string? applicationName, string
821822
creationFlags, lpEnvironment, currentDirectory, si, ext ?? new Dictionary<string, object?>());
822823
}
823824

824-
internal static void ResumeAndWait(ProcessInfo processInfo)
825+
internal static void ResumeAndWait(
826+
ProcessInfo processInfo,
827+
CancellationToken cancellationToken)
825828
{
826829
// Thanks to Raymond for these details https://devblogs.microsoft.com/oldnewthing/20130405-00/?p=4743
827830
int compPortSize = Marshal.SizeOf(typeof(Helpers.JOBOBJECT_ASSOCIATE_COMPLETION_PORT));
@@ -845,7 +848,12 @@ internal static void ResumeAndWait(ProcessInfo processInfo)
845848

846849
// Resume the thread and wait until it has exited.
847850
Kernel32.ResumeThread(processInfo.Thread);
848-
Kernel32.WaitForSingleObject(processInfo.Process, 0xFFFFFFFF);
851+
852+
const uint canceledCode = 0xFFFFFFFF;
853+
using var cancelRegistration = cancellationToken.Register(() =>
854+
{
855+
Kernel32.PostQueuedCompletionStatus(ioPort, canceledCode, IntPtr.Zero);
856+
});
849857

850858
// Continue to poll the job until it receives JOB_OBJECT_MSG_ACTIVE_PROCESS_ZERO (4) that indicates
851859
// all other processes in that job have finished.
@@ -854,7 +862,7 @@ internal static void ResumeAndWait(ProcessInfo processInfo)
854862
{
855863
Kernel32.GetQueuedCompletionStatus(ioPort, 0xFFFFFFFF, out completionCode, out var completionKey,
856864
out var overlapped);
857-
} while (completionCode != 4);
865+
} while (completionCode != 4 && completionCode != canceledCode);
858866
}
859867

860868
internal static Dictionary<string, string> ConvertEnvironmentBlock(SafeHandle block)

‎tests/ProcessEx.Tests.ps1‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1239,6 +1239,71 @@ Describe "Start-ProcessEx" {
12391239
}
12401240
}
12411241

1242+
It "Cancels process -Wait" {
1243+
$stdout = [System.IO.Pipes.AnonymousPipeServerStream]::new("In", "Inheritable")
1244+
$stdin = [System.IO.Pipes.AnonymousPipeServerStream]::new("Out", "Inheritable")
1245+
try {
1246+
$si = New-StartupInfo -StandardOutput $stdout.ClientSafePipeHandle -StandardInput $stdin.ClientSafePipeHandle
1247+
$procParams = @{
1248+
FilePath = 'powershell.exe'
1249+
ArgumentList = '-Command', "-"
1250+
StartupInfo = $si
1251+
Wait = $true
1252+
PassThru = $true
1253+
}
1254+
1255+
$projectPath = Split-Path -Path $PSScriptRoot -Parent
1256+
$ps = [PowerShell]::Create()
1257+
[void]$ps.AddCommand("Import-Module").AddParameter("Name", "$projectPath\output\ProcessEx").AddStatement()
1258+
$task = $ps.AddCommand('Start-ProcessEx').AddParameters($procParams).BeginInvoke()
1259+
try {
1260+
$stdinWriter = [IO.StreamWriter]::new($stdin)
1261+
$stdinWriter.AutoFlush = $true
1262+
$stdoutReader = [IO.StreamReader]::new($stdout)
1263+
1264+
$stdinWriter.WriteLine('$pid')
1265+
[int]$childPid1 = $stdoutReader.ReadLine()
1266+
1267+
# Wait until we know the process has started (read the input) before disposing the client copies
1268+
$stdin.DisposeLocalCopyOfClientHandle()
1269+
$stdout.DisposeLocalCopyOfClientHandle()
1270+
1271+
$stdinWriter.WriteLine('(Start-Process powershell.exe -PassThru).Id')
1272+
[int]$childPid2 = $stdoutReader.ReadLine()
1273+
$stdoutReader.Dispose()
1274+
1275+
try {
1276+
$spawnedProcess = Get-Process -Id $childPid1 -ErrorAction SilentlyContinue
1277+
$spawnedProcess | Should -Not -Be $null
1278+
$task.IsCompleted | Should -Be $false
1279+
1280+
$stdinWriter.WriteLine('exit')
1281+
$stdinWriter.Dispose()
1282+
$spawnedProcess.WaitForExit(5000)
1283+
$task.IsCompleted | Should -Be $false
1284+
1285+
Wait-Process -Id $childPid1 -Timeout 5 -ErrorAction SilentlyContinue
1286+
$task.IsCompleted | Should -Be $false
1287+
1288+
$null = $ps.BeginStop($null, $null)
1289+
$task.AsyncWaitHandle.WaitOne(5000)
1290+
$task.IsCompleted | Should -Be $true
1291+
}
1292+
finally {
1293+
Stop-Process -Id $childPid1 -Force -ErrorAction SilentlyContinue
1294+
Stop-Process -Id $childPid2 -Force -ErrorAction SilentlyContinue
1295+
}
1296+
}
1297+
finally {
1298+
$ps.EndInvoke($task)
1299+
}
1300+
}
1301+
finally {
1302+
$stdout.Dispose()
1303+
$stdin.Dispose()
1304+
}
1305+
}
1306+
12421307
It "Fails with suspend and -Wait" {
12431308
$err = $null
12441309
Start-ProcessEx pwsh -CreationFlags Suspended -Wait -ErrorAction SilentlyContinue -ErrorVariable err

‎tests/ProcessWith.Tests.ps1‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,11 @@ Describe "Start-ProcessWith" {
2323
([IO.Path]::GettempPath()), 'C:\Windows\TEMP' | ForEach-Object {
2424
$acl = Get-Acl -LiteralPath $_
2525
$acl.AddAccessRule($acl.AccessRuleFactory($user,
26-
[System.Security.AccessControl.FileSystemRights]::FullControl,
27-
$false,
28-
[System.Security.AccessControl.InheritanceFlags]"ContainerInherit, ObjectInherit",
29-
[System.Security.AccessControl.PropagationFlags]::None,
30-
[System.Security.AccessControl.AccessControlType]::Allow))
26+
[System.Security.AccessControl.FileSystemRights]::FullControl,
27+
$false,
28+
[System.Security.AccessControl.InheritanceFlags]"ContainerInherit, ObjectInherit",
29+
[System.Security.AccessControl.PropagationFlags]::None,
30+
[System.Security.AccessControl.AccessControlType]::Allow))
3131
Set-Acl -LiteralPath $_ -AclObject $acl
3232
}
3333

@@ -187,7 +187,7 @@ Describe "Start-ProcessWith" {
187187
[System.Security.Principal.WellKnownSidType]::LocalSystemSid, $null)
188188
$systemName = $systemSid.Translate([System.Security.Principal.NTAccount]).Value
189189

190-
$actual = Get-Process -IncludeUserName | Where-Object UserName -eq $systemName | ForEach-Object {
190+
$actual = Get-Process -IncludeUserName | Where-Object UserName -EQ $systemName | ForEach-Object {
191191
$systemToken = Get-ProcessToken -Process $_ -Access Duplicate, Impersonate, Query -ErrorAction SilentlyContinue
192192
if (-not $systemToken) { return }
193193

0 commit comments

Comments
 (0)