From 52ac05ee175c1116e5ec50950111cb155a5bc465 Mon Sep 17 00:00:00 2001
From: ChsBuffer <33744752+chsbuffer@users.noreply.github.com>
Date: Mon, 26 Oct 2020 17:36:50 +0800
Subject: [PATCH] feat: add ChildProcessTracker
---
Netch/Controllers/HTTPController.cs | 5 +-
Netch/Controllers/MainController.cs | 23 ++++-
Netch/Utils/ChildProcessTracker.cs | 145 ++++++++++++++++++++++++++++
3 files changed, 171 insertions(+), 2 deletions(-)
create mode 100644 Netch/Utils/ChildProcessTracker.cs
diff --git a/Netch/Controllers/HTTPController.cs b/Netch/Controllers/HTTPController.cs
index 83a5c798..1ff5fdfd 100644
--- a/Netch/Controllers/HTTPController.cs
+++ b/Netch/Controllers/HTTPController.cs
@@ -31,7 +31,10 @@ namespace Netch.Controllers
try
{
- pPrivoxyController.Start(MainController.ServerController.Server, mode);
+ if (pPrivoxyController.Start(MainController.ServerController.Server, mode))
+ {
+ ChildProcessTracker.AddProcess(pPrivoxyController.Instance);
+ }
if (mode.Type == 3) NativeMethods.SetGlobal($"127.0.0.1:{Global.Settings.HTTPLocalPort}", IEProxyExceptions);
}
diff --git a/Netch/Controllers/MainController.cs b/Netch/Controllers/MainController.cs
index 468e334d..5bf97acc 100644
--- a/Netch/Controllers/MainController.cs
+++ b/Netch/Controllers/MainController.cs
@@ -122,6 +122,14 @@ namespace Netch.Controllers
Global.MainForm.StatusText(i18N.TranslateFormat("Starting {0}", controller.Name));
if (controller.Start(in server, mode))
{
+ if (controller is Guard guard)
+ {
+ if (guard.Instance != null)
+ {
+ ChildProcessTracker.AddProcess(guard.Instance);
+ }
+ }
+
UsingPorts.Add(StatusPortInfoText.Socks5Port = controller.Socks5LocalPort());
StatusPortInfoText.ShareLan = controller.LocalAddress == "0.0.0.0";
@@ -143,7 +151,20 @@ namespace Netch.Controllers
if (ModeController != null)
{
Global.MainForm.StatusText(i18N.TranslateFormat("Starting {0}", ModeController.Name));
- return await Task.Run(() => ModeController.Start(mode));
+ if (await Task.Run(() => ModeController.Start(mode)))
+ {
+ if (ModeController is Guard guard)
+ {
+ if (guard.Instance != null)
+ {
+ ChildProcessTracker.AddProcess(guard.Instance);
+ }
+ }
+
+ return true;
+ }
+
+ return false;
}
return true;
diff --git a/Netch/Utils/ChildProcessTracker.cs b/Netch/Utils/ChildProcessTracker.cs
new file mode 100644
index 00000000..750954cd
--- /dev/null
+++ b/Netch/Utils/ChildProcessTracker.cs
@@ -0,0 +1,145 @@
+using System;
+using System.ComponentModel;
+using System.Diagnostics;
+using System.Runtime.InteropServices;
+
+namespace Netch.Utils
+{
+ ///
+ /// Allows processes to be automatically killed if this parent process unexpectedly quits.
+ /// This feature requires Windows 8 or greater. On Windows 7, nothing is done.
+ /// References:
+ /// https://stackoverflow.com/a/4657392/386091
+ /// https://stackoverflow.com/a/9164742/386091
+ public static class ChildProcessTracker
+ {
+ ///
+ /// Add the process to be tracked. If our current process is killed, the child processes
+ /// that we are tracking will be automatically killed, too. If the child process terminates
+ /// first, that's fine, too.
+ ///
+ public static void AddProcess(Process process)
+ {
+ if (s_jobHandle != IntPtr.Zero)
+ {
+ var success = AssignProcessToJobObject(s_jobHandle, process.Handle);
+ if (!success && !process.HasExited)
+ throw new Win32Exception();
+ }
+ }
+
+ static ChildProcessTracker()
+ {
+ // This feature requires Windows 8 or later. To support Windows 7 requires
+ // registry settings to be added if you are using Visual Studio plus an
+ // app.manifest change.
+ // https://stackoverflow.com/a/4232259/386091
+ // https://stackoverflow.com/a/9507862/386091
+ /*if (Environment.OSVersion.Version < new Version(6, 2))
+ return;*/
+
+ // The job name is optional (and can be null) but it helps with diagnostics.
+ // If it's not null, it has to be unique. Use SysInternals' Handle command-line
+ // utility: handle -a ChildProcessTracker
+ var jobName = "ChildProcessTracker" + Process.GetCurrentProcess().Id;
+ s_jobHandle = CreateJobObject(IntPtr.Zero, jobName);
+
+ var info = new JOBOBJECT_BASIC_LIMIT_INFORMATION();
+
+ // This is the key flag. When our process is killed, Windows will automatically
+ // close the job handle, and when that happens, we want the child processes to
+ // be killed, too.
+ info.LimitFlags = JOBOBJECTLIMIT.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
+
+ var extendedInfo = new JOBOBJECT_EXTENDED_LIMIT_INFORMATION
+ {
+ BasicLimitInformation = info
+ };
+
+ var length = Marshal.SizeOf(typeof(JOBOBJECT_EXTENDED_LIMIT_INFORMATION));
+ var extendedInfoPtr = Marshal.AllocHGlobal(length);
+ try
+ {
+ Marshal.StructureToPtr(extendedInfo, extendedInfoPtr, false);
+
+ if (!SetInformationJobObject(s_jobHandle, JobObjectInfoType.ExtendedLimitInformation,
+ extendedInfoPtr, (uint) length))
+ {
+ throw new Win32Exception();
+ }
+ }
+ finally
+ {
+ Marshal.FreeHGlobal(extendedInfoPtr);
+ }
+ }
+
+ [DllImport("kernel32.dll", CharSet = CharSet.Unicode)]
+ private static extern IntPtr CreateJobObject(IntPtr lpJobAttributes, string name);
+
+ [DllImport("kernel32.dll")]
+ private static extern bool SetInformationJobObject(IntPtr job, JobObjectInfoType infoType,
+ IntPtr lpJobObjectInfo, uint cbJobObjectInfoLength);
+
+ [DllImport("kernel32.dll", SetLastError = true)]
+ private static extern bool AssignProcessToJobObject(IntPtr job, IntPtr process);
+
+ // Windows will automatically close any open job handles when our process terminates.
+ // This can be verified by using SysInternals' Handle utility. When the job handle
+ // is closed, the child processes will be killed.
+ private static readonly IntPtr s_jobHandle;
+ }
+
+ public enum JobObjectInfoType
+ {
+ AssociateCompletionPortInformation = 7,
+ BasicLimitInformation = 2,
+ BasicUIRestrictions = 4,
+ EndOfJobTimeInformation = 6,
+ ExtendedLimitInformation = 9,
+ SecurityLimitInformation = 5,
+ GroupInformation = 11
+ }
+
+ [StructLayout(LayoutKind.Sequential)]
+ public struct JOBOBJECT_BASIC_LIMIT_INFORMATION
+ {
+ public Int64 PerProcessUserTimeLimit;
+ public Int64 PerJobUserTimeLimit;
+ public JOBOBJECTLIMIT LimitFlags;
+ public UIntPtr MinimumWorkingSetSize;
+ public UIntPtr MaximumWorkingSetSize;
+ public UInt32 ActiveProcessLimit;
+ public Int64 Affinity;
+ public UInt32 PriorityClass;
+ public UInt32 SchedulingClass;
+ }
+
+ [Flags]
+ public enum JOBOBJECTLIMIT : uint
+ {
+ JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x2000
+ }
+
+ [StructLayout(LayoutKind.Sequential)]
+ public struct IO_COUNTERS
+ {
+ public UInt64 ReadOperationCount;
+ public UInt64 WriteOperationCount;
+ public UInt64 OtherOperationCount;
+ public UInt64 ReadTransferCount;
+ public UInt64 WriteTransferCount;
+ public UInt64 OtherTransferCount;
+ }
+
+ [StructLayout(LayoutKind.Sequential)]
+ public struct JOBOBJECT_EXTENDED_LIMIT_INFORMATION
+ {
+ public JOBOBJECT_BASIC_LIMIT_INFORMATION BasicLimitInformation;
+ public IO_COUNTERS IoInfo;
+ public UIntPtr ProcessMemoryLimit;
+ public UIntPtr JobMemoryLimit;
+ public UIntPtr PeakProcessMemoryUsed;
+ public UIntPtr PeakJobMemoryUsed;
+ }
+}
\ No newline at end of file