Make sure that base::MessagePumpForUI from different modules are isolated from each other and add protection from shatter attacks by placing |this| pointer to the used data associated with the message-only window (instead of blindly trusting the value of WPARAM).

BUG=124091


Review URL: http://codereview.chromium.org/10134001

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@133134 0039d316-1c4b-4281-b951-d872f2087c98

4 files changed, 78 insertions, 29 deletions

diff --git a/base/message_pump_win.cc b/base/message_pump_win.cc
index 9484b29..72657dc 100644
--- a/base/message_pump_win.cc
+++ b/base/message_pump_win.cc
@@ -8,11 +8,20 @@
 
 #include "base/message_loop.h"
 #include "base/metrics/histogram.h"
+#include "base/process_util.h"
+#include "base/stringprintf.h"
 #include "base/win/wrapped_window_proc.h"
 
+namespace {
+
+// The ID of the timer used by the UI message pump.
+const int kMessagePumpTimerId = 0;
+
+}  // namespace
+
 namespace base {
 
-static const wchar_t kWndClass[] = L"Chrome_MessagePumpWindow";
+static const wchar_t kWndClassFormat[] = L"Chrome_MessagePumpWindow%p";
 
 // Message sent to get an additional time slice for pumping (processing) another
 // task (a series of such messages creates a continuous task pump).
@@ -82,13 +91,19 @@ int MessagePumpWin::GetCurrentDelay() const {
 //-----------------------------------------------------------------------------
 // MessagePumpForUI public:
 
-MessagePumpForUI::MessagePumpForUI() {
+MessagePumpForUI::MessagePumpForUI()
+    : atom_(0),
+      instance_(NULL),
+      message_hwnd_(NULL) {
   InitMessageWnd();
 }
 
 MessagePumpForUI::~MessagePumpForUI() {
-  DestroyWindow(message_hwnd_);
-  UnregisterClass(kWndClass, GetModuleHandle(NULL));
+  if (message_hwnd_ != NULL)
+    DestroyWindow(message_hwnd_);
+
+  if (atom_ != 0)
+    UnregisterClass(reinterpret_cast<const char16*>(atom_), instance_);
 }
 
 void MessagePumpForUI::ScheduleWork() {
@@ -96,7 +111,7 @@ void MessagePumpForUI::ScheduleWork() {
     return;  // Someone else continued the pumping.
 
   // Make sure the MessagePump does some work for us.
-  PostMessage(message_hwnd_, kMsgHaveWork, reinterpret_cast<WPARAM>(this), 0);
+  PostMessage(message_hwnd_, kMsgHaveWork, 0, 0);
 }
 
 void MessagePumpForUI::ScheduleDelayedWork(const TimeTicks& delayed_work_time) {
@@ -129,7 +144,7 @@ void MessagePumpForUI::ScheduleDelayedWork(const TimeTicks& delayed_work_time) {
 
   // Create a WM_TIMER event that will wake us up to check for any pending
   // timers (in case we are running within a nested, external sub-pump).
-  SetTimer(message_hwnd_, reinterpret_cast<UINT_PTR>(this), delay_msec, NULL);
+  SetTimer(message_hwnd_, kMessagePumpTimerId, delay_msec, NULL);
 }
 
 void MessagePumpForUI::PumpOutPendingPaintMessages() {
@@ -162,13 +177,19 @@ void MessagePumpForUI::PumpOutPendingPaintMessages() {
 // static
 LRESULT CALLBACK MessagePumpForUI::WndProcThunk(
     HWND hwnd, UINT message, WPARAM wparam, LPARAM lparam) {
-  switch (message) {
-    case kMsgHaveWork:
-      reinterpret_cast<MessagePumpForUI*>(wparam)->HandleWorkMessage();
-      break;
-    case WM_TIMER:
-      reinterpret_cast<MessagePumpForUI*>(wparam)->HandleTimerMessage();
-      break;
+  // Retrieve |this| from the user data, associated with the window.
+  MessagePumpForUI* self = reinterpret_cast<MessagePumpForUI*>(
+      GetWindowLongPtr(hwnd, GWLP_USERDATA));
+  if (self != NULL) {
+    switch (message) {
+      case kMsgHaveWork:
+        self->HandleWorkMessage();
+        break;
+      case WM_TIMER:
+        DCHECK(wparam == kMessagePumpTimerId);
+        self->HandleTimerMessage();
+        break;
+    }
   }
   return DefWindowProc(hwnd, message, wparam, lparam);
 }
@@ -211,7 +232,7 @@ void MessagePumpForUI::DoRunLoop() {
     // don't want to disturb that timer if it is already in flight.  However,
     // if we did do all remaining delayed work, then lets kill the WM_TIMER.
     if (more_work_is_plausible && delayed_work_time_.is_null())
-      KillTimer(message_hwnd_, reinterpret_cast<UINT_PTR>(this));
+      KillTimer(message_hwnd_, kMessagePumpTimerId);
     if (state_->should_quit)
       break;
 
@@ -230,18 +251,38 @@ void MessagePumpForUI::DoRunLoop() {
 }
 
 void MessagePumpForUI::InitMessageWnd() {
-  HINSTANCE hinst = GetModuleHandle(NULL);
+  // Register a unique window class for each instance of UI pump.
+  string16 class_name = base::StringPrintf(kWndClassFormat, this);
+  WNDPROC window_procedure = &base::win::WrappedWindowProc<WndProcThunk>;
 
-  WNDCLASSEX wc = {0};
+  // RegisterClassEx uses a handle of the module containing the window procedure
+  // to distinguish identically named classes registered in different modules.
+  instance_ = GetModuleFromAddress(window_procedure);
+
+  WNDCLASSEXW wc = {0};
   wc.cbSize = sizeof(wc);
-  wc.lpfnWndProc = base::win::WrappedWindowProc<WndProcThunk>;
-  wc.hInstance = hinst;
-  wc.lpszClassName = kWndClass;
-  RegisterClassEx(&wc);
+  wc.lpfnWndProc = window_procedure;
+  wc.hInstance = instance_;
+  wc.lpszClassName = class_name.c_str();
+  atom_ = RegisterClassEx(&wc);
+  if (atom_ == 0) {
+    DCHECK(atom_);
+    return;
+  }
+
+  // Create the message-only window.
+  message_hwnd_ = CreateWindow(
+      reinterpret_cast<const char16*>(atom_), 0, 0, 0, 0, 0, 0, HWND_MESSAGE, 0,
+      instance_, 0);
+  if (message_hwnd_ == NULL) {
+    DCHECK(message_hwnd_);
+    return;
+  }
 
-  message_hwnd_ =
-      CreateWindow(kWndClass, 0, 0, 0, 0, 0, 0, HWND_MESSAGE, 0, hinst, 0);
-  DCHECK(message_hwnd_);
+  // Store |this| so that the window procedure could retrieve it later.
+  SetWindowLongPtr(message_hwnd_,
+                   GWLP_USERDATA,
+                   reinterpret_cast<LONG_PTR>(this));
 }
 
 void MessagePumpForUI::WaitForWork() {
@@ -300,7 +341,7 @@ void MessagePumpForUI::HandleWorkMessage() {
 }
 
 void MessagePumpForUI::HandleTimerMessage() {
-  KillTimer(message_hwnd_, reinterpret_cast<UINT_PTR>(this));
+  KillTimer(message_hwnd_, kMessagePumpTimerId);
 
   // If we are being called outside of the context of Run, then don't do
   // anything.  This could correspond to a MessageBox call or something of
diff --git a/base/message_pump_win.h b/base/message_pump_win.h
index f5a00f3..7514bae 100644
--- a/base/message_pump_win.h
+++ b/base/message_pump_win.h
@@ -154,6 +154,12 @@ class BASE_EXPORT MessagePumpForUI : public MessagePumpWin {
   bool ProcessMessageHelper(const MSG& msg);
   bool ProcessPumpReplacementMessage();
 
+  // Atom representing the message-only window class.
+  ATOM atom_;
+
+  // Instance of the module containing the window procedure.
+  HMODULE instance_;
+
   // A hidden message-only window.
   HWND message_hwnd_;
 };
diff --git a/base/process_util.h b/base/process_util.h
index 1aab778..f2e53cb 100644
--- a/base/process_util.h
+++ b/base/process_util.h
@@ -146,7 +146,8 @@ BASE_EXPORT ProcessId GetCurrentProcId();
 BASE_EXPORT ProcessHandle GetCurrentProcessHandle();
 
 #if defined(OS_WIN)
-// Returns the module handle to which an address belongs.
+// Returns the module handle to which an address belongs. The reference counter
+// of the module is not incremented.
 BASE_EXPORT HMODULE GetModuleFromAddress(void* address);
 #endif
 
diff --git a/base/process_util_win.cc b/base/process_util_win.cc
index 4fab4c6..9b2e100 100644
--- a/base/process_util_win.cc
+++ b/base/process_util_win.cc
@@ -176,13 +176,14 @@ ProcessHandle GetCurrentProcessHandle() {
 }
 
 HMODULE GetModuleFromAddress(void* address) {
-  HMODULE hinst = NULL;
-  if (!::GetModuleHandleExA(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS,
+  HMODULE instance = NULL;
+  if (!::GetModuleHandleExA(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS |
+                            GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
                             static_cast<char*>(address),
-                            &hinst)) {
+                            &instance)) {
     NOTREACHED();
   }
-  return hinst;
+  return instance;
 }
 
 bool OpenProcessHandle(ProcessId pid, ProcessHandle* handle) {