From 2447a299c83f6647165dd371b3446024617c916f Mon Sep 17 00:00:00 2001
From: Duncan Ogilvie <mr.exodia.tpodt@gmail.com>
Date: Mon, 19 Oct 2020 22:57:19 +0200
Subject: [PATCH] Improve debug privilege functionality

---
 TitanEngine/Global.Engine.cpp        | 18 ++----------------
 TitanEngine/TitanEngine.Debugger.cpp |  9 +++++----
 2 files changed, 7 insertions(+), 20 deletions(-)

diff --git a/TitanEngine/Global.Engine.cpp b/TitanEngine/Global.Engine.cpp
index d938aca..1c936b4 100644
--- a/TitanEngine/Global.Engine.cpp
+++ b/TitanEngine/Global.Engine.cpp
@@ -2066,24 +2066,10 @@ DWORD EngineSetDebugPrivilege(HANDLE hProcess, bool bEnablePrivilege)
 
 HANDLE EngineOpenProcess(DWORD dwDesiredAccess, bool bInheritHandle, DWORD dwProcessId)
 {
-    if(engineEnableDebugPrivilege)
-        EngineSetDebugPrivilege(GetCurrentProcess(), true);
-    HANDLE hProcess = OpenProcess(dwDesiredAccess, bInheritHandle, dwProcessId);
-    DWORD dwLastError = GetLastError();
-    if(engineEnableDebugPrivilege)
-        EngineSetDebugPrivilege(GetCurrentProcess(), false);
-    SetLastError(dwLastError);
-    return hProcess;
+    return OpenProcess(dwDesiredAccess, bInheritHandle, dwProcessId);
 }
 
 HANDLE EngineOpenThread(DWORD dwDesiredAccess, bool bInheritHandle, DWORD dwThreadId)
 {
-    if(engineEnableDebugPrivilege)
-        EngineSetDebugPrivilege(GetCurrentProcess(), true);
-    HANDLE hThread = OpenThread(dwDesiredAccess, bInheritHandle, dwThreadId);
-    DWORD dwLastError = GetLastError();
-    if(engineEnableDebugPrivilege)
-        EngineSetDebugPrivilege(GetCurrentProcess(), false);
-    SetLastError(dwLastError);
-    return hThread;
+    return OpenThread(dwDesiredAccess, bInheritHandle, dwThreadId);
 }
diff --git a/TitanEngine/TitanEngine.Debugger.cpp b/TitanEngine/TitanEngine.Debugger.cpp
index f2345fa..88705bb 100644
--- a/TitanEngine/TitanEngine.Debugger.cpp
+++ b/TitanEngine/TitanEngine.Debugger.cpp
@@ -79,8 +79,6 @@ __declspec(dllexport) void* TITCALL InitDebugW(wchar_t* szFileName, wchar_t* szC
     }
     if(CreateProcessW(szFileNameCreateProcess, szCommandLineCreateProcess, NULL, NULL, false, DEBUG_PROCESS | DEBUG_ONLY_THIS_PROCESS | DebugConsoleFlag | CREATE_NEW_CONSOLE, NULL, szCurrentFolder, &dbgStartupInfo, &dbgProcessInformation))
     {
-        if(engineEnableDebugPrivilege)
-            EngineSetDebugPrivilege(GetCurrentProcess(), false);
         DebugAttachedToProcess = false;
         DebugAttachedProcessCallBack = NULL;
         return &dbgProcessInformation;
@@ -551,8 +549,6 @@ __declspec(dllexport) bool TITCALL AttachDebugger(DWORD ProcessId, bool KillOnEx
         }
         if((engineSafeAttach ? DebugActiveProcess_ : DebugActiveProcess)(ProcessId))
         {
-            if(engineEnableDebugPrivilege)
-                EngineSetDebugPrivilege(GetCurrentProcess(), false);
             funcDebugSetProcessKillOnExit = GetProcAddress(GetModuleHandleA("kernel32.dll"), "DebugSetProcessKillOnExit");
             if(funcDebugSetProcessKillOnExit != NULL)
             {
@@ -570,6 +566,11 @@ __declspec(dllexport) bool TITCALL AttachDebugger(DWORD ProcessId, bool KillOnEx
             return true;
         }
     }
+    if (engineEnableDebugPrivilege)
+    {
+        EngineSetDebugPrivilege(GetCurrentProcess(), false);
+        DebugRemoveDebugPrivilege = false;
+    }
     return false;
 }