chrome_frame/test_utils.cc

   1 // Copyright (c) 2012 The Chromium Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style license that can be
   3 // found in the LICENSE file.
   4
   5 #include "chrome_frame/test_utils.h"
   6
   7 #include <atlbase.h>
   8 #include <atlwin.h>
   9 #include <shellapi.h>
  10 #include <winternl.h>
  11
  12 #include <algorithm>
  13
  14 #include "base/command_line.h"
  15 #include "base/file_path.h"
  16 #include "base/file_util.h"
  17 #include "base/logging.h"
  18 #include "base/path_service.h"
  19 #include "base/process_util.h"
  20 #include "base/string_util.h"
  21 #include "base/stringprintf.h"
  22 #include "base/utf_string_conversions.h"
  23 #include "base/win/scoped_handle.h"
  24 #include "chrome/common/chrome_paths.h"
  25 #include "chrome/common/chrome_switches.h"
  26 #include "chrome_frame/test/chrome_frame_test_utils.h"
  27 #include "testing/gtest/include/gtest/gtest.h"
  28
  29 const wchar_t kChromeFrameDllName[] = L"npchrome_frame.dll";
  30 const wchar_t kChromeLauncherExeName[] = L"chrome_launcher.exe";
  31 // How long to wait for DLLs to register or unregister themselves before killing
  32 // the registrar.
  33 const int64 kDllRegistrationTimeoutMs = 30 * 1000;
  34
  35 base::FilePath GetChromeFrameBuildPath() {
  36   base::FilePath build_path;
  37   PathService::Get(chrome::DIR_APP, &build_path);
  38
  39   base::FilePath dll_path = build_path.Append(kChromeFrameDllName);
  40
  41   if (!file_util::PathExists(dll_path)) {
  42     // Well, dang.. try looking in the current directory.
  43     dll_path = build_path.Append(kChromeFrameDllName);
  44   }
  45
  46   if (!file_util::PathExists(dll_path)) {
  47     // No luck, return something empty.
  48     dll_path = base::FilePath();
  49   }
  50
  51   return dll_path;
  52 }
  53
  54 const wchar_t ScopedChromeFrameRegistrar::kCallRegistrationEntrypointSwitch[] =
  55     L"--call-registration-entrypoint";
  56
  57 bool ScopedChromeFrameRegistrar::register_chrome_path_provider_ = false;
  58
  59 // static
  60 void ScopedChromeFrameRegistrar::RegisterDefaults() {
  61   if (!register_chrome_path_provider_) {
  62     chrome::RegisterPathProvider();
  63     register_chrome_path_provider_ = true;
  64   }
  65 }
  66
  67 // Registers or unregisters the DLL at |path| by calling out to the current
  68 // executable with --call-registration-entrypoint.  Loading the DLL into the
  69 // test process is problematic for component=shared_library builds since
  70 // singletons in base.dll aren't designed to handle multiple initialization.
  71 // Use of rundll32.exe is problematic since it does not return useful error
  72 // information.
  73 // static
  74 void ScopedChromeFrameRegistrar::DoRegistration(
  75     const string16& path,
  76     RegistrationType registration_type,
  77     RegistrationOperation registration_operation) {
  78   static const char* const kEntrypoints[] = {
  79     "DllRegisterServer",
  80     "DllUnregisterServer",
  81     "DllRegisterUserServer",
  82     "DllUnregisterUserServer",
  83   };
  84
  85   DCHECK(!path.empty());
  86   DCHECK(registration_type == PER_USER || registration_type == SYSTEM_LEVEL);
  87   DCHECK(registration_operation == REGISTER ||
  88          registration_operation == UNREGISTER);
  89
  90   int entrypoint_index = 0;
  91   base::LaunchOptions launch_options;
  92   base::ProcessHandle process_handle = INVALID_HANDLE_VALUE;
  93   int exit_code = -1;
  94
  95   if (registration_type == PER_USER)
  96     entrypoint_index += 2;
  97   if (registration_operation == UNREGISTER)
  98     entrypoint_index += 1;
  99   string16 registration_command(ASCIIToUTF16("\""));
 100   registration_command +=
 101       CommandLine::ForCurrentProcess()->GetProgram().value();
 102   registration_command += ASCIIToUTF16("\" ");
 103   registration_command += kCallRegistrationEntrypointSwitch;
 104   registration_command += ASCIIToUTF16(" \"");
 105   registration_command += path;
 106   registration_command += ASCIIToUTF16("\" ");
 107   registration_command += ASCIIToUTF16(kEntrypoints[entrypoint_index]);
 108   launch_options.wait = true;
 109   if (!base::LaunchProcess(registration_command, launch_options,
 110                            &process_handle)) {
 111     PLOG(FATAL)
 112         << "Failed to register or unregister DLL with command: "
 113         << registration_command;
 114   } else {
 115     base::win::ScopedHandle rundll32(process_handle);
 116     if (!base::WaitForExitCodeWithTimeout(
 117             process_handle, &exit_code,
 118             base::TimeDelta::FromMilliseconds(kDllRegistrationTimeoutMs))) {
 119       LOG(ERROR) << "Timeout waiting to register or unregister DLL with "
 120                     "command: " << registration_command;
 121       base::KillProcess(process_handle, 0, false);
 122       NOTREACHED() << "Aborting test due to registration failure.";
 123     }
 124   }
 125   if (exit_code != 0) {
 126     if (registration_operation == REGISTER) {
 127       LOG(ERROR)
 128           << "DLL registration failed (exit code: 0x" << std::hex << exit_code
 129           << ", command: " << registration_command
 130           << "). Make sure you are running as Admin.";
 131       ::ExitProcess(1);
 132     } else {
 133       LOG(WARNING)
 134           << "DLL unregistration failed (exit code: 0x" << std::hex << exit_code
 135           << ", command: " << registration_command << ").";
 136     }
 137   }
 138 }
 139
 140 // static
 141 void ScopedChromeFrameRegistrar::RegisterAtPath(
 142     const std::wstring& path, RegistrationType registration_type) {
 143   DoRegistration(path, registration_type, REGISTER);
 144 }
 145
 146 // static
 147 void ScopedChromeFrameRegistrar::UnregisterAtPath(
 148     const std::wstring& path, RegistrationType registration_type) {
 149   DoRegistration(path, registration_type, UNREGISTER);
 150 }
 151
 152 base::FilePath ScopedChromeFrameRegistrar::GetReferenceChromeFrameDllPath() {
 153   base::FilePath reference_build_dir;
 154   PathService::Get(chrome::DIR_APP, &reference_build_dir);
 155
 156   reference_build_dir = reference_build_dir.DirName();
 157   reference_build_dir = reference_build_dir.DirName();
 158
 159   reference_build_dir = reference_build_dir.AppendASCII("chrome_frame");
 160   reference_build_dir = reference_build_dir.AppendASCII("tools");
 161   reference_build_dir = reference_build_dir.AppendASCII("test");
 162   reference_build_dir = reference_build_dir.AppendASCII("reference_build");
 163   reference_build_dir = reference_build_dir.AppendASCII("chrome");
 164   reference_build_dir = reference_build_dir.AppendASCII("servers");
 165   reference_build_dir = reference_build_dir.Append(kChromeFrameDllName);
 166   return reference_build_dir;
 167 }
 168
 169 // static
 170 void ScopedChromeFrameRegistrar::RegisterAndExitProcessIfDirected() {
 171   // This method is invoked before any Chromium helpers have been initialized.
 172   // Take pains to use only Win32 and CRT functions.
 173   int argc = 0;
 174   const wchar_t* const* argv = ::CommandLineToArgvW(::GetCommandLine(), &argc);
 175   if (argc < 2 || ::lstrcmp(argv[1], kCallRegistrationEntrypointSwitch) != 0)
 176     return;
 177   if (argc != 4) {
 178     printf("Usage: %S %S <path to dll> <entrypoint>\n", argv[0],
 179            kCallRegistrationEntrypointSwitch);
 180     return;
 181   }
 182
 183   // The only way to leave from here on down is ExitProcess.
 184   const wchar_t* dll_path = argv[2];
 185   const wchar_t* wide_entrypoint = argv[3];
 186   char entrypoint[256];
 187   HRESULT exit_code = 0;
 188   int entrypoint_len = lstrlen(wide_entrypoint);
 189   if (entrypoint_len <= 0 || entrypoint_len >= arraysize(entrypoint)) {
 190     exit_code = HRESULT_FROM_WIN32(ERROR_PROC_NOT_FOUND);
 191   } else {
 192     // Convert wide to narrow. Since the entrypoint must be a narrow string
 193     // anyway, it is safe to truncate each character like this.
 194     std::copy(wide_entrypoint, wide_entrypoint + entrypoint_len + 1,
 195               &entrypoint[0]);
 196     HMODULE dll_module = ::LoadLibrary(dll_path);
 197     if (dll_module == NULL) {
 198       exit_code = HRESULT_FROM_WIN32(::GetLastError());
 199     } else {
 200       typedef HRESULT (STDAPICALLTYPE *RegisterFp)();
 201       RegisterFp register_func =
 202           reinterpret_cast<RegisterFp>(::GetProcAddress(dll_module,
 203                                                         entrypoint));
 204       if (register_func == NULL) {
 205         exit_code = HRESULT_FROM_WIN32(::GetLastError());
 206       } else {
 207         exit_code = register_func();
 208       }
 209       ::FreeLibrary(dll_module);
 210     }
 211   }
 212
 213   ::ExitProcess(exit_code);
 214 }
 215
 216 // Non-statics
 217
 218 ScopedChromeFrameRegistrar::ScopedChromeFrameRegistrar(
 219     const std::wstring& path, RegistrationType registration_type)
 220     : registration_type_(registration_type) {
 221   if (!register_chrome_path_provider_) {
 222     // Register paths needed by the ScopedChromeFrameRegistrar.
 223     chrome::RegisterPathProvider();
 224     register_chrome_path_provider_ = true;
 225   }
 226   original_dll_path_ = path;
 227   RegisterChromeFrameAtPath(original_dll_path_);
 228 }
 229
 230 ScopedChromeFrameRegistrar::ScopedChromeFrameRegistrar(
 231     RegistrationType registration_type)
 232     : registration_type_(registration_type) {
 233   if (!register_chrome_path_provider_) {
 234     // Register paths needed by the ScopedChromeFrameRegistrar.
 235     chrome::RegisterPathProvider();
 236     register_chrome_path_provider_ = true;
 237   }
 238   original_dll_path_ = GetChromeFrameBuildPath().value();
 239   RegisterChromeFrameAtPath(original_dll_path_);
 240 }
 241
 242 ScopedChromeFrameRegistrar::~ScopedChromeFrameRegistrar() {
 243   if (base::FilePath(original_dll_path_) !=
 244       base::FilePath(new_chrome_frame_dll_path_)) {
 245     RegisterChromeFrameAtPath(original_dll_path_);
 246   } else if (registration_type_ == PER_USER) {
 247     UnregisterAtPath(new_chrome_frame_dll_path_, registration_type_);
 248     HWND chrome_frame_helper_window =
 249         FindWindow(L"ChromeFrameHelperWindowClass", NULL);
 250     if (IsWindow(chrome_frame_helper_window)) {
 251       PostMessage(chrome_frame_helper_window, WM_CLOSE, 0, 0);
 252     } else {
 253       base::KillProcesses(L"chrome_frame_helper.exe", 0, NULL);
 254     }
 255   }
 256 }
 257
 258 void ScopedChromeFrameRegistrar::RegisterChromeFrameAtPath(
 259     const std::wstring& path) {
 260   RegisterAtPath(path, registration_type_);
 261   new_chrome_frame_dll_path_ = path;
 262 }
 263
 264 void ScopedChromeFrameRegistrar::RegisterReferenceChromeFrameBuild() {
 265   RegisterChromeFrameAtPath(GetReferenceChromeFrameDllPath().value());
 266 }
 267
 268 std::wstring ScopedChromeFrameRegistrar::GetChromeFrameDllPath() const {
 269   return new_chrome_frame_dll_path_;
 270 }
 271
 272 // TODO(robertshield): The following could be factored out into its own file.
 273 namespace {
 274
 275 typedef LONG WINAPI
 276 NtQueryInformationProcess(
 277   IN HANDLE ProcessHandle,
 278   IN PROCESSINFOCLASS ProcessInformationClass,
 279   OUT PVOID ProcessInformation,
 280   IN ULONG ProcessInformationLength,
 281   OUT PULONG ReturnLength OPTIONAL
 282 );
 283
 284 // Get the function pointer to NtQueryInformationProcess in NTDLL.DLL
 285 static bool GetQIP(NtQueryInformationProcess** qip_func_ptr) {
 286   static NtQueryInformationProcess* qip_func =
 287       reinterpret_cast<NtQueryInformationProcess*>(
 288           GetProcAddress(GetModuleHandle(L"ntdll.dll"),
 289           "NtQueryInformationProcess"));
 290   DCHECK(qip_func) << "Could not get pointer to NtQueryInformationProcess.";
 291   *qip_func_ptr = qip_func;
 292   return qip_func != NULL;
 293 }
 294
 295 // Get the command line of a process
 296 bool GetCommandLineForProcess(uint32 process_id, std::wstring* cmd_line) {
 297   DCHECK(process_id != 0);
 298   DCHECK(cmd_line);
 299
 300   // Open the process
 301   base::win::ScopedHandle process_handle(::OpenProcess(
 302       PROCESS_QUERY_INFORMATION | PROCESS_VM_READ,
 303       false,
 304       process_id));
 305   if (!process_handle) {
 306     DLOG(ERROR) << "Failed to open process " << process_id << ", last error = "
 307                 << GetLastError();
 308   }
 309
 310   // Obtain Process Environment Block
 311   NtQueryInformationProcess* qip_func = NULL;
 312   if (process_handle) {
 313     GetQIP(&qip_func);
 314   }
 315
 316   // Read the address of the process params from the peb.
 317   DWORD process_params_address = 0;
 318   if (qip_func) {
 319     PROCESS_BASIC_INFORMATION info = { 0 };
 320     // NtQueryInformationProcess returns an NTSTATUS for whom negative values
 321     // are negative. Just check for that instead of pulling in DDK macros.
 322     if ((qip_func(process_handle.Get(),
 323                   ProcessBasicInformation,
 324                   &info,
 325                   sizeof(info),
 326                   NULL)) < 0) {
 327       DLOG(ERROR) << "Failed to invoke NtQueryProcessInformation, last error = "
 328                   << GetLastError();
 329     } else {
 330       BYTE* peb = reinterpret_cast<BYTE*>(info.PebBaseAddress);
 331
 332       // The process command line parameters are (or were once) located at
 333       // the base address of the PEB + 0x10 for 32 bit processes. 64 bit
 334       // processes have a different PEB struct as per
 335       // http://msdn.microsoft.com/en-us/library/aa813706(VS.85).aspx.
 336       // TODO(robertshield): See about doing something about this.
 337       SIZE_T bytes_read = 0;
 338       if (!::ReadProcessMemory(process_handle.Get(),
 339                                peb + 0x10,
 340                                &process_params_address,
 341                                sizeof(process_params_address),
 342                                &bytes_read)) {
 343         DLOG(ERROR) << "Failed to read process params address, last error = "
 344                     << GetLastError();
 345       }
 346     }
 347   }
 348
 349   // Copy all the process parameters into a buffer.
 350   bool success = false;
 351   std::wstring buffer;
 352   if (process_params_address) {
 353     SIZE_T bytes_read;
 354     RTL_USER_PROCESS_PARAMETERS params = { 0 };
 355     if (!::ReadProcessMemory(process_handle.Get(),
 356                              reinterpret_cast<void*>(process_params_address),
 357                              &params,
 358                              sizeof(params),
 359                              &bytes_read)) {
 360       DLOG(ERROR) << "Failed to read RTL_USER_PROCESS_PARAMETERS, "
 361                   << "last error = " << GetLastError();
 362     } else {
 363       // Read the command line parameter
 364       const int max_cmd_line_len = std::min(
 365           static_cast<int>(params.CommandLine.MaximumLength),
 366           4096);
 367       buffer.resize(max_cmd_line_len + 1);
 368       if (!::ReadProcessMemory(process_handle.Get(),
 369                                params.CommandLine.Buffer,
 370                                &buffer[0],
 371                                max_cmd_line_len,
 372                                &bytes_read)) {
 373         DLOG(ERROR) << "Failed to copy process command line, "
 374                     << "last error = " << GetLastError();
 375       } else {
 376         *cmd_line = buffer;
 377         success = true;
 378       }
 379     }
 380   }
 381
 382   return success;
 383 }
 384
 385 // Used to filter processes by process ID.
 386 class ArgumentFilter : public base::ProcessFilter {
 387  public:
 388   explicit ArgumentFilter(const std::wstring& argument)
 389       : argument_to_find_(argument) {}
 390
 391   // Returns true to indicate set-inclusion and false otherwise.  This method
 392   // should not have side-effects and should be idempotent.
 393   virtual bool Includes(const base::ProcessEntry& entry) const {
 394     bool found = false;
 395     std::wstring command_line;
 396     if (GetCommandLineForProcess(entry.pid(), &command_line)) {
 397       std::wstring::const_iterator it =
 398           std::search(command_line.begin(),
 399                       command_line.end(),
 400                       argument_to_find_.begin(),
 401                       argument_to_find_.end(),
 402           base::CaseInsensitiveCompareASCII<wchar_t>());
 403       found = (it != command_line.end());
 404     }
 405     return found;
 406   }
 407
 408  protected:
 409   std::wstring argument_to_find_;
 410 };
 411
 412 }  // namespace
 413
 414 bool KillAllNamedProcessesWithArgument(const std::wstring& process_name,
 415                                        const std::wstring& argument) {
 416   return base::KillProcesses(process_name, 0, &ArgumentFilter(argument));
 417 }
 418
 419 bool IsWorkstationLocked() {
 420   bool is_locked = true;
 421   HDESK input_desk = ::OpenInputDesktop(0, 0, GENERIC_READ);
 422   if (input_desk)  {
 423     wchar_t name[256] = {0};
 424     DWORD needed = 0;
 425     if (::GetUserObjectInformation(input_desk,
 426       UOI_NAME,
 427       name,
 428       sizeof(name),
 429       &needed)) {
 430         is_locked = lstrcmpi(name, L"default") != 0;
 431     }
 432     ::CloseDesktop(input_desk);
 433   }
 434   return is_locked;
 435 }