[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..855503454c18a21e9f1fcca7c5c2d93f52b8b398
--- /dev/null
+++ b/chrome/common/conflicts/module_watcher_win.cc
@@ -0,0 +1,272 @@
+// 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 <tlhelp32.h>
+#include <windows.h>
+#include <winternl.h>  // For UNICODE_STRING.
+
+#include <memory>
+
+#include "base/observer_list_threadsafe.h"
+#include "base/strings/string_piece.h"
+#include "base/win/scoped_handle.h"
+
+namespace {

[Patrick Monette, 2016/11/03 20:18:15]

Usually, the unnamed namespace is put inside the file namespace (conflicts).

Unless it was causing problems with the typedefs below.

[chrisha, 2016/11/04 15:32:42]

Done.

+
+// 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 otification callback function.
+typedef VOID(CALLBACK* PLDR_DLL_NOTIFICATION_FUNCTION)(
+    ULONG notification_reason,
+    const LDR_DLL_NOTIFICATION_DATA* notification_data,
+    PVOID context);
+
+// Signatures of the functions used for registering DLL notification callbacks.
+typedef NTSTATUS(NTAPI* LdrRegisterDllNotificationFunc)(
+    ULONG flags,
+    PLDR_DLL_NOTIFICATION_FUNCTION notification_function,
+    PVOID context,
+    PVOID* cookie);
+typedef NTSTATUS(NTAPI* LdrUnregisterDllNotificationFunc)(PVOID cookie);
+
+// Names of the DLL notification registration functions. These are exported by
+// ntdll.
+constexpr char kLdrRegisterDllNotification[] = "LdrRegisterDllNotification";
+constexpr char kLdrUnregisterDllNotification[] = "LdrUnregisterDllNotification";
+
+// Helper function for converting a UNICODE_STRING to a string piece.
+base::StringPiece16 ToStringPiece(const UNICODE_STRING* str) {
+  CHECK_NE(static_cast<const UNICODE_STRING*>(NULL), str);

[Patrick Monette, 2016/11/03 20:18:15]

This should probably be a DCHECK. Also s/NULL/nullptr.

[chrisha, 2016/11/04 15:32:42]

(This is cut and pasted code from elsewhere, oops.)

Actually, we have a policy to not explcitly CHECK or DCHECK pointers, as
dereferencing them is the same as advertising the fact that we expect them to be
non-null, and it will deterministically crash anyways.

Removed.

+  return base::StringPiece16(str->Buffer, str->Length / sizeof(wchar_t));
+}
+
+HMODULE GetNtDll() {
+  HMODULE ntdll = ::GetModuleHandle(L"ntdll.dll");
+  CHECK_NE(static_cast<HMODULE>(NULL), ntdll);

[Patrick Monette, 2016/11/03 20:18:15]

Same

[chrisha, 2016/11/04 15:32:42]

Removed this helper function entirely, as the CHECK is redundant.

+  return ntdll;
+}
+
+}  // namespace
+
+namespace conflicts {
+
+// Static helper class for forwarding system notifications to the ModuleWatcher.
+// The notifications here will be called on the thread that is doing the DLL
+// load, which could be on a variety of threads in Chrome, including threads
+// injected by third parties.
+class ModuleWatcher::Bridge {
+ public:
+  // This is invoked by the OS. This forwards to OnModuleLoaded or
+  // OnModuleUnloaded depending on |notification_reason|.
+  static void CALLBACK
+  LoaderNotificationCallback(ULONG notification_reason,
+                             const LDR_DLL_NOTIFICATION_DATA* notification_data,
+                             PVOID context);
+
+  // Invoked by LoaderNotification for LDR_DLL_NOTIFICATION_REASON_LOADED
+  // events.
+  static void OnModuleLoad(const LDR_DLL_LOADED_NOTIFICATION_DATA& load_data,
+                           ModuleWatcher* module_watcher);
+
+  // Invoked by LoaderNotification for LDR_DLL_NOTIFICATION_REASON_UNLOADED
+  // events.
+  static void OnModuleUnload(
+      const LDR_DLL_UNLOADED_NOTIFICATION_DATA& unload_data,
+      ModuleWatcher* module_watcher);
+};
+
+// static
+void ModuleWatcher::Bridge::LoaderNotificationCallback(
+    ULONG notification_reason,
+    const LDR_DLL_NOTIFICATION_DATA* notification_data,
+    PVOID context) {
+  ModuleWatcher* module_watcher = reinterpret_cast<ModuleWatcher*>(context);
+  switch (notification_reason) {
+    case LDR_DLL_NOTIFICATION_REASON_LOADED:
+      return OnModuleLoad(notification_data->Loaded, module_watcher);
+
+    case LDR_DLL_NOTIFICATION_REASON_UNLOADED:
+      return OnModuleUnload(notification_data->Unloaded, module_watcher);
+
+    default:
+      NOTREACHED() << "Unknown LDR_DLL_NOTIFICATION_REASON: "
+                   << notification_reason;
+  }
+}
+
+// static
+void ModuleWatcher::Bridge::OnModuleLoad(
+    const LDR_DLL_LOADED_NOTIFICATION_DATA& load_data,
+    ModuleWatcher* module_watcher) {
+  ModuleEvent event = {};

[Patrick Monette, 2016/11/03 20:18:15]

All 4 members of ModuleEvent are always initialized right after creation.

Add a constructor or use uniform initialization syntax:
ModuleEvent event = {
  MODULE_LOADED,
  base::FilePath(ToStringPiece(load_data.FullDllName)),
  load_data.DllBase,
  load_data.SizeOfImage,
};

[chrisha, 2016/11/04 15:32:42]

Done.

+  event.event_type = MODULE_LOADED;
+  event.module_path = base::FilePath(ToStringPiece(load_data.FullDllName));
+  event.load_address = load_data.DllBase;
+  event.size = load_data.SizeOfImage;
+  return module_watcher->Notify(event);
+}
+
+// static
+void ModuleWatcher::Bridge::OnModuleUnload(
+    const LDR_DLL_UNLOADED_NOTIFICATION_DATA& unload_data,
+    ModuleWatcher* module_watcher) {
+  ModuleEvent event = {};
+  event.event_type = MODULE_UNLOADED;
+  event.module_path = base::FilePath(ToStringPiece(unload_data.FullDllName));
+  event.load_address = unload_data.DllBase;
+  event.size = unload_data.SizeOfImage;
+  return module_watcher->Notify(event);
+}
+
+ModuleWatcher::ModuleWatcher()

[Patrick Monette, 2016/11/03 20:18:15]

I looked at the design doc. Won't each ModuleWatcher only have one observer?

If this is the case the logic could be simplified.

[chrisha, 2016/11/04 15:32:42]

Yeah, but I need the thread affinity, as I want the notifications on the file
thread (not mentioned in the design doc... I'll go update that).

Rather than do the thread bouncing myself I figured it was easier to use
ObserverListThreadSafe. Either way I write it, there's some amount of
complexity.

+    : started_(false), dll_notification_cookie_(nullptr) {
+  observer_list_ = new base::ObserverListThreadSafe<Observer>();
+}
+
+ModuleWatcher::~ModuleWatcher() {}

[Patrick Monette, 2016/11/03 20:18:15]

= default;

[chrisha, 2016/11/04 15:32:42]

Always forgetting about that new feature...

Done

+
+std::unique_ptr<ModuleWatcher::Observer> ModuleWatcher::RegisterCallback(
+    const OnModuleEventCallback& callback) {
+  // Create a new Observer. This will register itself with the |observer_list_|.
+  return std::unique_ptr<Observer>(new Observer(callback, this));
+}
+
+void ModuleWatcher::Start() {
+  base::AutoLock lock(lock_);
+  DCHECK(!started_);
+  if (!RegisterDllNotificationCallback())
+    return;
+  EnumerateAlreadyLoadedModules();
+  started_ = true;
+}
+
+void ModuleWatcher::Stop() {
+  base::AutoLock lock(lock_);
+  DCHECK(started_);
+  UnregisterDllNotificationCallback();
+}
+
+bool ModuleWatcher::IsRunning() {
+  base::AutoLock lock(lock_);
+  return started_;
+}
+
+bool ModuleWatcher::RegisterDllNotificationCallback() {

[Patrick Monette, 2016/11/03 20:18:15]

lock_.AssertAquired() in all relevant functions?

[chrisha, 2016/11/04 15:32:42]

Done.

+  LdrRegisterDllNotificationFunc reg_fn =
+      reinterpret_cast<LdrRegisterDllNotificationFunc>(
+          ::GetProcAddress(GetNtDll(), kLdrRegisterDllNotification));
+
+  if (reg_fn == NULL)

[Patrick Monette, 2016/11/03 20:18:15]

if (!reg_fn)

[chrisha, 2016/11/04 15:32:42]

Done.

+    return false;
+
+  NTSTATUS status = reg_fn(0, &Bridge::LoaderNotificationCallback, this,
+                           &dll_notification_cookie_);
+  return status == 0;
+}
+
+bool ModuleWatcher::UnregisterDllNotificationCallback() {

[Patrick Monette, 2016/11/03 20:18:15]

The return value is unused. Not sure if it's worth removing though.

[chrisha, 2016/11/04 15:32:42]

This is a thin wrapper around the system function, so I wanted to preserve the
original semantics. This apparently can fail, although there's no meaningful
action for us to take in that case. I'm ambivalent.

+  LdrUnregisterDllNotificationFunc unreg_fn =
+      reinterpret_cast<LdrUnregisterDllNotificationFunc>(
+          ::GetProcAddress(GetNtDll(), kLdrUnregisterDllNotification));
+
+  if (unreg_fn == NULL)

[Patrick Monette, 2016/11/03 20:18:15]

if (!unreg_fn)

[chrisha, 2016/11/04 15:32:42]

Done.

+    return false;
+
+  NTSTATUS status = unreg_fn(dll_notification_cookie_);
+  return status == 0;
+}
+
+bool ModuleWatcher::EnumerateAlreadyLoadedModules() {
+  // Get all modules in the current process.
+  base::win::ScopedHandle snap(
+      ::CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, ::GetCurrentProcessId()));
+  if (!snap.Get())

[Patrick Monette, 2016/11/03 20:18:15]

if (!snap.IsValid())

[chrisha, 2016/11/04 15:32:42]

Done.

+    return false;
+
+  // Walk the module list.
+  MODULEENTRY32 module = {sizeof(module)};
+  if (!::Module32First(snap.Get(), &module))
+    return false;
+
+  do {
+    ModuleEvent event = {};
+    event.event_type = MODULE_ALREADY_LOADED;
+    event.module_path = base::FilePath(module.szExePath);
+    event.load_address = module.modBaseAddr;
+    event.size = module.modBaseSize;
+    Notify(event);
+  } while (::Module32Next(snap.Get(), &module));
+
+  return true;
+}
+
+void ModuleWatcher::Notify(const ModuleEvent& event) {
+  // Forward the notification to each registered observer.
+  observer_list_->Notify(FROM_HERE, &Observer::Notify, event);
+}
+
+ModuleWatcher::Observer::Observer(const OnModuleEventCallback& callback,
+                                  ModuleWatcher* module_watcher)
+    : callback_(callback), module_watcher_(module_watcher) {
+  module_watcher_->observer_list_->AddObserver(this);
+}
+
+ModuleWatcher::Observer::~Observer() {
+  module_watcher_->observer_list_->RemoveObserver(this);
+}
+
+void ModuleWatcher::Observer::Notify(const ModuleEvent& event) {
+  // Forward the observer to the callback. This all happens on the same thread
+  // on which the Observer was created.
+  callback_.Run(event);
+}
+
+}  // namespace conflicts