[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 <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 conflicts {
+
+namespace {
+
+// 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 {

[sky, 2016/11/04 16:27:24]

using for all these?

[chrisha, 2016/11/04 19:15:49]

I'm never sure what the right solution is here. These are actually Windows
definitions, so I've declared them in the Windows style. I'm ambivalent, but
generally prefer to have OS things look like OS things, and Chrome things look
like Chrome things.

+  // 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 wchar_t kNtDll[] = L"ntdll.dll";
+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) {
+  return base::StringPiece16(str->Buffer, str->Length / sizeof(wchar_t));
+}
+
+HMODULE GetNtDll() {
+  HMODULE ntdll = ::GetModuleHandle(L"ntdll.dll");
+  return ntdll;
+}
+
+}  // namespace
+
+// 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);
+};

[sky, 2016/11/04 16:27:24]

DISALLOW_IMPLICIT_CONSTRUCTORS

+
+// 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 = {MODULE_LOADED,
+                       base::FilePath(ToStringPiece(load_data.FullDllName)),
+                       load_data.DllBase, 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 = {MODULE_UNLOADED,
+                       base::FilePath(ToStringPiece(unload_data.FullDllName)),
+                       unload_data.DllBase, unload_data.SizeOfImage};
+  return module_watcher->Notify(event);
+}
+
+ModuleWatcher::ModuleWatcher()
+    : started_(false), dll_notification_cookie_(nullptr) {
+  observer_list_ = new base::ObserverListThreadSafe<Observer>();
+}
+
+ModuleWatcher::~ModuleWatcher() = default;
+
+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() {
+  lock_.AssertAcquired();
+  LdrRegisterDllNotificationFunc reg_fn =
+      reinterpret_cast<LdrRegisterDllNotificationFunc>(::GetProcAddress(
+          ::GetModuleHandle(kNtDll), kLdrRegisterDllNotification));
+
+  if (!reg_fn)
+    return false;
+
+  NTSTATUS status = reg_fn(0, &Bridge::LoaderNotificationCallback, this,
+                           &dll_notification_cookie_);
+  return status == 0;
+}
+
+bool ModuleWatcher::UnregisterDllNotificationCallback() {
+  lock_.AssertAcquired();
+  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.
+  base::win::ScopedHandle snap(
+      ::CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, ::GetCurrentProcessId()));
+  if (!snap.IsValid())
+    return false;
+
+  // Walk the module list.
+  MODULEENTRY32 module = {sizeof(module)};
+  if (!::Module32First(snap.Get(), &module))
+    return false;
+
+  do {
+    ModuleEvent event = {MODULE_ALREADY_LOADED,
+                         base::FilePath(module.szExePath), module.modBaseAddr,
+                         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