Fix race condition in Windows suspend process

tgstation · Feb 25, 2024 · 0acd57d · 0acd57d
1 parent 77d8278
commit 0acd57d
Showing 1 changed file with 31 additions and 18 deletions.
diff --git a/src/Tgstation.Server.Host/System/WindowsProcessFeatures.cs b/src/Tgstation.Server.Host/System/WindowsProcessFeatures.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.Diagnostics;
 using System.IO;
 using System.Linq;
@@ -68,28 +69,40 @@ public void SuspendProcess(global::System.Diagnostics.Process process)
 		{
 			ArgumentNullException.ThrowIfNull(process);
 
-			process.Refresh();
-			foreach (ProcessThread thread in process.Threads)
-			{
-				var threadId = (uint)thread.Id;
-				logger.LogTrace("Suspending thread {threadId}...", threadId);
-				var pOpenThread = NativeMethods.OpenThread(NativeMethods.ThreadAccess.SuspendResume, false, threadId);
-				if (pOpenThread == IntPtr.Zero)
-				{
-					logger.LogDebug(new Win32Exception(), "Failed to open thread {threadId}!", threadId);
-					continue;
-				}
+			var suspendedThreadIds = new HashSet<uint>();
+			var suspendedNewThreads = false;
 
-				try
-				{
-					if (NativeMethods.SuspendThread(pOpenThread) == UInt32.MaxValue)
-						throw new Win32Exception();
-				}
-				finally
+			do
+			{
+				process.Refresh();
+				foreach (ProcessThread thread in process.Threads)
 				{
-					NativeMethods.CloseHandle(pOpenThread);
+					var threadId = (uint)thread.Id;
+
+					if (!suspendedThreadIds.Add(threadId))
+						continue;
+
+					suspendedNewThreads = true;
+					logger.LogTrace("Suspending thread {threadId}...", threadId);
+					var pOpenThread = NativeMethods.OpenThread(NativeMethods.ThreadAccess.SuspendResume, false, threadId);
+					if (pOpenThread == IntPtr.Zero)
+					{
+						logger.LogDebug(new Win32Exception(), "Failed to open thread {threadId}!", threadId);
+						continue;
+					}
+
+					try
+					{
+						if (NativeMethods.SuspendThread(pOpenThread) == UInt32.MaxValue)
+							throw new Win32Exception();
+					}
+					finally
+					{
+						NativeMethods.CloseHandle(pOpenThread);
+					}
 				}
 			}
+			while (suspendedNewThreads);
 		}
 
 		/// <inheritdoc />