[Win] Create ModuleWatcher.

chrome/common/conflicts/module_watcher_win.cc

Index: chrome/common/conflicts/module_watcher_win.cc
diff --git a/chrome/common/conflicts/module_watcher_win.cc b/chrome/common/conflicts/module_watcher_win.cc
new file mode 100644
index 0000000000000000000000000000000000000000..edd671cfda4146d7b777c068a5b34c2fdf618e82
--- /dev/null
+++ b/chrome/common/conflicts/module_watcher_win.cc
@@ -0,0 +1,237 @@
+// 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/strings/utf_string_conversions.h"
+#include "base/win/scoped_handle.h"
+
+namespace {
+
+// Global lock for ensuring synchronization of destruction and notifications.
+base::Lock module_watcher_lock;

[grt (UTC plus 2), 2016/11/10 08:52:12]

you need to use a leaky lazy instance here to avoid creating a static dtor:
base::LazyInstance<base::Lock>::Leaky g_module_watcher_lock =
    LAZY_INSTANCE_INITIALIZER;

[grt (UTC plus 2), 2016/11/10 08:52:12]

nit: g_foo here and below

[chrisha, 2016/11/14 16:06:44]

Done.

[chrisha, 2016/11/14 16:06:45]

Done.

+// Global pointer to the singleton ModuleWatcher, if one exists. Under
+// |module_watcher_lock|.
+ModuleWatcher* module_watcher_instance = nullptr;
+
+// 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.
+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.
+typedef 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;
+} LDR_DLL_LOADED_NOTIFICATION_DATA, *PLDR_DLL_LOADED_NOTIFICATION_DATA;
+
+// Structure that is used for module unload notifications.
+typedef 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;
+} LDR_DLL_UNLOADED_NOTIFICATION_DATA, *PLDR_DLL_UNLOADED_NOTIFICATION_DATA;
+
+// Union that is used for notifications.
+typedef union _LDR_DLL_NOTIFICATION_DATA {
+  LDR_DLL_LOADED_NOTIFICATION_DATA Loaded;
+  LDR_DLL_UNLOADED_NOTIFICATION_DATA Unloaded;
+} LDR_DLL_NOTIFICATION_DATA, *PLDR_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);
+
+// 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) {
+  std::string s;
+  base::WideToUTF8(str->Buffer, str->Length / sizeof(wchar_t), &s);
+  return s;
+}
+
+// static

[grt (UTC plus 2), 2016/11/10 08:52:12]

remove

[chrisha, 2016/11/14 16:06:45]

Done.

+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(module_watcher_lock);
+  if (module_watcher_instance)
+    return nullptr;
+
+  // This thread acquired the right to create a ModuleWatcher, so do so.
+  module_watcher_instance = new ModuleWatcher(callback);
+  return std::unique_ptr<ModuleWatcher>(module_watcher_instance);

[grt (UTC plus 2), 2016/11/10 08:52:12]

nit:
  return base::WrapUnique<ModuleWatcher>(module_watcher_instance);

[chrisha, 2016/11/14 16:06:45]

Done.

+}
+
+ModuleWatcher::ModuleWatcher(const OnModuleEventCallback& callback)

[grt (UTC plus 2), 2016/11/10 08:52:12]

nit: move down below LoaderNotificationCallback to maintain ordering of
declarations in .h file

[chrisha, 2016/11/14 16:06:45]

Done.

+    : callback_(callback) {
+  RegisterDllNotificationCallback();
+  EnumerateAlreadyLoadedModules();
+}
+
+ModuleWatcher::~ModuleWatcher() {
+  // As soon as |module_watcher_instance| is null any dispatched callbacks
+  // will be silently absorbed by LoaderNotificationCallback.
+  base::AutoLock lock(module_watcher_lock);
+  DCHECK_EQ(module_watcher_instance, this);
+  module_watcher_instance = nullptr;
+  UnregisterDllNotificationCallback();
+}
+
+bool ModuleWatcher::RegisterDllNotificationCallback() {
+  LdrRegisterDllNotificationFunc reg_fn =
+      reinterpret_cast<LdrRegisterDllNotificationFunc>(::GetProcAddress(
+          ::GetModuleHandle(kNtDll), kLdrRegisterDllNotification));
+
+  if (!reg_fn)
+    return false;
+
+  NTSTATUS status = reg_fn(0, reinterpret_cast<PLDR_DLL_NOTIFICATION_FUNCTION>(
+                                  &ModuleWatcher::LoaderNotificationCallback),

[grt (UTC plus 2), 2016/11/10 08:52:12]

nit: omit "ModuleWatcher::"

[chrisha, 2016/11/14 16:06:45]

Done.

+                           this, &dll_notification_cookie_);
+  return status == 0;
+}
+
+bool ModuleWatcher::UnregisterDllNotificationCallback() {
+  LdrUnregisterDllNotificationFunc unreg_fn =
+      reinterpret_cast<LdrUnregisterDllNotificationFunc>(::GetProcAddress(
+          ::GetModuleHandle(kNtDll), kLdrUnregisterDllNotification));
+
+  if (!unreg_fn)
+    return false;
+
+  NTSTATUS status = unreg_fn(dll_notification_cookie_);
+  return status == 0;
+}
+
+bool 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;
+  for (size_t i = 0; i < 5; ++i) {
+    snap.Set(::CreateToolhelp32Snapshot(TH32CS_SNAPMODULE | TH32CS_SNAPMODULE32,
+                                        ::GetCurrentProcessId()));

[grt (UTC plus 2), 2016/11/09 08:37:52]

is GetCurrentProcessId free? if no, move outside of loop

[chrisha, 2016/11/14 16:06:45]

Done.

+    if (snap.IsValid())
+      break;
+    if (::GetLastError() == ERROR_BAD_LENGTH)

[grt (UTC plus 2), 2016/11/09 08:37:52]

personal style: i'd invert the logic here and do away with the continue
statement.

[chrisha, 2016/11/14 16:06:44]

Done.

+      continue;
+    return false;
+  }
+  if (!snap.IsValid())
+    return false;
+
+  // Walk the module list.
+  MODULEENTRY32 module = {sizeof(module)};
+  if (!::Module32First(snap.Get(), &module))

