From 5e4193cb13bf53ca17ef55d181ee86e87c7c11ca Mon Sep 17 00:00:00 2001
From: Victor Zarubkin <v.s.zarubkin@gmail.com>
Date: Sun, 11 Dec 2016 03:31:22 +0300
Subject: [PATCH] (ETW) setDebugPrivelege() rename

---
 easy_profiler_core/event_trace_win.cpp | 46 +++++++++++++++++++-------
 easy_profiler_core/event_trace_win.h   |  3 +-
 2 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/easy_profiler_core/event_trace_win.cpp b/easy_profiler_core/event_trace_win.cpp
index 0ac63d5..388770b 100644
--- a/easy_profiler_core/event_trace_win.cpp
+++ b/easy_profiler_core/event_trace_win.cpp
@@ -246,15 +246,14 @@ namespace profiler {
         m_lowPriority.store(_value, ::std::memory_order_release);
     }
 
-    bool EasyEventTracer::setDebugPrivilege()
+    bool setPrivilege(HANDLE hToken, LPCSTR _privelegeName)
     {
         bool success = false;
 
-        HANDLE hToken = nullptr;
-        if (OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken))
+        if (hToken)
         {
             LUID privilegyId;
-            if (LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &privilegyId))
+            if (LookupPrivilegeValue(NULL, _privelegeName, &privilegyId))
             {
                 TOKEN_PRIVILEGES tokenPrivilege;
                 tokenPrivilege.PrivilegeCount = 1;
@@ -262,18 +261,45 @@ namespace profiler {
                 tokenPrivilege.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
                 success = AdjustTokenPrivileges(hToken, FALSE, &tokenPrivilege, sizeof(TOKEN_PRIVILEGES), NULL, NULL) != FALSE;
             }
-
-            CloseHandle(hToken);
         }
 
 #if EASY_LOG_ENABLED != 0
         if (!success)
-            ::std::cerr << "Warning: EasyProfiler failed to set Debug privelege for the application. Some context switch events could not get process name.\n";
+            ::std::cerr << "Warning: EasyProfiler failed to set " << _privelegeName << " privelege for the application.\n";
 #endif
 
         return success;
     }
 
+    void EasyEventTracer::setProcessPrivileges()
+    {
+        static bool alreadySet = false;
+        if (alreadySet)
+            return;
+
+        alreadySet = true;
+
+        HANDLE hToken = nullptr;
+        if (OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken))
+        {
+#if EASY_LOG_ENABLED != 0
+            const bool success = setPrivilege(hToken, SE_DEBUG_NAME);
+            if (!success)
+                ::std::cerr << "Warning: Some context switch events could not get process name.\n";
+#else
+            setPrivilege(hToken, SE_DEBUG_NAME);
+#endif
+            
+            CloseHandle(hToken);
+        }
+#if EASY_LOG_ENABLED != 0
+        else
+        {
+            ::std::cerr << "Warning: EasyProfiler failed to open process to adjust priveleges.\n";
+        }
+#endif
+    }
+
     ::profiler::EventTracingEnableStatus EasyEventTracer::startTrace(bool _force, int _step)
     {
         auto startTraceResult = StartTrace(&m_sessionHandle, KERNEL_LOGGER_NAME, props());
@@ -355,12 +381,8 @@ namespace profiler {
         /*
         Trying to set debug privilege for current process
         to be able to get other process information (process name).
-
-        Also it seems that debug privelege lets you to launch
-        event tracing without Administrator access rights.
         */
-        if (!m_bPrivilegeSet)
-            m_bPrivilegeSet = setDebugPrivilege();
+        EasyEventTracer::setProcessPrivileges();
 
         // Clear properties
         memset(&m_properties, 0, sizeof(m_properties));
diff --git a/easy_profiler_core/event_trace_win.h b/easy_profiler_core/event_trace_win.h
index 780d52d..21c27f3 100644
--- a/easy_profiler_core/event_trace_win.h
+++ b/easy_profiler_core/event_trace_win.h
@@ -85,7 +85,6 @@ namespace profiler {
         TRACEHANDLE         m_sessionHandle = INVALID_PROCESSTRACE_HANDLE;
         TRACEHANDLE          m_openedHandle = INVALID_PROCESSTRACE_HANDLE;
         bool                     m_bEnabled = false;
-        bool                m_bPrivilegeSet = false;
 
     public:
 
@@ -97,6 +96,7 @@ namespace profiler {
         ::profiler::EventTracingEnableStatus enable(bool _force = false);
         void disable();
         void setLowPriority(bool _value);
+        static void setProcessPrivileges();
 
     private:
 
@@ -108,7 +108,6 @@ namespace profiler {
         }
 
         ::profiler::EventTracingEnableStatus startTrace(bool _force, int _step = 0);
-        bool setDebugPrivilege();
 
     }; // END of class EasyEventTracer.