Procgov returns the exit code of the launched monitored process (#66)

procgov/ProcessModule.cs

@@ -40,7 +40,7 @@ Win32Job OpenOrCreateJob()
                     CheckWin32Result(PInvoke.IsProcessInJob(targetProcessHandle, jobHandle, out var jobNameMatches)) && jobNameMatches)
                 {
                     SetProcessEnvironmentVariables(pid, session.AdditionalEnvironmentVars);
-                    return new Win32Job(jobHandle, jobName, session.ClockTimeLimitInMilliseconds);
+                    return new Win32Job(jobHandle, jobName, null, session.ClockTimeLimitInMilliseconds);
                 }
                 else
                 {
@@ -170,7 +170,7 @@ void AssignProcessToExistingJobObject(int processId, Win32Job job, bool checkIfA
                     SetProcessEnvironmentVariables(jobProcessId, session.AdditionalEnvironmentVars);
                     AccountPrivilegeModule.EnablePrivileges((uint)jobProcessId, jobProcessHandle, session.Privileges, TraceEventType.Error);
 
-                    job = new Win32Job(jobHandle, jobName, session.ClockTimeLimitInMilliseconds);
+                    job = new Win32Job(jobHandle, jobName, null, session.ClockTimeLimitInMilliseconds);
 
                     logger.TraceEvent(TraceEventType.Verbose, 0,
                         $"Procgov job already exists ('{jobName}') for process {jobProcessId} and we will use it for other processes.");
@@ -209,8 +209,7 @@ public static unsafe Win32Job StartProcessAndAssignToJobObject(
     {
         var pi = new PROCESS_INFORMATION();
         var si = new STARTUPINFOW();
-        var processCreationFlags = PROCESS_CREATION_FLAGS.CREATE_UNICODE_ENVIRONMENT |
-                                   PROCESS_CREATION_FLAGS.CREATE_SUSPENDED;
+        var processCreationFlags = PROCESS_CREATION_FLAGS.CREATE_UNICODE_ENVIRONMENT | PROCESS_CREATION_FLAGS.CREATE_SUSPENDED;
         if (session.SpawnNewConsoleWindow)
         {
             processCreationFlags |= PROCESS_CREATION_FLAGS.CREATE_NEW_CONSOLE;
@@ -230,7 +229,7 @@ public static unsafe Win32Job StartProcessAndAssignToJobObject(
             }
         }
 
-        using var processHandle = new SafeFileHandle(pi.hProcess, true);
+        var processHandle = new SafeFileHandle(pi.hProcess, true);
         try
         {
             var job = Win32JobModule.CreateJobObjectAndAssignProcess(processHandle, jobName, session.PropagateOnChildProcesses,
@@ -280,7 +279,7 @@ public static unsafe Win32Job StartProcessUnderDebuggerAndAssignToJobObject(
             }
         }
 
-        using var processHandle = new SafeFileHandle(pi.hProcess, true);
+        var processHandle = new SafeFileHandle(pi.hProcess, true);
         try
         {
             CheckWin32Result(PInvoke.DebugSetProcessKillOnExit(false));

procgov/Win32Job.cs

@@ -1,33 +1,24 @@
-using System;
-using System.Diagnostics;
-using System.Runtime.InteropServices;
+using System.Runtime.InteropServices;
 
 namespace ProcessGovernor;
 
-public sealed class Win32Job : IDisposable
+public record Win32Job(SafeHandle JobHandle, string JobName, SafeHandle? ProcessHandle = null, long ClockTimeLimitInMilliseconds = 0L) : IDisposable
 {
-    private readonly SafeHandle hJob;
-    private readonly string jobName;
-    private readonly Stopwatch? stopWatch;
-    private readonly long clockTimeLimitInMilliseconds;
+    private readonly DateTime startTimeUtc = DateTime.UtcNow;
 
-    public Win32Job(SafeHandle hJob, string jobName, long clockTimeLimitInMilliseconds = 0L)
-    {
-        this.hJob = hJob;
-        this.jobName = jobName;
-
-        this.clockTimeLimitInMilliseconds = clockTimeLimitInMilliseconds;
-        stopWatch = clockTimeLimitInMilliseconds > 0 ? Stopwatch.StartNew() : null;
-    }
-
-    public SafeHandle JobHandle => hJob;
-
-    public string JobName => jobName;
+    public bool IsTimedOut => ClockTimeLimitInMilliseconds > 0 
+        && DateTime.UtcNow.Subtract(startTimeUtc).TotalMilliseconds >= ClockTimeLimitInMilliseconds;
 
-    public bool IsTimedOut => stopWatch != null && stopWatch.ElapsedMilliseconds > clockTimeLimitInMilliseconds;
+    // When we are monitoring only a specific process, we will wait for its termination. Otherwise,
+    // we will wait for the job object to be signaled.
+    public SafeHandle WaitHandle => ProcessHandle ?? JobHandle;
 
     public void Dispose()
     {
-        hJob.Dispose();
+        JobHandle.Dispose();
+        if (ProcessHandle is { } h && !h.IsInvalid)
+        {
+            h.Dispose();
+        }
     }
 }

procgov/Win32JobModule.cs

@@ -19,7 +19,12 @@ internal static class Win32JobModule
     public static unsafe Win32Job CreateJobObjectAndAssignProcess(SafeHandle processHandle, string jobName, bool propagateOnChildProcesses,
         long clockTimeLimitInMilliseconds)
     {
-        var job = CreateJobObject(jobName, clockTimeLimitInMilliseconds);
+        var securityAttributes = new SECURITY_ATTRIBUTES();
+        securityAttributes.nLength = (uint)Marshal.SizeOf(securityAttributes);
+
+        var hJob = CheckWin32Result(PInvoke.CreateJobObject(securityAttributes, $"Local\\{jobName}"));
+        var job = new Win32Job(hJob, jobName, processHandle, clockTimeLimitInMilliseconds);
+
         AssignProcess(job, processHandle, propagateOnChildProcesses);
 
         return job;
@@ -31,7 +36,7 @@ public static unsafe Win32Job CreateJobObject(string jobName, long clockTimeLimi
         securityAttributes.nLength = (uint)Marshal.SizeOf(securityAttributes);
 
         var hJob = CheckWin32Result(PInvoke.CreateJobObject(securityAttributes, $"Local\\{jobName}"));
-        var job = new Win32Job(hJob, jobName, clockTimeLimitInMilliseconds);
+        var job = new Win32Job(hJob, jobName, null, clockTimeLimitInMilliseconds);
 
         return job;
     }
@@ -266,11 +271,16 @@ public static unsafe int WaitForTheJobToComplete(Win32Job job, CancellationToken
     {
         while (!ct.IsCancellationRequested)
         {
-            switch (PInvoke.WaitForSingleObject(job.JobHandle, 200 /* ms */))
+            switch (PInvoke.WaitForSingleObject(job.WaitHandle, 200 /* ms */))
             {
                 case WAIT_EVENT.WAIT_OBJECT_0:
-                    logger.TraceInformation("End of job time limit passed - terminating.");
-                    return 2;
+                    logger.TraceInformation("Job or process got signaled.");
+                    if (job.ProcessHandle is { } h && !h.IsInvalid)
+                    {
+                        PInvoke.GetExitCodeProcess(h, out var exitCode);
+                        return (int)exitCode;
+                    }
+                    return 0;
                 case WAIT_EVENT.WAIT_FAILED:
                     throw new Win32Exception();
                 default:
@@ -289,7 +299,7 @@ public static unsafe int WaitForTheJobToComplete(Win32Job job, CancellationToken
                     {
                         logger.TraceInformation("Clock time limit passed - terminating.");
                         PInvoke.TerminateJobObject(job.JobHandle, 1);
-                        return 1;
+                        return 0;
                     }
                     else
                     {
@@ -298,6 +308,6 @@ public static unsafe int WaitForTheJobToComplete(Win32Job job, CancellationToken
             }
         }
 
-        return 3;
+        return 0;
     }
 }