Skip to content

Commit

Permalink
Refactoring Windows Hooks in IE driver
Browse files Browse the repository at this point in the history
Windows hooks are now wrapped in a single class, eliminating a fair amount
of duplicate code.
  • Loading branch information
jimevans committed Jun 18, 2015
1 parent 0c4324b commit 840a86f
Show file tree
Hide file tree
Showing 10 changed files with 332 additions and 215 deletions.
89 changes: 32 additions & 57 deletions cpp/iedriver/CommandHandlers/ScreenshotCommandHandler.h
Original file line number Diff line number Diff line change
Expand Up @@ -18,23 +18,13 @@
#define WEBDRIVER_IE_SCREENSHOTCOMMANDHANDLER_H_

#include "../Browser.h"
#include "../HookProcessor.h"
#include "../IECommandHandler.h"
#include "../IECommandExecutor.h"
#include "logging.h"
#include <atlimage.h>
#include <atlenc.h>

// Define a shared data segment. Variables in this segment can be
// shared across processes that load this DLL.
#pragma data_seg("SHARED")
HHOOK next_hook = NULL;
HWND ie_window_handle = NULL;
int max_width = 0;
int max_height = 0;
#pragma data_seg()

#pragma comment(linker, "/section:SHARED,RWS")

namespace webdriver {

class ScreenshotCommandHandler : public IECommandHandler {
Expand Down Expand Up @@ -111,7 +101,7 @@ class ScreenshotCommandHandler : public IECommandHandler {
HRESULT CaptureBrowser(BrowserHandle browser) {
LOG(TRACE) << "Entering ScreenshotCommandHandler::CaptureBrowser";

ie_window_handle = browser->GetTopLevelWindowHandle();
HWND ie_window_handle = browser->GetTopLevelWindowHandle();
HWND content_window_handle = browser->GetContentWindowHandle();

CComPtr<IHTMLDocument2> document;
Expand Down Expand Up @@ -139,24 +129,31 @@ class ScreenshotCommandHandler : public IECommandHandler {
LOG(DEBUG) << "Initial chrome sizes are (w, h): "
<< chrome_width << ", " << chrome_height;

max_width = document_info.width + chrome_width;
max_height = document_info.height + chrome_height;
// Technically, we could use a custom structure here, and save a
// few bytes, but RECT is already a well-known structure, and
// one that is included as part of the Windows SDK, so we'll
// leverage it.
RECT max_image_dimensions;
max_image_dimensions.left = 0;
max_image_dimensions.top = 0;
max_image_dimensions.right = document_info.width + chrome_width;
max_image_dimensions.bottom = document_info.height + chrome_height;

// For some reason, this technique does not allow the user to resize
// the browser window to greater than SIZE_LIMIT x SIZE_LIMIT. This is pretty
// big, so we'll cap the allowable screenshot size to that.
//
// GDI+ limit after which it may report Generic error for some image types
int SIZE_LIMIT = 65534;
if (max_height > SIZE_LIMIT) {
if (max_image_dimensions.bottom > SIZE_LIMIT) {
LOG(WARN) << "Required height is greater than limit. Truncating screenshot height.";
max_height = SIZE_LIMIT;
document_info.height = max_height - chrome_height;
max_image_dimensions.bottom = SIZE_LIMIT;
document_info.height = max_image_dimensions.bottom - chrome_height;
}
if (max_width > SIZE_LIMIT) {
if (max_image_dimensions.right > SIZE_LIMIT) {
LOG(WARN) << "Required width is greater than limit. Truncating screenshot width.";
max_width = SIZE_LIMIT;
document_info.width = max_width - chrome_width;
max_image_dimensions.right = SIZE_LIMIT;
document_info.width = max_image_dimensions.right - chrome_width;
}

long original_width = browser->GetWidth();
Expand All @@ -175,10 +172,11 @@ class ScreenshotCommandHandler : public IECommandHandler {
::ShowWindow(ie_window_handle, SW_SHOWNORMAL);
}

this->InstallWindowsHook();

browser->SetWidth(max_width);
browser->SetHeight(max_height);
HookProcessor hook(ie_window_handle);
hook.InstallWindowsHook("ScreenshotWndProc", WH_CALLWNDPROC);
hook.PushData(sizeof(max_image_dimensions), &max_image_dimensions);
browser->SetWidth(max_image_dimensions.right);
browser->SetHeight(max_image_dimensions.bottom);

// Capture the window's canvas to a DIB.
BOOL created = this->image_->Create(document_info.width,
Expand All @@ -196,7 +194,7 @@ class ScreenshotCommandHandler : public IECommandHandler {
LOG(WARN) << "PrintWindow API is not able to get content window screenshot";
}

this->UninstallWindowsHook();
hook.UninstallWindowsHook();

// Restore the browser to the original dimensions.
if (is_maximized) {
Expand Down Expand Up @@ -330,33 +328,6 @@ class ScreenshotCommandHandler : public IECommandHandler {
*width = window_rect.right - window_rect.left;
*height = window_rect.bottom - window_rect.top;
}

void InstallWindowsHook() {
LOG(TRACE) << "Entering ScreenshotCommandHandler::InstallWindowsHook";

HINSTANCE instance_handle = _AtlBaseModule.GetModuleInstance();

FARPROC hook_procedure_address = ::GetProcAddress(instance_handle, "ScreenshotWndProc");
if (hook_procedure_address == NULL || hook_procedure_address == 0) {
LOGERR(WARN) << "Unable to get address of hook procedure to catch WM_GETMINMAXINFO";
return;
}
HOOKPROC hook_procedure = reinterpret_cast<HOOKPROC>(hook_procedure_address);

// Install the Windows hook.
DWORD thread_id = ::GetWindowThreadProcessId(ie_window_handle, NULL);
next_hook = ::SetWindowsHookEx(WH_CALLWNDPROC,
hook_procedure,
instance_handle,
thread_id);
if (next_hook == NULL) {
LOGERR(WARN) << "Unable to set windows hook to catch WM_GETMINMAXINFO";
}
}

void UninstallWindowsHook() {
::UnhookWindowsHookEx(next_hook);
}
};

} // namespace webdriver
Expand Down Expand Up @@ -387,9 +358,10 @@ LRESULT CALLBACK MinMaxInfoHandler(HWND hwnd,

if (WM_GETMINMAXINFO == message) {
MINMAXINFO* minMaxInfo = reinterpret_cast<MINMAXINFO*>(lParam);

minMaxInfo->ptMaxTrackSize.x = max_width;
minMaxInfo->ptMaxTrackSize.y = max_height;
RECT max_size;
webdriver::HookProcessor::CopyDataFromBuffer(sizeof(RECT), reinterpret_cast<void*>(&max_size));
minMaxInfo->ptMaxTrackSize.x = max_size.right;
minMaxInfo->ptMaxTrackSize.y = max_size.bottom;

// We're not going to pass this message onto the original message
// processor, so we should return 0, per the documentation for
Expand Down Expand Up @@ -417,7 +389,10 @@ LRESULT CALLBACK MinMaxInfoHandler(HWND hwnd,
// See the discussion here: http://www.codeguru.com/forum/showthread.php?p=1889928
LRESULT CALLBACK ScreenshotWndProc(int nCode, WPARAM wParam, LPARAM lParam) {
CWPSTRUCT* call_window_proc_struct = reinterpret_cast<CWPSTRUCT*>(lParam);
if (WM_GETMINMAXINFO == call_window_proc_struct->message) {
if (WM_COPYDATA == call_window_proc_struct->message) {
COPYDATASTRUCT* data = reinterpret_cast<COPYDATASTRUCT*>(call_window_proc_struct->lParam);
webdriver::HookProcessor::CopyDataToBuffer(data->cbData, data->lpData);
} else if (WM_GETMINMAXINFO == call_window_proc_struct->message) {
// Inject our own message processor into the process so we can modify
// the WM_GETMINMAXINFO message. It is not possible to modify the
// message from this hook, so the best we can do is inject a function
Expand All @@ -430,7 +405,7 @@ LRESULT CALLBACK ScreenshotWndProc(int nCode, WPARAM wParam, LPARAM lParam) {
reinterpret_cast<HANDLE>(proc));
}

return ::CallNextHookEx(next_hook, nCode, wParam, lParam);
return ::CallNextHookEx(NULL, nCode, wParam, lParam);
}

#ifdef __cplusplus
Expand Down
172 changes: 172 additions & 0 deletions cpp/iedriver/HookProcessor.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,172 @@
// Licensed to the Software Freedom Conservancy (SFC) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The SFC licenses this file
// to you under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "HookProcessor.h"
#include "logging.h"

#define MAX_BUFFER_SIZE 32768

// Define a shared data segment. Variables in this segment can be
// shared across processes that load this DLL.
#pragma data_seg("SHARED")
int data_buffer_size = MAX_BUFFER_SIZE;
char data_buffer[MAX_BUFFER_SIZE];
#pragma data_seg()

#pragma comment(linker, "/section:SHARED,RWS")

namespace webdriver {

class CopyDataHolderWindow : public CWindowImpl<CopyDataHolderWindow> {
public:
DECLARE_WND_CLASS(L"CopyDataHolderWindow")
BEGIN_MSG_MAP(CopyDataHolderWindow)
END_MSG_MAP()

LRESULT CopyData(HWND destination_window_handle,
int data_size,
void* pointer_to_data) {
if (data_size > data_buffer_size) {
LOG(WARN) << "Destination data buffer not large enough";
}

// Copy the contents of local memory into a temporary structure
// for transport across the process boundary.
std::vector<char> buffer(data_size);
memcpy_s(&buffer[0], data_size, pointer_to_data, data_size);

// Send the data across using SendMessage with the WM_COPYDATA
// message. N.B., the window procedure in the other process *must*
// have a handler for the WM_COPYDATA message, which copies the
// content from the message payload into a local buffer. The
// HookProcessor::CopyDataToBuffer method provides a common
// implementation to copy the data into the shared buffer location.
COPYDATASTRUCT data;
data.dwData = 1;
data.cbData = static_cast<int>(buffer.size());
data.lpData = &buffer[0];
LRESULT result = ::SendMessage(destination_window_handle,
WM_COPYDATA,
reinterpret_cast<WPARAM>(this->m_hWnd),
reinterpret_cast<LPARAM>(&data));
return result;
}
};

HookProcessor::HookProcessor(HWND window_handle) {
this->window_handle_ = window_handle;
}

HookProcessor::~HookProcessor(void) {
}

bool HookProcessor::InstallWindowsHook(const std::string& hook_proc_name,
const int hook_proc_type) {
LOG(TRACE) << "Entering HookProcessor::InstallWindowsHook";

HINSTANCE instance_handle = _AtlBaseModule.GetModuleInstance();

FARPROC hook_procedure_address = ::GetProcAddress(instance_handle,
hook_proc_name.c_str());
if (hook_procedure_address == NULL) {
LOGERR(WARN) << "Unable to get address of hook procedure named "
<< hook_proc_name;
return false;
}
HOOKPROC hook_procedure = reinterpret_cast<HOOKPROC>(hook_procedure_address);

// Install the Windows hook.
DWORD thread_id = 0;
if (this->window_handle_ != NULL) {
thread_id = ::GetWindowThreadProcessId(this->window_handle_, NULL);
}
this->hook_procedure_handle_ = ::SetWindowsHookEx(hook_proc_type,
hook_procedure,
instance_handle,
thread_id);
if (this->hook_procedure_handle_ == NULL) {
LOGERR(WARN) << "Unable to set windows hook";
return false;
}
return true;
}

void HookProcessor::UninstallWindowsHook() {
LOG(TRACE) << "Entering HookProcessor::UninstallWindowsHook";
if (this->hook_procedure_handle_ != NULL) {
::UnhookWindowsHookEx(this->hook_procedure_handle_);
}
}

bool HookProcessor::PushData(int data_size,
void* pointer_to_data) {
LOG(TRACE) << "Entering HookProcessor::PushData";
if (this->hook_procedure_handle_ == NULL) {
LOG(WARN) << "No hook procedure has been set";
return false;
}
CopyDataHolderWindow holder;
holder.Create(/*HWND*/ HWND_MESSAGE,
/*_U_RECT rect*/ CWindow::rcDefault,
/*LPCTSTR szWindowName*/ NULL,
/*DWORD dwStyle*/ NULL,
/*DWORD dwExStyle*/ NULL,
/*_U_MENUorID MenuOrID*/ 0U,
/*LPVOID lpCreateParam*/ NULL);
LRESULT result = holder.CopyData(this->window_handle_,
data_size,
pointer_to_data);
LOG(INFO) << "SendMessage result? " << result;
holder.DestroyWindow();
return true;
}

int HookProcessor::GetDataBufferSize() {
return data_buffer_size;
}

void HookProcessor::SetDataBufferSize(int size) {
data_buffer_size = size;
}

void HookProcessor::CopyDataToBuffer(int source_data_size, void* source) {
// clear the shared buffer before putting data into it
ClearBuffer();
if (source_data_size < data_buffer_size) {
data_buffer_size = source_data_size;
}
memcpy_s(data_buffer,
data_buffer_size,
source,
data_buffer_size);
}

void HookProcessor::CopyDataFromBuffer(int destination_data_size, void* destination) {
if (data_buffer_size >= destination_data_size) {
destination_data_size = data_buffer_size;
}
memcpy_s(destination, destination_data_size, data_buffer, destination_data_size);
// clear the shared buffer after taking data out of it.
ClearBuffer();
}

void HookProcessor::ClearBuffer() {
// Zero out the shared buffer
data_buffer_size = MAX_BUFFER_SIZE;
memset(data_buffer, 0, MAX_BUFFER_SIZE);
}

} // namespace webdriver
46 changes: 46 additions & 0 deletions cpp/iedriver/HookProcessor.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
// Licensed to the Software Freedom Conservancy (SFC) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The SFC licenses this file
// to you under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef WEBDRIVER_HOOKPROCESSOR_H_
#define WEBDRIVER_HOOKPROCESSOR_H_

namespace webdriver {

class HookProcessor {
public:
HookProcessor(HWND window_handle);
~HookProcessor(void);

static int GetDataBufferSize(void);
static void SetDataBufferSize(int size);
static void CopyDataToBuffer(int source_data_size, void* source);
static void CopyDataFromBuffer(int destination_data_size, void* destination);

bool InstallWindowsHook(const std::string& hook_proc_name,
const int hook_proc_type);
void UninstallWindowsHook(void);
bool PushData(int data_size,
void* pointer_to_data);
private:
static void ClearBuffer(void);
HWND window_handle_;
HHOOK hook_procedure_handle_;
};

} // namespace webdriver

#endif // WEBDRIVER_HOOKPROCESSOR_H_

2 changes: 1 addition & 1 deletion cpp/iedriver/IEDriver.def
Original file line number Diff line number Diff line change
Expand Up @@ -3,4 +3,4 @@ EXPORTS
ScreenshotWndProc @1
KeyboardHookProc @2
MouseHookProc @3
OverrideWndProc @4
SetProxyWndProc @4
Loading

0 comments on commit 840a86f

Please sign in to comment.