[Win] Create ModuleWatcher.

chrome/common/conflicts/module_watcher_win.cc

+// Copyright 2016 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "chrome/common/conflicts/module_watcher_win.h"
+
+#include <windows.h>
+#include <tlhelp32.h>
+#include <winternl.h>  // For UNICODE_STRING.
+
+#include "base/lazy_instance.h"
+#include "base/memory/ptr_util.h"
+#include "base/strings/utf_string_conversions.h"
+#include "base/win/scoped_handle.h"
+
+// These structures and functions are documented in MSDN, see
+// http://msdn.microsoft.com/en-us/library/gg547638(v=vs.85).aspx
+// there are however no headers or import libraries available in the
+// Platform SDK. They are declared outside of the anonymous namespace to
+// allow them to be forward declared in the header file.
+enum {
+  // The DLL was loaded. The NotificationData parameter points to a
+  // LDR_DLL_LOADED_NOTIFICATION_DATA structure.
+  LDR_DLL_NOTIFICATION_REASON_LOADED = 1,
+  // The DLL was unloaded. The NotificationData parameter points to a
+  // LDR_DLL_UNLOADED_NOTIFICATION_DATA structure.
+  LDR_DLL_NOTIFICATION_REASON_UNLOADED = 2,
+};
+
+// Structure that is used for module load notifications.
+struct LDR_DLL_LOADED_NOTIFICATION_DATA {
+  // Reserved.
+  ULONG Flags;
+  // The full path name of the DLL module.
+  PCUNICODE_STRING FullDllName;
+  // The base file name of the DLL module.
+  PCUNICODE_STRING BaseDllName;
+  // A pointer to the base address for the DLL in memory.
+  PVOID DllBase;
+  // The size of the DLL image, in bytes.
+  ULONG SizeOfImage;
+};
+using PLDR_DLL_LOADED_NOTIFICATION_DATA = LDR_DLL_LOADED_NOTIFICATION_DATA*;
+
+// Structure that is used for module unload notifications.
+struct LDR_DLL_UNLOADED_NOTIFICATION_DATA {
+  // Reserved.
+  ULONG Flags;
+  // The full path name of the DLL module.
+  PCUNICODE_STRING FullDllName;
+  // The base file name of the DLL module.
+  PCUNICODE_STRING BaseDllName;
+  // A pointer to the base address for the DLL in memory.
+  PVOID DllBase;
+  // The size of the DLL image, in bytes.
+  ULONG SizeOfImage;
+};
+using PLDR_DLL_UNLOADED_NOTIFICATION_DATA = LDR_DLL_UNLOADED_NOTIFICATION_DATA*;
+
+// Union that is used for notifications.
+union LDR_DLL_NOTIFICATION_DATA {
+  LDR_DLL_LOADED_NOTIFICATION_DATA Loaded;
+  LDR_DLL_UNLOADED_NOTIFICATION_DATA Unloaded;
+};
+using PLDR_DLL_NOTIFICATION_DATA = LDR_DLL_NOTIFICATION_DATA*;
+
+// Signature of the notification callback function.
+using PLDR_DLL_NOTIFICATION_FUNCTION =
+    VOID(CALLBACK*)(ULONG notification_reason,
+                    const LDR_DLL_NOTIFICATION_DATA* notification_data,
+                    PVOID context);
+
+// Signatures of the functions used for registering DLL notification callbacks.
+using LdrRegisterDllNotificationFunc =
+    NTSTATUS(NTAPI*)(ULONG flags,
+                     PLDR_DLL_NOTIFICATION_FUNCTION notification_function,
+                     PVOID context,
+                     PVOID* cookie);
+using LdrUnregisterDllNotificationFunc = NTSTATUS(NTAPI*)(PVOID cookie);
+
+namespace {
+
+// Global lock for ensuring synchronization of destruction and notifications.
+base::LazyInstance<base::Lock>::Leaky g_module_watcher_lock =

[grt (UTC plus 2), 2016/11/14 20:30:31]

#include "base/synchronization/lock.h"

[chrisha, 2016/11/14 21:38:39]

Done.

+    LAZY_INSTANCE_INITIALIZER;
+// Global pointer to the singleton ModuleWatcher, if one exists. Under
+// |module_watcher_lock|.
+ModuleWatcher* g_module_watcher_instance = nullptr;
+
+// Names of the DLL notification registration functions. These are exported by
+// ntdll.
+constexpr wchar_t kNtDll[] = L"ntdll.dll";
+constexpr char kLdrRegisterDllNotification[] = "LdrRegisterDllNotification";
+constexpr char kLdrUnregisterDllNotification[] = "LdrUnregisterDllNotification";
+
+// Helper function for converting a UNICODE_STRING to a UTF8 std::string.
+std::string ToString(const UNICODE_STRING* str) {

[grt (UTC plus 2), 2016/11/14 20:30:31]

#include <string>

[chrisha, 2016/11/14 21:38:39]

Done.

+  std::string s;
+  base::WideToUTF8(str->Buffer, str->Length / sizeof(wchar_t), &s);
+  return s;
+}
+
+template <typename NotificationDataType>
+void OnModuleEvent(mojom::ModuleEventType event_type,
+                   const NotificationDataType& notification_data,
+                   const ModuleWatcher::OnModuleEventCallback& callback) {
+  mojom::ModuleEvent event;
+  event.event_type = event_type;
+  event.module_path = ToString(notification_data.FullDllName);
+  event.load_address = reinterpret_cast<uintptr_t>(notification_data.DllBase);
+  event.size = notification_data.SizeOfImage;
+  callback.Run(event);
+}
+
+}  // namespace
+
+// static
+std::unique_ptr<ModuleWatcher> ModuleWatcher::Create(
+    const OnModuleEventCallback& callback) {
+  // If a ModuleWatcher already exists then bail out.
+  base::AutoLock lock(g_module_watcher_lock.Get());
+  if (g_module_watcher_instance)
+    return nullptr;
+
+  // This thread acquired the right to create a ModuleWatcher, so do so.
+  g_module_watcher_instance = new ModuleWatcher(callback);
+  return base::WrapUnique(g_module_watcher_instance);
+}
+
+ModuleWatcher::~ModuleWatcher() {
+  // As soon as |g_module_watcher_instance| is null any dispatched callbacks
+  // will be silently absorbed by LoaderNotificationCallback.
+  base::AutoLock lock(g_module_watcher_lock.Get());
+  DCHECK_EQ(g_module_watcher_instance, this);
+  g_module_watcher_instance = nullptr;
+  UnregisterDllNotificationCallback();
+}
+
+void ModuleWatcher::RegisterDllNotificationCallback() {
+  LdrRegisterDllNotificationFunc reg_fn =
+      reinterpret_cast<LdrRegisterDllNotificationFunc>(::GetProcAddress(
+          ::GetModuleHandle(kNtDll), kLdrRegisterDllNotification));
+  if (!reg_fn)

[grt (UTC plus 2), 2016/11/14 20:30:31]

nit: without the return value, i think this reads better as:
  if (reg_fn)
    reg_fn(...);
but it's personal preference so whatever works for you

[chrisha, 2016/11/14 21:38:39]

Done.

+    return;
+
+  reg_fn(0, &LoaderNotificationCallback, this, &dll_notification_cookie_);
+}
+
+void ModuleWatcher::UnregisterDllNotificationCallback() {
+  LdrUnregisterDllNotificationFunc unreg_fn =
+      reinterpret_cast<LdrUnregisterDllNotificationFunc>(::GetProcAddress(
+          ::GetModuleHandle(kNtDll), kLdrUnregisterDllNotification));
+  if (!unreg_fn)
+    return;
+
+  unreg_fn(dll_notification_cookie_);
+}
+
+void ModuleWatcher::EnumerateAlreadyLoadedModules() {
+  // Get all modules in the current process. According to MSDN,
+  // CreateToolhelp32Snapshot should be retried as long as its returning
+  // ERROR_BAD_LENGTH. To avoid locking up here a retry limit is enforced.
+  base::win::ScopedHandle snap;
+  DWORD process_id = ::GetCurrentProcessId();
+  for (size_t i = 0; i < 5; ++i) {
+    snap.Set(::CreateToolhelp32Snapshot(TH32CS_SNAPMODULE | TH32CS_SNAPMODULE32,
+                                        process_id));
+    if (snap.IsValid())
+      break;
+    if (::GetLastError() != ERROR_BAD_LENGTH)
+      return;
+  }
+  if (!snap.IsValid())
+    return;
+
+  // Walk the module list.
+  MODULEENTRY32 module = {sizeof(module)};
+  for (BOOL result = ::Module32First(snap.Get(), &module); result != FALSE;
+       result = ::Module32Next(snap.Get(), &module)) {
+    std::string s;

[grt (UTC plus 2), 2016/11/14 20:30:31]

move this outside of the loop to save on heap churn

[chrisha, 2016/11/14 21:38:39]

Done.

+    base::WideToUTF8(module.szExePath, ::wcslen(module.szExePath), &s);
+
+    mojom::ModuleEvent event;
+    event.event_type = mojom::ModuleEventType::MODULE_ALREADY_LOADED;
+    event.module_path = s;
+    event.load_address = reinterpret_cast<uintptr_t>(module.modBaseAddr);
+    event.size = module.modBaseSize;
+
+    callback_.Run(event);
+  }
+
+  return;
+}
+
+// static
+ModuleWatcher::OnModuleEventCallback ModuleWatcher::GetCallbackForContext(
+    void* context) {
+  base::AutoLock lock(g_module_watcher_lock.Get());
+  if (context != g_module_watcher_instance)
+    return ModuleWatcher::OnModuleEventCallback();

[grt (UTC plus 2), 2016/11/14 20:30:31]

omit "ModuleWatcher::"

[chrisha, 2016/11/14 21:38:39]

Done.

+  return g_module_watcher_instance->callback_;
+}
+
+// static
+void __stdcall ModuleWatcher::LoaderNotificationCallback(
+    unsigned long notification_reason,
+    const LDR_DLL_NOTIFICATION_DATA* notification_data,
+    void* context) {
+  auto callback = GetCallbackForContext(context);
+  if (callback.is_null())
+    return;
+
+  switch (notification_reason) {
+    case LDR_DLL_NOTIFICATION_REASON_LOADED:
+      OnModuleEvent(mojom::ModuleEventType::MODULE_LOADED,
+                    notification_data->Loaded, callback);
+      break;
+
+    case LDR_DLL_NOTIFICATION_REASON_UNLOADED:
+      OnModuleEvent(mojom::ModuleEventType::MODULE_UNLOADED,
+                    notification_data->Unloaded, callback);
+      break;
+
+    default:
+      // This is unexpected, but not a reason to crash.
+      NOTREACHED() << "Unknown LDR_DLL_NOTIFICATION_REASON: "
+                   << notification_reason;
+  }
+}
+
+ModuleWatcher::ModuleWatcher(const OnModuleEventCallback& callback)
+    : callback_(callback) {
+  RegisterDllNotificationCallback();
+  EnumerateAlreadyLoadedModules();
+}