1  //
  2  // Copyright (c) Microsoft Corporation.  All rights reserved.
  3  //
  4  
  5  #pragma once
  6  #include "pch.h"
  7  
  8  #include "ETWTrace.h"
  9  
 10  #include <thread>
 11  
 12  #include <wil\stl.h>
 13  #include <wil\win32_helpers.h>
 14  
 15  namespace fs = std::filesystem;
 16  
 17  namespace
 18  {
 19      constexpr inline const wchar_t* DataDiagnosticsRegKey = L"Software\\Classes\\PowerToys";
 20      constexpr inline const wchar_t* DataDiagnosticsRegValueName = L"AllowDataDiagnostics";
 21      constexpr inline const wchar_t* ViewDataDiagnosticsRegValueName = L"DataDiagnosticsViewEnabled";
 22  
 23      inline std::wstring get_root_save_folder_location()
 24      {
 25          PWSTR local_app_path;
 26          winrt::check_hresult(SHGetKnownFolderPath(FOLDERID_LocalAppData, 0, NULL, &local_app_path));
 27          std::wstring result{ local_app_path };
 28          CoTaskMemFree(local_app_path);
 29  
 30          result += L"\\Microsoft\\PowerToys";
 31          std::filesystem::path save_path(result);
 32          if (!std::filesystem::exists(save_path))
 33          {
 34              std::filesystem::create_directories(save_path);
 35          }
 36          return result;
 37      }
 38  
 39      bool IsDataDiagnosticsEnabled()
 40      {
 41          HKEY key{};
 42          if (RegOpenKeyExW(HKEY_CURRENT_USER,
 43                            DataDiagnosticsRegKey,
 44                            0,
 45                            KEY_READ,
 46                            &key) != ERROR_SUCCESS)
 47          {
 48              return false;
 49          }
 50  
 51          DWORD isDataDiagnosticsEnabled = 0;
 52          DWORD size = sizeof(isDataDiagnosticsEnabled);
 53  
 54          if (RegGetValueW(
 55                  HKEY_CURRENT_USER,
 56                  DataDiagnosticsRegKey,
 57                  DataDiagnosticsRegValueName,
 58                  RRF_RT_REG_DWORD,
 59                  nullptr,
 60                  &isDataDiagnosticsEnabled,
 61                  &size) != ERROR_SUCCESS)
 62          {
 63              RegCloseKey(key);
 64              return false;
 65          }
 66          RegCloseKey(key);
 67  
 68          return isDataDiagnosticsEnabled;
 69      }
 70  
 71      bool isViewDataDiagnosticEnabled()
 72      {
 73          HKEY key{};
 74          if (RegOpenKeyExW(HKEY_CURRENT_USER,
 75                            DataDiagnosticsRegKey,
 76                            0,
 77                            KEY_READ,
 78                            &key) != ERROR_SUCCESS)
 79          {
 80              return false;
 81          }
 82  
 83          DWORD isDataDiagnosticsEnabled = 0;
 84          DWORD size = sizeof(isDataDiagnosticsEnabled);
 85  
 86          if (RegGetValueW(
 87                  HKEY_CURRENT_USER,
 88                  DataDiagnosticsRegKey,
 89                  ViewDataDiagnosticsRegValueName,
 90                  RRF_RT_REG_DWORD,
 91                  nullptr,
 92                  &isDataDiagnosticsEnabled,
 93                  &size) != ERROR_SUCCESS)
 94          {
 95              RegCloseKey(key);
 96              return false;
 97          }
 98          RegCloseKey(key);
 99  
100          return isDataDiagnosticsEnabled == 1;
101      }
102  
103  }
104  
105  namespace Shared
106  {
107      namespace Trace
108      {
109          ETWTrace::ETWTrace()
110          {
111              GUID id;
112              if (SUCCEEDED(CLSIDFromString(PowerToysProviderGUID, &id)))
113              {
114                  m_providerGUID = id;
115              }
116  
117              fs::path outputFolder = get_root_save_folder_location();
118              m_etwFolder = (outputFolder / c_etwFolderName);
119          }
120  
121          ETWTrace::ETWTrace(const std::wstring& etlFileNameOverride) :
122              ETWTrace()
123          {
124              m_etlFileNameOverride = etlFileNameOverride;
125          }
126  
127          ETWTrace::~ETWTrace()
128          {
129              Flush();
130              Stop();
131              m_etwFolder.clear();
132              m_providerGUID = {};
133          }
134  
135          void ETWTrace::UpdateState(bool tracing)
136          {
137              if (tracing)
138              {
139                  Start();
140              }
141              else
142              {
143                  Stop();
144              }
145          }
146  
147          void ETWTrace::Flush()
148          {
149              if (m_tracing)
150              {
151                  Control(EVENT_TRACE_CONTROL_FLUSH);
152                  // Control(EVENT_TRACE_CONTROL_INCREMENT_FILE);
153              }
154          }
155  
156          void ETWTrace::CreateEtwFolderIfNeeded()
157          {
158              if (!std::filesystem::exists(m_etwFolder))
159              {
160                  std::filesystem::create_directories(m_etwFolder);
161              }
162              else if (!std::filesystem::is_directory(m_etwFolder))
163              {
164                  std::filesystem::remove(m_etwFolder);
165                  std::filesystem::create_directory(m_etwFolder);
166              }
167  
168              THROW_HR_IF(E_UNEXPECTED, !std::filesystem::exists(m_etwFolder));
169          }
170  
171          void ETWTrace::InitEventTraceProperties()
172          {
173              const std::filesystem::path exePath(wil::GetModuleFileNameW<std::wstring>(nullptr));
174              const auto exeName = exePath.stem().wstring();
175  
176              auto now = std::chrono::system_clock::now();
177              auto timeNow = std::chrono::system_clock::to_time_t(now);
178              std::wstringstream dateTime;
179              struct tm timeInfo
180              {
181              };
182              errno_t err = localtime_s(&timeInfo, &timeNow);
183              if (err == 0)
184              {
185                  dateTime << std::put_time(&timeInfo, L"-%m-%d-%Y__%H_%M_%S");
186              }
187  
188              if (m_etlFileNameOverride.empty())
189              {
190                  m_sessionName = wil::str_printf<std::wstring>(L"%ws-%d%ws", exeName.c_str(), GetCurrentProcessId(), dateTime.str().c_str());
191              }
192              else
193              {
194                  m_sessionName = wil::str_printf<std::wstring>(L"%ws-%d%ws", m_etlFileNameOverride.c_str(), GetCurrentProcessId(), dateTime.str().c_str());
195              }
196  
197              std::replace(m_sessionName.begin(), m_sessionName.end(), '.', '_');
198  
199              const ULONG etwSessionNameCharCount = static_cast<ULONG>(m_sessionName.size() + 1);
200              const ULONG etwSessionNameByteSize = etwSessionNameCharCount * sizeof(m_sessionName[0]);
201  
202              auto etlFileNameFormattedCounter = m_sessionName + c_etwNewFileFormattedCounter;
203              std::filesystem::path etlFilePath = m_etwFolder / etlFileNameFormattedCounter;
204              etlFilePath.replace_extension(c_etwFileNameEnd);
205              THROW_HR_IF(E_UNEXPECTED, etlFilePath.empty());
206  
207              const auto etlFilePathStr = etlFilePath.wstring();
208              // std::string/wstring returns number of characters not including the null terminator, so add +1 for that.
209              const ULONG etwFilePathCharCount = static_cast<ULONG>(etlFilePathStr.size() + 1);
210              const ULONG etwFilePathByteSize = etwFilePathCharCount * sizeof(etlFilePathStr[0]);
211  
212              const ULONG bufferSizeInBytes = sizeof(EVENT_TRACE_PROPERTIES) + etwSessionNameByteSize + etwFilePathByteSize;
213              auto eventTracePropertiesBuffer = std::make_unique<unsigned char[]>(bufferSizeInBytes);
214              ZeroMemory(eventTracePropertiesBuffer.get(), bufferSizeInBytes);
215              auto eventTraceProperties = reinterpret_cast<EVENT_TRACE_PROPERTIES*>(eventTracePropertiesBuffer.get());
216  
217              eventTraceProperties->Wnode.BufferSize = bufferSizeInBytes;
218              eventTraceProperties->Wnode.Flags = WNODE_FLAG_TRACED_GUID;
219              eventTraceProperties->Wnode.ClientContext = 1;
220              eventTraceProperties->Wnode.Guid = m_providerGUID;
221              eventTraceProperties->BufferSize = 4; // 4KB, the minimum size
222              eventTraceProperties->LogFileMode = EVENT_TRACE_PRIVATE_LOGGER_MODE | EVENT_TRACE_PRIVATE_IN_PROC | EVENT_TRACE_FILE_MODE_NEWFILE;
223              eventTraceProperties->MaximumFileSize = 1; // 1 MB
224  
225              // LoggerName is placed at the end of EVENT_TRACE_PROPERTIES structure
226              eventTraceProperties->LoggerNameOffset = sizeof(EVENT_TRACE_PROPERTIES);
227              wcsncpy_s(reinterpret_cast<LPWSTR>(eventTracePropertiesBuffer.get() + eventTraceProperties->LoggerNameOffset), etwSessionNameCharCount, m_sessionName.c_str(), etwSessionNameCharCount);
228  
229              // LogFileName is placed at the end of the Logger Name
230              eventTraceProperties->LogFileNameOffset = eventTraceProperties->LoggerNameOffset + etwSessionNameByteSize;
231              wcsncpy_s(reinterpret_cast<LPWSTR>(eventTracePropertiesBuffer.get() + eventTraceProperties->LogFileNameOffset), etwFilePathCharCount, etlFilePathStr.c_str(), etwFilePathCharCount);
232  
233              m_eventTracePropertiesBuffer = std::move(eventTracePropertiesBuffer);
234          }
235  
236          void ETWTrace::Start()
237          {
238              if (m_tracing)
239              {
240                  return;
241              }
242  
243              if (!IsDataDiagnosticsEnabled())
244              {
245                  return;
246              }
247  
248              if (!isViewDataDiagnosticEnabled())
249              {
250                  return;
251              }
252  
253              CreateEtwFolderIfNeeded();
254              InitEventTraceProperties();
255  
256              auto eventTraceProperties = reinterpret_cast<EVENT_TRACE_PROPERTIES*>(m_eventTracePropertiesBuffer.get());
257              THROW_IF_WIN32_ERROR(StartTrace(&m_traceHandle, m_sessionName.c_str(), eventTraceProperties));
258              Enable(EVENT_CONTROL_CODE_ENABLE_PROVIDER);
259  
260              m_tracing = true;
261  
262              m_flushing_thread = std::thread([this] { FlushWorker(); });
263          }
264  
265          void ETWTrace::Stop()
266          {
267              if (!m_tracing)
268              {
269                  return;
270              }
271  
272              Enable(EVENT_CONTROL_CODE_DISABLE_PROVIDER);
273  
274              // ControlTrace with EVENT_TRACE_CONTROL_STOP on the trace handle,
275              // which is equivalent to calling CloseTrace() on the trace handle.
276              Control(EVENT_TRACE_CONTROL_STOP);
277  
278              m_traceHandle = INVALID_PROCESSTRACE_HANDLE;
279              m_eventTracePropertiesBuffer.reset();
280              m_tracing = false;
281              m_terminate_flushing_thread.notify_one();
282              m_flushing_thread.join();
283          }
284  
285          void ETWTrace::Control(ULONG traceControlCode)
286          {
287              auto eventTraceProperties = reinterpret_cast<EVENT_TRACE_PROPERTIES*>(m_eventTracePropertiesBuffer.get());
288              const ULONG result = ControlTrace(m_traceHandle, m_sessionName.c_str(), eventTraceProperties, traceControlCode);
289              THROW_IF_FAILED(HRESULT_FROM_WIN32(result));
290          }
291  
292          void ETWTrace::Enable(ULONG eventControlCode)
293          {
294              // Control the main provider
295              THROW_IF_WIN32_ERROR(EnableTraceEx2(m_traceHandle, &m_providerGUID, eventControlCode, TRACE_LEVEL_VERBOSE, 0, 0, 0, nullptr));
296          }
297  
298          void ETWTrace::FlushWorker()
299          {
300              std::unique_lock<std::mutex> lock(m_mutex);
301              while (m_tracing)
302              {
303                  m_terminate_flushing_thread.wait_for(lock,
304                                  std::chrono::seconds(30),
305                                  [this]() { return !m_tracing.load(); });
306                  Flush();
307              }
308          }
309      }
310  }