Improve NtSetInformationProcess hook accuracy

HookLibrary/HookHelper.cpp

@@ -176,6 +176,43 @@ bool IsValidThreadHandle(HANDLE hThread)
 	}
 }
 
+bool HasDebugPrivileges(HANDLE hProcess)
+{
+	HANDLE hToken;
+	LUID luid;
+	if (!OpenProcessToken(hProcess, TOKEN_QUERY, &hToken))
+		return false;
+	if (!LookupPrivilegeValue(nullptr, SE_DEBUG_NAME, &luid))
+		return false;
+
+	DWORD size = 0;
+	bool hasDebugPrivileges = false;
+	void* wheresMalloc = NULL;
+	GetTokenInformation(hToken, TokenPrivileges, NULL, 0, &size);
+
+	if (size > 0 && (wheresMalloc = VirtualAlloc(NULL, size, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE)) != NULL)
+	{
+		TOKEN_PRIVILEGES* tokenPrivileges = (TOKEN_PRIVILEGES*)wheresMalloc;
+		if (GetTokenInformation(hToken, TokenPrivileges, tokenPrivileges, size, &size))
+		{
+			for (unsigned int i = 0; i < tokenPrivileges->PrivilegeCount; ++i)
+			{
+				LUID_AND_ATTRIBUTES luidAndAttributes = tokenPrivileges->Privileges[i];
+				if (luidAndAttributes.Luid.LowPart == luid.LowPart && luidAndAttributes.Luid.HighPart == luid.HighPart)
+				{
+					hasDebugPrivileges = ((luidAndAttributes.Attributes & SE_PRIVILEGE_ENABLED) != 0) ||
+										((luidAndAttributes.Attributes & SE_PRIVILEGE_ENABLED_BY_DEFAULT) != 0 &&
+										(luidAndAttributes.Attributes & SE_PRIVILEGE_REMOVED) == 0);
+				}
+			}
+		}
+		VirtualFree(tokenPrivileges, size, MEM_RELEASE);
+	}
+
+	CloseHandle(hToken);
+	return hasDebugPrivileges;
+}
+
 OSVERSIONINFO versionInfo = { 0 };
 
 bool IsAtleastVista()

HookLibrary/HookHelper.h

@@ -5,6 +5,7 @@
 bool IsValidHandle(HANDLE hHandle);
 bool IsValidThreadHandle(HANDLE hThread);
 bool IsValidProcessHandle(HANDLE hProcess);
+bool HasDebugPrivileges(HANDLE hProcess);
 DWORD GetExplorerProcessId();
 DWORD GetCsrssProcessId();
 DWORD GetProcessIdByName(const WCHAR * processName);

HookLibrary/HookedFunctions.cpp

@@ -123,21 +123,54 @@ NTSTATUS NTAPI HookedNtSetInformationProcess(HANDLE ProcessHandle, PROCESSINFOCL
 {
     if (ProcessHandle == NtCurrentProcess || GetCurrentProcessId() == GetProcessIdByProcessHandle(ProcessHandle))
     {
-        if (ProcessInformationClass == ProcessBreakOnTermination && ProcessInformation != 0 && sizeof(ULONG) == ProcessInformationLength)
+        if (ProcessInformationClass == ProcessBreakOnTermination)
         {
+			if (ProcessInformationLength != sizeof(ULONG))
+			{
+				return STATUS_INFO_LENGTH_MISMATCH;
+			}
+
+			// NtSetInformationProcess will happily dereference this pointer
+			if (ProcessInformation == NULL)
+			{
+				return STATUS_ACCESS_VIOLATION;
+			}
+
+			// A process must have debug privileges enabled to set the ProcessBreakOnTermination flag
+			if (!HasDebugPrivileges(NtCurrentProcess))
+			{
+				return STATUS_PRIVILEGE_NOT_HELD;
+			}
+
             ValueProcessBreakOnTermination = *((ULONG *)ProcessInformation);
             return STATUS_SUCCESS;
         }
 
 		//PROCESS_HANDLE_TRACING_ENABLE -> ULONG, PROCESS_HANDLE_TRACING_ENABLE_EX -> ULONG,ULONG
-		if (ProcessInformationClass == ProcessHandleTracing && ProcessInformation != 0)
+		if (ProcessInformationClass == ProcessHandleTracing)
 		{
-			if (ProcessInformationLength != sizeof(ULONG) && ProcessInformationLength != (sizeof(ULONG)*2))
+			bool enable = ProcessInformationLength != 0; // A length of 0 is valid and indicates we should disable tracing
+			if (enable)
 			{
-				return STATUS_INFO_LENGTH_MISMATCH;
+				if (ProcessInformationLength != sizeof(ULONG) && ProcessInformationLength != (sizeof(ULONG) * 2))
+				{
+					return STATUS_INFO_LENGTH_MISMATCH;
+				}
+
+				// NtSetInformationProcess will happily dereference this pointer
+				if (ProcessInformation == NULL)
+				{
+					return STATUS_ACCESS_VIOLATION;
+				}
+
+				PPROCESS_HANDLE_TRACING_ENABLE_EX phtEx = (PPROCESS_HANDLE_TRACING_ENABLE_EX)ProcessInformation;
+				if (phtEx->Flags != 0)
+				{
+					return STATUS_INVALID_PARAMETER;
+				}
 			}
 
-			IsProcessHandleTracingEnabled = true;
+			IsProcessHandleTracingEnabled = enable;
 			return STATUS_SUCCESS;
 		}
     }

HookLibrary/HookedFunctions.h

@@ -3,6 +3,7 @@
 #include <windows.h>
 
 
+#define STATUS_PRIVILEGE_NOT_HELD ((NTSTATUS)0xC0000061L)
 #define STATUS_PORT_NOT_SET ((NTSTATUS)0xC0000353L)
 #define STATUS_HANDLE_NOT_CLOSABLE ((NTSTATUS)0xC0000235L)
 #define STATUS_INSUFFICIENT_RESOURCES ((NTSTATUS)0xC000009AL)

ntdll/ntdll.h

@@ -522,6 +522,12 @@ typedef struct _PS_ATTRIBUTE_LIST
     PS_ATTRIBUTE Attributes[1];
 } PS_ATTRIBUTE_LIST, *PPS_ATTRIBUTE_LIST;
 
+typedef struct _PROCESS_HANDLE_TRACING_ENABLE_EX // PROCESS_HANDLE_TRACING_ENABLE is simply this struct with only the first field
+{
+	ULONG Flags;
+	ULONG TotalSlots;
+} PROCESS_HANDLE_TRACING_ENABLE_EX, *PPROCESS_HANDLE_TRACING_ENABLE_EX;
+
 //Source: http://processhacker.sourceforge.net
 typedef enum _PROCESSINFOCLASS
 {