[grt (UTC plus 2), 2016/11/09 08:37:52]

both functions are documented as handing success/failure the same way, so you
could switch this to:
  for (BOOL result = ::Module32First(snap.Get(), &module); result != FALSE;
       result = ::Module32Next(snap.Get(), &module)) {
    ...
  }
  return ::GetLastError() == ERROR_NO_MORE_FILES;

then again, the return value of this is ignored, so you could get rid of the
last statement. in fact, the return value is ignored for all uses of these three
member functions. why not simplify and get rid of the bool?

[chrisha, 2016/11/14 16:06:44]

Yeah, I initially was CHECKing on these, but obviously that's a little too
aggressive. I don't expect these to fail, but having that case be DCHECKed
doesn't tell me much either. So removed all the bools.

+    return false;
+
+  do {
+    std::string s;
+    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);
+  } while (::Module32Next(snap.Get(), &module));
+
+  return true;
+}
+
+// static
+void __stdcall ModuleWatcher::LoaderNotificationCallback(
+    uint32_t notification_reason,
+    const void* notification_data,
+    void* context) {
+  auto* data =
+      reinterpret_cast<const LDR_DLL_NOTIFICATION_DATA*>(notification_data);
+
+  // Get a copy of the callback to invoke.
+  ModuleWatcher* module_watcher = reinterpret_cast<ModuleWatcher*>(context);

[grt (UTC plus 2), 2016/11/10 08:52:12]

|module_watcher| shouldn't actually be used as one outside of the lock. to avoid
a future where someone inadvertently adds code to use it, wdyt about introducing
a function in the unnamed namespace along these lines:

ModuleWatcher::OnModuleEventCallback
GetCallbackForLoaderNotificationContext(void* context) {
   acquire the lock, compare the instance, and either return an empty callback
or the instance's
}

and then change this to:

  OnModuleEventCallback callback =
GetCallbackForLoaderNotificationContext(context);
  if (callback.is_null())
    return;

or something along those lines. maybe shrink the name of the fn to avoid
wrapping. :-)

[chrisha, 2016/11/14 16:06:45]

Great idea! Done.

+  ModuleWatcher::OnModuleEventCallback callback;

[grt (UTC plus 2), 2016/11/10 08:52:12]

nit: omit "ModuleWatcher::"

[chrisha, 2016/11/14 16:06:45]

Done.

+  {
+    base::AutoLock lock(module_watcher_lock);
+    if (module_watcher_instance != module_watcher)
+      return;
+    callback = module_watcher->callback_;
+  }
+
+  switch (notification_reason) {
+    case LDR_DLL_NOTIFICATION_REASON_LOADED:
+      OnModuleEvent(mojom::ModuleEventType::MODULE_LOADED, data->Loaded,
+                    callback);
+      break;
+
+    case LDR_DLL_NOTIFICATION_REASON_UNLOADED:
+      OnModuleEvent(mojom::ModuleEventType::MODULE_UNLOADED, data->Unloaded,
+                    callback);
+      break;
+
+    default:
+      // This is unexpected, but not a reason to crash.
+      DCHECK(false) << "Unknown LDR_DLL_NOTIFICATION_REASON: "

[grt (UTC plus 2), 2016/11/09 08:37:52]

NOTREACHED()

[chrisha, 2016/11/09 15:38:07]

Oops, I thought NOTREACHED was a fatal thing. Didn't realize it was effectively
a #define for DCHECK(false).

+                    << notification_reason;
+  }
+}