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 FilePath GetChromeFrameBuildPath() {
  36   FilePath build_path;
  37   PathService::Get(chrome::DIR_APP, &build_path);
  38
  39   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 = 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(process_handle, &exit_code,
 117                                           kDllRegistrationTimeoutMs)) {
 118       LOG(ERROR) << "Timeout waiting to register or unregister DLL with "
 119                     "command: " << registration_command;
 120       base::KillProcess(process_handle, 0, false);
 121       NOTREACHED() << "Aborting test due to registration failure.";
 122     }
 123   }
 124   if (exit_code != 0) {
 125     if (registration_operation == REGISTER) {
 126       LOG(FATAL)
 127           << "DLL registration failed (exit code: 0x" << std::hex << exit_code
 128           << ", command: " << registration_command
 129           << "). Make sure you are running as Admin.";
 130     } else {
 131       LOG(WARNING)
 132           << "DLL unregistration failed (exit code: 0x" << std::hex << exit_code
 133           << ", command: " << registration_command << ").";
 134     }
 135   }
 136 }
 137
 138 // static
 139 void ScopedChromeFrameRegistrar::RegisterAtPath(
 140     const std::wstring& path, RegistrationType registration_type) {
 141   DoRegistration(path, registration_type, REGISTER);
 142 }
 143
 144 // static
 145 void ScopedChromeFrameRegistrar::UnregisterAtPath(
 146     const std::wstring& path, RegistrationType registration_type) {
 147   DoRegistration(path, registration_type, UNREGISTER);
 148 }
 149
 150 FilePath ScopedChromeFrameRegistrar::GetReferenceChromeFrameDllPath() {
 151   FilePath reference_build_dir;
 152   PathService::Get(chrome::DIR_APP, &reference_build_dir);
 153
 154   reference_build_dir = reference_build_dir.DirName();
 155   reference_build_dir = reference_build_dir.DirName();
 156
 157   reference_build_dir = reference_build_dir.AppendASCII("chrome_frame");
 158   reference_build_dir = reference_build_dir.AppendASCII("tools");
 159   reference_build_dir = reference_build_dir.AppendASCII("test");
 160   reference_build_dir = reference_build_dir.AppendASCII("reference_build");
 161   reference_build_dir = reference_build_dir.AppendASCII("chrome");
 162   reference_build_dir = reference_build_dir.AppendASCII("servers");
 163   reference_build_dir = reference_build_dir.Append(kChromeFrameDllName);
 164   return reference_build_dir;
 165 }
 166
 167 // static
 168 void ScopedChromeFrameRegistrar::RegisterAndExitProcessIfDirected() {
 169   // This method is invoked before any Chromium helpers have been initialized.
 170   // Take pains to use only Win32 and CRT functions.
 171   int argc = 0;
 172   const wchar_t* const* argv = ::CommandLineToArgvW(::GetCommandLine(), &argc);
 173   if (argc < 2 || ::lstrcmp(argv[1], kCallRegistrationEntrypointSwitch) != 0)
 174     return;
 175   if (argc != 4) {
 176     printf("Usage: %S %S <path to dll> <entrypoint>\n", argv[0],
 177            kCallRegistrationEntrypointSwitch);
 178     return;
 179   }
 180
 181   // The only way to leave from here on down is ExitProcess.
 182   const wchar_t* dll_path = argv[2];
 183   const wchar_t* wide_entrypoint = argv[3];
 184   char entrypoint[256];
 185   HRESULT exit_code = 0;
 186   int entrypoint_len = lstrlen(wide_entrypoint);
 187   if (entrypoint_len <= 0 || entrypoint_len >= arraysize(entrypoint)) {
 188     exit_code = HRESULT_FROM_WIN32(ERROR_PROC_NOT_FOUND);
 189   } else {
 190     // Convert wide to narrow. Since the entrypoint must be a narrow string
 191     // anyway, it is safe to truncate each character like this.
 192     std::copy(wide_entrypoint, wide_entrypoint + entrypoint_len + 1,
 193               &entrypoint[0]);
 194     HMODULE dll_module = ::LoadLibrary(dll_path);
 195     if (dll_module == NULL) {
 196       exit_code = HRESULT_FROM_WIN32(::GetLastError());
 197     } else {
 198       typedef HRESULT (STDAPICALLTYPE *RegisterFp)();
 199       RegisterFp register_func =
 200           reinterpret_cast<RegisterFp>(::GetProcAddress(dll_module,
 201                                                         entrypoint));
 202       if (register_func == NULL) {
 203         exit_code = HRESULT_FROM_WIN32(::GetLastError());
 204       } else {
 205         exit_code = register_func();
 206       }
 207       ::FreeLibrary(dll_module);
 208     }
 209   }
 210
 211   ::ExitProcess(exit_code);
 212 }
 213
 214 // Non-statics
 215
 216 ScopedChromeFrameRegistrar::ScopedChromeFrameRegistrar(
 217     const std::wstring& path, RegistrationType registration_type)
 218     : registration_type_(registration_type) {
 219   if (!register_chrome_path_provider_) {
 220     // Register paths needed by the ScopedChromeFrameRegistrar.
 221     chrome::RegisterPathProvider();
 222     register_chrome_path_provider_ = true;
 223   }
 224   original_dll_path_ = path;
 225   RegisterChromeFrameAtPath(original_dll_path_);
 226 }
 227
 228 ScopedChromeFrameRegistrar::ScopedChromeFrameRegistrar(
 229     RegistrationType registration_type)
 230     : registration_type_(registration_type) {
 231   if (!register_chrome_path_provider_) {
 232     // Register paths needed by the ScopedChromeFrameRegistrar.
 233     chrome::RegisterPathProvider();
 234     register_chrome_path_provider_ = true;
 235   }
 236   original_dll_path_ = GetChromeFrameBuildPath().value();
 237   RegisterChromeFrameAtPath(original_dll_path_);
 238 }
 239
 240 ScopedChromeFrameRegistrar::~ScopedChromeFrameRegistrar() {
 241   if (FilePath(original_dll_path_) != FilePath(new_chrome_frame_dll_path_)) {
 242     RegisterChromeFrameAtPath(original_dll_path_);
 243   } else if (registration_type_ == PER_USER) {
 244     UnregisterAtPath(new_chrome_frame_dll_path_, registration_type_);
 245     HWND chrome_frame_helper_window =
 246         FindWindow(L"ChromeFrameHelperWindowClass", NULL);
 247     if (IsWindow(chrome_frame_helper_window)) {
 248       PostMessage(chrome_frame_helper_window, WM_CLOSE, 0, 0);
 249     } else {
 250       chrome_frame_test::KillProcesses(L"chrome_frame_helper.exe", 0, false);
 251     }
 252   }
 253 }
 254
 255 void ScopedChromeFrameRegistrar::RegisterChromeFrameAtPath(
 256     const std::wstring& path) {
 257   RegisterAtPath(path, registration_type_);
 258   new_chrome_frame_dll_path_ = path;
 259 }
 260
 261 void ScopedChromeFrameRegistrar::RegisterReferenceChromeFrameBuild() {
 262   RegisterChromeFrameAtPath(GetReferenceChromeFrameDllPath().value());
 263 }
 264
 265 std::wstring ScopedChromeFrameRegistrar::GetChromeFrameDllPath() const {
 266   return new_chrome_frame_dll_path_;
 267 }
 268
 269 // TODO(robertshield): The following could be factored out into its own file.
 270 namespace {
 271
 272 typedef LONG WINAPI
 273 NtQueryInformationProcess(
 274   IN HANDLE ProcessHandle,
 275   IN PROCESSINFOCLASS ProcessInformationClass,
 276   OUT PVOID ProcessInformation,
 277   IN ULONG ProcessInformationLength,
 278   OUT PULONG ReturnLength OPTIONAL
 279 );
 280
 281 // Get the function pointer to NtQueryInformationProcess in NTDLL.DLL
 282 static bool GetQIP(NtQueryInformationProcess** qip_func_ptr) {
 283   static NtQueryInformationProcess* qip_func =
 284       reinterpret_cast<NtQueryInformationProcess*>(
 285           GetProcAddress(GetModuleHandle(L"ntdll.dll"),
 286           "NtQueryInformationProcess"));
 287   DCHECK(qip_func) << "Could not get pointer to NtQueryInformationProcess.";
 288   *qip_func_ptr = qip_func;
 289   return qip_func != NULL;
 290 }
 291
 292 // Get the command line of a process
 293 bool GetCommandLineForProcess(uint32 process_id, std::wstring* cmd_line) {
 294   DCHECK(process_id != 0);
 295   DCHECK(cmd_line);
 296
 297   // Open the process
 298   base::win::ScopedHandle process_handle(::OpenProcess(
 299       PROCESS_QUERY_INFORMATION | PROCESS_VM_READ,
 300       false,
 301       process_id));
 302   if (!process_handle) {
 303     DLOG(ERROR) << "Failed to open process " << process_id << ", last error = "
 304                 << GetLastError();
 305   }
 306
 307   // Obtain Process Environment Block
 308   NtQueryInformationProcess* qip_func = NULL;
 309   if (process_handle) {
 310     GetQIP(&qip_func);
 311   }
 312
 313   // Read the address of the process params from the peb.
 314   DWORD process_params_address = 0;
 315   if (qip_func) {
 316     PROCESS_BASIC_INFORMATION info = { 0 };
 317     // NtQueryInformationProcess returns an NTSTATUS for whom negative values
 318     // are negative. Just check for that instead of pulling in DDK macros.
 319     if ((qip_func(process_handle.Get(),
 320                   ProcessBasicInformation,
 321                   &info,
 322                   sizeof(info),
 323                   NULL)) < 0) {
 324       DLOG(ERROR) << "Failed to invoke NtQueryProcessInformation, last error = "
 325                   << GetLastError();
 326     } else {
 327       BYTE* peb = reinterpret_cast<BYTE*>(info.PebBaseAddress);
 328
 329       // The process command line parameters are (or were once) located at
 330       // the base address of the PEB + 0x10 for 32 bit processes. 64 bit
 331       // processes have a different PEB struct as per
 332       // http://msdn.microsoft.com/en-us/library/aa813706(VS.85).aspx.
 333       // TODO(robertshield): See about doing something about this.
 334       SIZE_T bytes_read = 0;
 335       if (!::ReadProcessMemory(process_handle.Get(),
 336                                peb + 0x10,
 337                                &process_params_address,
 338                                sizeof(process_params_address),
 339                                &bytes_read)) {
 340         DLOG(ERROR) << "Failed to read process params address, last error = "
 341                     << GetLastError();
 342       }
 343     }
 344   }
 345
 346   // Copy all the process parameters into a buffer.
 347   bool success = false;
 348   std::wstring buffer;
 349   if (process_params_address) {
 350     SIZE_T bytes_read;
 351     RTL_USER_PROCESS_PARAMETERS params = { 0 };
 352     if (!::ReadProcessMemory(process_handle.Get(),
 353                              reinterpret_cast<void*>(process_params_address),
 354                              &params,
 355                              sizeof(params),
 356                              &bytes_read)) {
 357       DLOG(ERROR) << "Failed to read RTL_USER_PROCESS_PARAMETERS, "
 358                   << "last error = " << GetLastError();
 359     } else {
 360       // Read the command line parameter
 361       const int max_cmd_line_len = std::min(
 362           static_cast<int>(params.CommandLine.MaximumLength),
 363           4096);
 364       buffer.resize(max_cmd_line_len + 1);
 365       if (!::ReadProcessMemory(process_handle.Get(),
 366                                params.CommandLine.Buffer,
 367                                &buffer[0],
 368                                max_cmd_line_len,
 369                                &bytes_read)) {
 370         DLOG(ERROR) << "Failed to copy process command line, "
 371                     << "last error = " << GetLastError();
 372       } else {
 373         *cmd_line = buffer;
 374         success = true;
 375       }
 376     }
 377   }
 378
 379   return success;
 380 }
 381
 382 // Used to filter processes by process ID.
 383 class ArgumentFilter : public base::ProcessFilter {
 384  public:
 385   explicit ArgumentFilter(const std::wstring& argument)
 386       : argument_to_find_(argument) {}
 387
 388   // Returns true to indicate set-inclusion and false otherwise.  This method
 389   // should not have side-effects and should be idempotent.
 390   virtual bool Includes(const base::ProcessEntry& entry) const {
 391     bool found = false;
 392     std::wstring command_line;
 393     if (GetCommandLineForProcess(entry.pid(), &command_line)) {
 394       std::wstring::const_iterator it =
 395           std::search(command_line.begin(),
 396                       command_line.end(),
 397                       argument_to_find_.begin(),
 398                       argument_to_find_.end(),
 399           base::CaseInsensitiveCompareASCII<wchar_t>());
 400       found = (it != command_line.end());
 401     }
 402     return found;
 403   }
 404
 405  protected:
 406   std::wstring argument_to_find_;
 407 };
 408
 409 }  // namespace
 410
 411 bool KillAllNamedProcessesWithArgument(const std::wstring& process_name,
 412                                        const std::wstring& argument) {
 413   bool result = true;
 414   ArgumentFilter filter(argument);
 415   base::NamedProcessIterator iter(process_name, &filter);
 416   while (const base::ProcessEntry* entry = iter.NextProcessEntry()) {
 417     if (!base::KillProcessById(entry->pid(), 0, true)) {
 418       result = false;
 419     }
 420   }
 421
 422   return result;
 423 }
 424
 425 bool IsWorkstationLocked() {
 426   bool is_locked = true;
 427   HDESK input_desk = ::OpenInputDesktop(0, 0, GENERIC_READ);
 428   if (input_desk)  {
 429     wchar_t name[256] = {0};
 430     DWORD needed = 0;
 431     if (::GetUserObjectInformation(input_desk,
 432       UOI_NAME,
 433       name,
 434       sizeof(name),
 435       &needed)) {
 436         is_locked = lstrcmpi(name, L"default") != 0;
 437     }
 438     ::CloseDesktop(input_desk);
 439   }
 440   return is_locked;
 441 }