// Copyright 2011 Software Freedom Conservancy | |
// Licensed 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_IE_SCREENSHOTCOMMANDHANDLER_H_ | |
#define WEBDRIVER_IE_SCREENSHOTCOMMANDHANDLER_H_ | |
#include "../Browser.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 { | |
public: | |
ScreenshotCommandHandler(void) { | |
this->image_ = NULL; | |
} | |
virtual ~ScreenshotCommandHandler(void) { | |
} | |
protected: | |
void ExecuteInternal(const IECommandExecutor& executor, | |
const LocatorMap& locator_parameters, | |
const ParametersMap& command_parameters, | |
Response* response) { | |
BrowserHandle browser_wrapper; | |
int status_code = executor.GetCurrentBrowser(&browser_wrapper); | |
if (status_code != SUCCESS) { | |
response->SetErrorResponse(status_code, "Unable to get browser"); | |
return; | |
} | |
bool isSameColour = true; | |
HRESULT hr; | |
int i = 0; | |
do { | |
if (this->image_ != NULL) { | |
delete this->image_; | |
} | |
this->image_ = new CImage(); | |
hr = this->CaptureBrowser(browser_wrapper); | |
if (FAILED(hr)) { | |
LOGHR(WARN, hr) << "Failed to capture browser image at " << i << " try"; | |
delete this->image_; | |
this->image_ = NULL; | |
response->SetSuccessResponse(""); | |
return; | |
} | |
isSameColour = IsSameColour(); | |
if (isSameColour) { | |
::Sleep(2000); | |
LOG(DEBUG) << "Failed to capture non single color browser image at " << i << " try"; | |
} | |
i++; | |
} while (i < 4 && isSameColour); | |
std::string base64_screenshot = ""; | |
hr = this->GetBase64Data(base64_screenshot); | |
if (FAILED(hr)) { | |
LOGHR(WARN, hr) << "Unable to transform browser image to Base64 format"; | |
response->SetSuccessResponse(""); | |
return; | |
} | |
response->SetSuccessResponse(base64_screenshot); | |
delete this->image_; | |
this->image_ = NULL; | |
} | |
private: | |
ATL::CImage* image_; | |
HRESULT CaptureBrowser(BrowserHandle browser) { | |
ie_window_handle = browser->GetTopLevelWindowHandle(); | |
HWND content_window_handle = browser->GetWindowHandle(); | |
CComPtr<IHTMLDocument2> document; | |
browser->GetDocument(&document); | |
LocationInfo document_info; | |
bool result = DocumentHost::GetDocumentDimensions(document, &document_info); | |
if (!result) { | |
LOG(DEBUG) << "Unable to get document dimensions"; | |
return E_FAIL; | |
} | |
int chrome_width(0); | |
int chrome_height(0); | |
this->GetBrowserChromeDimensions(ie_window_handle, | |
content_window_handle, | |
&chrome_width, | |
&chrome_height); | |
max_width = document_info.width + chrome_width; | |
max_height = document_info.height + chrome_height; | |
// For some reason, this technique does not allow the user to resize | |
// the browser window to greater than 65536 x 65536. This is pretty | |
// big, so we'll cap the allowable screenshot size to that. | |
if (max_height > 65536) { | |
LOG(WARN) << L"Height greater than 65536 pixels. Truncating screenshot height to 65536."; | |
max_height = 65536; | |
document_info.height = max_height - chrome_height; | |
} | |
if (max_width > 65536) { | |
LOG(WARN) << L"Width greater than 65536 pixels. Truncating screenshot width to 65536."; | |
max_width = 65536; | |
document_info.width = max_width - chrome_width; | |
} | |
long original_width = browser->GetWidth(); | |
long original_height = browser->GetHeight(); | |
// The resize message is being ignored if the window appears to be | |
// maximized. There's likely a way to bypass that. The kludgy way | |
// is to unmaximize the window, then move on with setting the window | |
// to the dimensions we really want. This is okay because we revert | |
// back to the original dimensions afterward. | |
BOOL is_maximized = ::IsZoomed(ie_window_handle); | |
if (is_maximized) { | |
::ShowWindow(ie_window_handle, SW_SHOWNORMAL); | |
} | |
this->InstallWindowsHook(); | |
browser->SetWidth(max_width); | |
browser->SetHeight(max_height); | |
// Capture the window's canvas to a DIB. | |
BOOL created = this->image_->Create(document_info.width, | |
document_info.height, | |
/*numbers of bits per pixel = */ 32); | |
if (!created) { | |
LOG(WARN) << "Unable to create image"; | |
} | |
HDC device_context_handle = this->image_->GetDC(); | |
BOOL print_result = ::PrintWindow(content_window_handle, | |
device_context_handle, | |
PW_CLIENTONLY); | |
if (!print_result) { | |
LOG(WARN) << L"PrintWindow API returned FALSE"; | |
} | |
this->UninstallWindowsHook(); | |
// Restore the browser to the original dimensions. | |
if (is_maximized) { | |
::ShowWindow(ie_window_handle, SW_MAXIMIZE); | |
} else { | |
browser->SetHeight(original_height); | |
browser->SetWidth(original_width); | |
} | |
this->image_->ReleaseDC(); | |
return S_OK; | |
} | |
bool IsSameColour() { | |
COLORREF firstPixelColour = this->image_->GetPixel(0, 0); | |
for (int i = 0; i < this->image_->GetWidth(); i++) { | |
for (int j = 0; j < this->image_->GetHeight(); j++) { | |
if (firstPixelColour != this->image_->GetPixel(i, j)) { | |
return false; | |
} | |
} | |
} | |
return true; | |
} | |
HRESULT GetBase64Data(std::string& data) { | |
if (this->image_ == NULL) { | |
LOG(DEBUG) << "CImage was not initialized."; | |
return E_POINTER; | |
} | |
CComPtr<IStream> stream; | |
HRESULT hr = ::CreateStreamOnHGlobal(NULL, TRUE, &stream); | |
if (FAILED(hr)) { | |
LOGHR(WARN, hr) << "Error creating IStream"; | |
return hr; | |
} | |
hr = this->image_->Save(stream, Gdiplus::ImageFormatPNG); | |
if (FAILED(hr)) { | |
LOGHR(WARN, hr) << "Saving image failed"; | |
return hr; | |
} | |
// Get the size of the stream. | |
STATSTG statstg; | |
hr = stream->Stat(&statstg, STATFLAG_DEFAULT); | |
if (FAILED(hr)) { | |
LOGHR(WARN, hr) << "No stat on stream"; | |
return hr; | |
} | |
HGLOBAL global_memory_handle = NULL; | |
hr = ::GetHGlobalFromStream(stream, &global_memory_handle); | |
if (FAILED(hr)) { | |
LOGHR(WARN, hr) << "No HGlobal in stream"; | |
return hr; | |
} | |
// TODO: What if the file is bigger than max_int? | |
LOG(DEBUG) << "Size of stream: " << statstg.cbSize.QuadPart; | |
int length = ::Base64EncodeGetRequiredLength(static_cast<int>(statstg.cbSize.QuadPart), | |
ATL_BASE64_FLAG_NOCRLF); | |
if (length <= 0) { | |
LOG(WARN) << "Got zero or negative length from base64 required length"; | |
return E_FAIL; | |
} | |
BYTE* global_lock = reinterpret_cast<BYTE*>(::GlobalLock(global_memory_handle)); | |
if (global_lock == NULL) { | |
::GlobalUnlock(global_memory_handle); | |
LOG(WARN) << "Failure to lock memory"; | |
return E_FAIL; | |
} | |
char* data_array = new char[length + 1]; | |
if (!::Base64Encode(global_lock, | |
static_cast<int>(statstg.cbSize.QuadPart), | |
data_array, | |
&length, | |
ATL_BASE64_FLAG_NOCRLF)) { | |
delete[] data_array; | |
::GlobalUnlock(global_memory_handle); | |
LOG(WARN) << "Failure encoding to base64"; | |
return E_FAIL; | |
} | |
data_array[length] = '\0'; | |
data = data_array; | |
delete[] data_array; | |
::GlobalUnlock(global_memory_handle); | |
return S_OK; | |
} | |
void GetBrowserChromeDimensions(HWND top_level_window_handle, | |
HWND content_window_handle, | |
int* width, | |
int* height) { | |
int top_level_window_width = 0; | |
int top_level_window_height = 0; | |
this->GetWindowDimensions(top_level_window_handle, | |
&top_level_window_width, | |
&top_level_window_height); | |
int content_window_width = 0; | |
int content_window_height = 0; | |
this->GetWindowDimensions(content_window_handle, | |
&content_window_width, | |
&content_window_height); | |
*width = top_level_window_width - content_window_width; | |
*height = top_level_window_height - content_window_height; | |
} | |
void ScreenshotCommandHandler::GetWindowDimensions(HWND window_handle, | |
int* width, | |
int* height) { | |
RECT window_rect; | |
::GetWindowRect(window_handle, &window_rect); | |
*width = window_rect.right - window_rect.left; | |
*height = window_rect.bottom - window_rect.top; | |
} | |
void InstallWindowsHook() { | |
HINSTANCE instance_handle = _AtlBaseModule.GetModuleInstance(); | |
HOOKPROC hook_procedure = reinterpret_cast<HOOKPROC>(::GetProcAddress(instance_handle, | |
"ScreenshotWndProc")); | |
if (hook_procedure == NULL) { | |
LOG(WARN) << L"GetProcAddress return value was NULL"; | |
return; | |
} | |
// 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) { | |
DWORD error = ::GetLastError(); | |
LOG(WARN) << L"SetWindowsHookEx return value was NULL, actual error code was " << error; | |
} | |
} | |
void UninstallWindowsHook() { | |
::UnhookWindowsHookEx(next_hook); | |
} | |
}; | |
} // namespace webdriver | |
#ifdef __cplusplus | |
extern "C" { | |
#endif | |
// This function is our message processor that we inject into the IEFrame | |
// process. Its sole purpose is to process WM_GETMINMAXINFO messages and | |
// modify the max tracking size so that we can resize the IEFrame window to | |
// greater than the virtual screen resolution. All other messages are | |
// delegated to the original IEFrame message processor. This function | |
// uninjects itself immediately upon execution. | |
LRESULT CALLBACK MinMaxInfoHandler(HWND hwnd, | |
UINT message, | |
WPARAM wParam, | |
LPARAM lParam) { | |
// Grab a reference to the original message processor. | |
HANDLE original_message_proc = ::GetProp(hwnd, | |
L"__original_message_processor__"); | |
::RemoveProp(hwnd, L"__original_message_processor__"); | |
// Uninject this method. | |
::SetWindowLongPtr(hwnd, | |
GWLP_WNDPROC, | |
reinterpret_cast<LONG_PTR>(original_message_proc)); | |
if (WM_GETMINMAXINFO == message) { | |
MINMAXINFO* minMaxInfo = reinterpret_cast<MINMAXINFO*>(lParam); | |
minMaxInfo->ptMaxTrackSize.x = max_width; | |
minMaxInfo->ptMaxTrackSize.y = max_height; | |
// We're not going to pass this message onto the original message | |
// processor, so we should return 0, per the documentation for | |
// the WM_GETMINMAXINFO message. | |
return 0; | |
} | |
// All other messages should be handled by the original message processor. | |
return ::CallWindowProc(reinterpret_cast<WNDPROC>(original_message_proc), | |
hwnd, | |
message, | |
wParam, | |
lParam); | |
} | |
// Many thanks to sunnyandy for helping out with this approach. What we're | |
// doing here is setting up a Windows hook to see incoming messages to the | |
// IEFrame's message processor. Once we find one that's WM_GETMINMAXINFO, | |
// we inject our own message processor into the IEFrame process to handle | |
// that one message. WM_GETMINMAXINFO is sent on a resize event so the process | |
// can see how large a window can be. By modifying the max values, we can allow | |
// a window to be sized greater than the (virtual) screen resolution would | |
// otherwise allow. | |
// | |
// 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) { | |
// 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 | |
// that can. | |
LONG_PTR proc = ::SetWindowLongPtr(call_window_proc_struct->hwnd, | |
GWLP_WNDPROC, | |
reinterpret_cast<LONG_PTR>(MinMaxInfoHandler)); | |
::SetProp(call_window_proc_struct->hwnd, | |
L"__original_message_processor__", | |
reinterpret_cast<HANDLE>(proc)); | |
} | |
return ::CallNextHookEx(next_hook, nCode, wParam, lParam); | |
} | |
#ifdef __cplusplus | |
} | |
#endif | |
#endif // WEBDRIVER_IE_SCREENSHOTCOMMANDHANDLER_H_ |