From 6b07af9e0574fa885d8150e4fe40a425c2396a1d Mon Sep 17 00:00:00 2001
From: Victor Zarubkin <v.s.zarubkin@gmail.com>
Date: Sat, 17 Sep 2016 01:05:50 +0300
Subject: [PATCH] (ETW) Event tracing fix for situation when application was
 not launched from Visual Studio: to be able to get process name application
 requires debug privilege level.

---
 src/event_trace_win.cpp | 38 +++++++++++++++++++++++++++++++++++---
 src/event_trace_win.h   |  2 ++
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/src/event_trace_win.cpp b/src/event_trace_win.cpp
index 35e0bb9..5dac433 100644
--- a/src/event_trace_win.cpp
+++ b/src/event_trace_win.cpp
@@ -4,11 +4,12 @@
 #include <memory.h>
 #include <chrono>
 #include <unordered_map>
-#include "event_trace_win.h"
-#include "Psapi.h"
 #include "profiler/profiler.h"
 #include "profile_manager.h"
 
+#include "event_trace_win.h"
+#include <Psapi.h>
+
 //////////////////////////////////////////////////////////////////////////
 //////////////////////////////////////////////////////////////////////////
 
@@ -96,6 +97,8 @@ namespace profiler {
                     // But it works fine with PROCESS_QUERY_LIMITED_INFORMATION instead of PROCESS_QUERY_INFORMATION.
                     // 
                     // See https://msdn.microsoft.com/en-us/library/windows/desktop/ms683196(v=vs.85).aspx
+                    //auto hProc = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, pid);
+                    //if (hProc == nullptr)
                     auto hProc = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION | PROCESS_VM_READ, FALSE, pid);
                     if (hProc != nullptr)
                     {
@@ -115,7 +118,8 @@ namespace profiler {
                     }
                     else
                     {
-                        //printf("Can not OpenProcess(%u);\n", pid);
+                        //auto err = GetLastError();
+                        //printf("OpenProcess(%u) fail: GetLastError() == %u\n", pid, err);
                         pinfo->valid = -1;
                     }
                 }
@@ -165,6 +169,29 @@ namespace profiler {
         m_lowPriority.store(_value, ::std::memory_order_release);
     }
 
+    bool EasyEventTracer::setDebugPrivilege()
+    {
+        bool success = false;
+
+        HANDLE hToken = nullptr;
+        if (OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken))
+        {
+            LUID privilegyId;
+            if (LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &privilegyId))
+            {
+                TOKEN_PRIVILEGES tokenPrivilegy;
+                tokenPrivilegy.PrivilegeCount = 1;
+                tokenPrivilegy.Privileges[0].Luid = privilegyId;
+                tokenPrivilegy.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
+                success = AdjustTokenPrivileges(hToken, FALSE, &tokenPrivilegy, sizeof(TOKEN_PRIVILEGES), NULL, NULL) != FALSE;
+            }
+
+            CloseHandle(hToken);
+        }
+
+        return success;
+    }
+
     ::profiler::EventTracingEnableStatus EasyEventTracer::startTrace(bool _force, int _step)
     {
         auto startTraceResult = StartTrace(&m_sessionHandle, KERNEL_LOGGER_NAME, props());
@@ -208,6 +235,11 @@ namespace profiler {
         if (m_bEnabled)
             return EVENT_TRACING_LAUNCHED_SUCCESSFULLY;
 
+        // Trying to set debug privilege for current process
+        // to be able to get other process information (process name)
+        if (!m_bPrivilegeSet)
+            m_bPrivilegeSet = setDebugPrivilege();
+
         // Clear properties
         memset(&m_properties, 0, sizeof(m_properties));
         m_properties.base.Wnode.BufferSize = sizeof(m_properties);
diff --git a/src/event_trace_win.h b/src/event_trace_win.h
index 54f1215..c2e1a72 100644
--- a/src/event_trace_win.h
+++ b/src/event_trace_win.h
@@ -39,6 +39,7 @@ namespace profiler {
         TRACEHANDLE         m_sessionHandle = INVALID_PROCESSTRACE_HANDLE;
         TRACEHANDLE          m_openedHandle = INVALID_PROCESSTRACE_HANDLE;
         bool                     m_bEnabled = false;
+        bool                m_bPrivilegeSet = false;
 
     public:
 
@@ -59,6 +60,7 @@ namespace profiler {
         }
 
         ::profiler::EventTracingEnableStatus startTrace(bool _force, int _step = 0);
+        bool setDebugPrivilege();
 
     }; // END of class EasyEventTracer.