From f74c40671b5ddec0b8bd8327ee486e6a9499f294 Mon Sep 17 00:00:00 2001
From: Gerardo Grignoli <gerardog@gmail.com>
Date: Sat, 12 Nov 2022 23:45:06 -0300
Subject: [PATCH] Feature: JobObject for TokenSwitch mode

---
 src/gsudo/Native/JobObject.cs                 | 144 ++++++++++++++++++
 .../ProcessRenderers/TokenSwitchRenderer.cs   |  12 +-
 2 files changed, 154 insertions(+), 2 deletions(-)
 create mode 100644 src/gsudo/Native/JobObject.cs

diff --git a/src/gsudo/Native/JobObject.cs b/src/gsudo/Native/JobObject.cs
new file mode 100644
index 00000000..ea1236f2
--- /dev/null
+++ b/src/gsudo/Native/JobObject.cs
@@ -0,0 +1,144 @@
+using System;
+using System.Diagnostics;
+using System.Runtime.InteropServices;
+
+namespace gsudo.Native
+{
+    public class JobObject : IDisposable
+    {
+        [DllImport("kernel32.dll", CharSet = CharSet.Unicode)]
+        static extern IntPtr CreateJobObject(IntPtr a, string lpName);
+
+        [DllImport("kernel32.dll")]
+        static extern bool SetInformationJobObject(IntPtr hJob, JobObjectInfoType infoType, IntPtr lpJobObjectInfo, uint cbJobObjectInfoLength);
+
+        [DllImport("kernel32.dll", SetLastError = true)]
+        static extern bool AssignProcessToJobObject(IntPtr job, IntPtr process);
+
+        [DllImport("kernel32.dll", SetLastError = true)]
+        [return: MarshalAs(UnmanagedType.Bool)]
+        static extern bool CloseHandle(IntPtr hObject);
+
+        private IntPtr handle;
+        private bool disposed;
+
+        public JobObject()
+        {
+            handle = CreateJobObject(IntPtr.Zero, null);
+
+            var info = new JOBOBJECT_BASIC_LIMIT_INFORMATION
+            {
+                LimitFlags = 0x2000
+            };
+
+            var extendedInfo = new JOBOBJECT_EXTENDED_LIMIT_INFORMATION
+            {
+                BasicLimitInformation = info
+            };
+
+            int length = Marshal.SizeOf(typeof(JOBOBJECT_EXTENDED_LIMIT_INFORMATION));
+            IntPtr extendedInfoPtr = Marshal.AllocHGlobal(length);
+            Marshal.StructureToPtr(extendedInfo, extendedInfoPtr, false);
+
+            if (!SetInformationJobObject(handle, JobObjectInfoType.ExtendedLimitInformation, extendedInfoPtr, (uint)length))
+                throw new System.ComponentModel.Win32Exception();
+        }
+
+        public void Dispose()
+        {
+            Dispose(true);
+            GC.SuppressFinalize(this);
+        }
+
+        private void Dispose(bool disposing)
+        {
+            if (disposed)
+                return;
+
+            if (disposing) { }
+
+            Close();
+            disposed = true;
+        }
+
+        public void Close()
+        {
+            CloseHandle(handle);
+            handle = IntPtr.Zero;
+        }
+
+        public bool AddProcess(IntPtr processHandle)
+        {
+            return AssignProcessToJobObject(handle, processHandle);
+        }
+
+        public bool AddProcess(int processId)
+        {
+            return AddProcess(Process.GetProcessById(processId).Handle);
+        }
+
+    }
+
+    #region Helper classes
+
+    [StructLayout(LayoutKind.Sequential)]
+    struct IO_COUNTERS
+    {
+        public ulong ReadOperationCount;
+        public ulong WriteOperationCount;
+        public ulong OtherOperationCount;
+        public ulong ReadTransferCount;
+        public ulong WriteTransferCount;
+        public ulong OtherTransferCount;
+    }
+
+
+    [StructLayout(LayoutKind.Sequential)]
+    struct JOBOBJECT_BASIC_LIMIT_INFORMATION
+    {
+        public long PerProcessUserTimeLimit;
+        public long PerJobUserTimeLimit;
+        public uint LimitFlags;
+        public UIntPtr MinimumWorkingSetSize;
+        public UIntPtr MaximumWorkingSetSize;
+        public uint ActiveProcessLimit;
+        public UIntPtr Affinity;
+        public uint PriorityClass;
+        public uint SchedulingClass;
+    }
+
+/*
+    [StructLayout(LayoutKind.Sequential)]
+    public struct SECURITY_ATTRIBUTES
+    {
+        public uint nLength;
+        public IntPtr lpSecurityDescriptor;
+        public int bInheritHandle;
+    }
+*/
+
+    [StructLayout(LayoutKind.Sequential)]
+    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;
+    }
+
+    public enum JobObjectInfoType
+    {
+        AssociateCompletionPortInformation = 7,
+        BasicLimitInformation = 2,
+        BasicUIRestrictions = 4,
+        EndOfJobTimeInformation = 6,
+        ExtendedLimitInformation = 9,
+        SecurityLimitInformation = 5,
+        GroupInformation = 11
+    }
+
+    #endregion
+
+}
diff --git a/src/gsudo/ProcessRenderers/TokenSwitchRenderer.cs b/src/gsudo/ProcessRenderers/TokenSwitchRenderer.cs
index c96700a9..b51902ad 100644
--- a/src/gsudo/ProcessRenderers/TokenSwitchRenderer.cs
+++ b/src/gsudo/ProcessRenderers/TokenSwitchRenderer.cs
@@ -22,6 +22,7 @@ class TokenSwitchRenderer : IProcessRenderer
         private readonly SafeProcessHandle _process;
         private readonly ProcessApi.PROCESS_INFORMATION _processInformation;
         private readonly ManualResetEventSlim tokenSwitchSuccessEvent = new ManualResetEventSlim(false);
+        private JobObject _job;
 
         internal TokenSwitchRenderer(Connection connection, ElevationRequest elevationRequest)
         {
@@ -55,7 +56,14 @@ internal TokenSwitchRenderer(Connection connection, ElevationRequest elevationRe
                 args = elevationRequest.Arguments;
             }
 
-            _process = ProcessFactory.CreateProcessAsUserWithFlags(exeName, args, dwCreationFlags, out _processInformation);
+            _process = ProcessFactory.CreateProcessAsUserWithFlags(exeName, args, dwCreationFlags, out _processInformation);
+            
+            if (elevationRequest.Wait)
+            {
+                _job = new JobObject();
+                Logger.Instance.Log($"Created JobObject", LogLevel.Debug);
+                _job.AddProcess(_process.DangerousGetHandle());
+            }
 
             elevationRequest.TargetProcessId = _processInformation.dwProcessId;
             if (!elevationRequest.NewWindow)
@@ -126,7 +134,7 @@ public Task<int> GetResult()
 
 
         enum Mode { Normal, Error};
-        Mode CurrentMode = Mode.Normal;
+        Mode CurrentMode = Mode.Normal;
 
         static readonly string[] TOKENS = new string[] { "\0", Constants.TOKEN_ERROR, Constants.TOKEN_SUCCESS };