Use PROC_THREAD_ATTRIBUTE_JOB_LIST to atomically assign a process to …

…a job on creation
asmichi · Sep 8, 2021 · 5fbd271 · 5fbd271
1 parent f14e0a0
commit 5fbd271
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 22 deletions.
diff --git a/src/ChildProcess/Interop/Windows/Kernel32.CreateProcess.cs b/src/ChildProcess/Interop/Windows/Kernel32.CreateProcess.cs
@@ -12,7 +12,6 @@ namespace Asmichi.Interop.Windows
     internal static partial class Kernel32
     {
         // Process Creation Flags
-        public const int CREATE_SUSPENDED = 0x00000004;
         public const int CREATE_UNICODE_ENVIRONMENT = 0x00000400;
         public const int EXTENDED_STARTUPINFO_PRESENT = 0x00080000;
         public const int CREATE_NO_WINDOW = 0x08000000;

diff --git a/src/ChildProcess/Interop/Windows/Kernel32.JobObject.cs b/src/ChildProcess/Interop/Windows/Kernel32.JobObject.cs
@@ -21,11 +21,6 @@ public static extern SafeJobObjectHandle CreateJobObject(
              [In] IntPtr lpJobAttributes,
              [In] char[]? lpName);
 
-        [DllImport(DllName, SetLastError = true)]
-        public static extern bool AssignProcessToJobObject(
-            [In] SafeJobObjectHandle hJob,
-            [In] SafeProcessHandle hProcess);
-
         [DllImport(DllName, SetLastError = true)]
         public static extern unsafe bool SetInformationJobObject(
             [In] SafeJobObjectHandle hJob,

diff --git a/src/ChildProcess/Interop/Windows/Kernel32.ProcThreadAttributeList.cs b/src/ChildProcess/Interop/Windows/Kernel32.ProcThreadAttributeList.cs
@@ -11,6 +11,7 @@ namespace Asmichi.Interop.Windows
     internal static partial class Kernel32
     {
         public static readonly IntPtr PROC_THREAD_ATTRIBUTE_HANDLE_LIST = new IntPtr(0x20002);
+        public static readonly IntPtr PROC_THREAD_ATTRIBUTE_JOB_LIST = new IntPtr(0x2000d);
         public static readonly IntPtr PROC_THREAD_ATTRIBUTE_PSEUDOCONSOLE = new IntPtr(0x20016);
 
         [DllImport(DllName, SetLastError = true)]

diff --git a/src/ChildProcess/Interop/Windows/ProcThreadAttributeList.cs b/src/ChildProcess/Interop/Windows/ProcThreadAttributeList.cs
@@ -39,6 +39,21 @@ public unsafe void UpdateHandleList(IntPtr* handles, int count)
             }
         }
 
+        public unsafe void UpdateJobList(IntPtr* handles, int count)
+        {
+            if (!Kernel32.UpdateProcThreadAttribute(
+                _unmanaged,
+                0,
+                Kernel32.PROC_THREAD_ATTRIBUTE_JOB_LIST,
+                handles,
+                sizeof(IntPtr) * count,
+                IntPtr.Zero,
+                IntPtr.Zero))
+            {
+                throw new Win32Exception();
+            }
+        }
+
         public unsafe void UpdatePseudoConsole(IntPtr hPC)
         {
             if (!Kernel32.UpdateProcThreadAttribute(

diff --git a/src/ChildProcess/ProcessManagement/WindowsChildProcessStateHelper.cs b/src/ChildProcess/ProcessManagement/WindowsChildProcessStateHelper.cs
@@ -63,22 +63,23 @@ public unsafe IChildProcessStateHolder SpawnProcess(
                 var childStdOut = stdOut != null ? inheritableHandleStore.Add(stdOut) : null;
                 var childStdErr = stdErr != null ? inheritableHandleStore.Add(stdErr) : null;
 
+                IntPtr jobObjectHandles = jobObjectHandle.DangerousGetHandle();
+
                 Span<IntPtr> inheritableHandles = stackalloc IntPtr[inheritableHandleStore.Count];
                 inheritableHandleStore.DangerousGetHandles(inheritableHandles);
                 fixed (IntPtr* pInheritableHandles = inheritableHandles)
                 {
-                    using var attr = new ProcThreadAttributeList(2);
+                    using var attr = new ProcThreadAttributeList(3);
                     if (pseudoConsole is not null)
                     {
                         attr.UpdatePseudoConsole(pseudoConsole.Handle.DangerousGetHandle());
                     }
                     attr.UpdateHandleList(pInheritableHandles, inheritableHandles.Length);
+                    attr.UpdateJobList(&jobObjectHandles, 1);
 
-                    // Create the process suspended so that it will not create a grandchild process before we assign it to the job object.
                     const int CreationFlags =
                         Kernel32.CREATE_UNICODE_ENVIRONMENT
-                        | Kernel32.EXTENDED_STARTUPINFO_PRESENT
-                        | Kernel32.CREATE_SUSPENDED;
+                        | Kernel32.EXTENDED_STARTUPINFO_PRESENT;
 
                     int processId;
                     (processId, processHandle, threadHandle) = InvokeCreateProcess(
@@ -91,18 +92,6 @@ public unsafe IChildProcessStateHolder SpawnProcess(
                         childStdErr,
                         attr);
 
-                    if (!Kernel32.AssignProcessToJobObject(jobObjectHandle, processHandle))
-                    {
-                        // Normally this will not fail...
-                        throw new Win32Exception();
-                    }
-
-                    if (Kernel32.ResumeThread(threadHandle) == -1)
-                    {
-                        // Normally this will not fail...
-                        throw new Win32Exception();
-                    }
-
                     return new WindowsChildProcessState(processId, processHandle, jobObjectHandle, pseudoConsole, startInfo.AllowSignal);
                 }
             }