diff options
Diffstat (limited to 'ceee/testing/utils/gflag_utils.cc')
| -rw-r--r-- | ceee/testing/utils/gflag_utils.cc | 157 |
1 files changed, 157 insertions, 0 deletions
diff --git a/ceee/testing/utils/gflag_utils.cc b/ceee/testing/utils/gflag_utils.cc new file mode 100644 index 0000000..b30553a --- /dev/null +++ b/ceee/testing/utils/gflag_utils.cc @@ -0,0 +1,157 @@ +// Copyright (c) 2010 The Chromium Authors. All rights reserved. +// Use of this source code is governed by a BSD-style license that can be +// found in the LICENSE file. +// +// Implementation of the launch with gflags function. +#include "ceee/testing/utils/gflag_utils.h" + +#include "ceee/testing/utils/nt_internals.h" +#include <shlwapi.h> + +namespace testing { + +using nt_internals::NTSTATUS; +using nt_internals::NtQueryInformationProcess; + +HRESULT WriteProcessGFlags(HANDLE process, DWORD gflags) { + nt_internals::PROCESS_BASIC_INFORMATION basic_info = {}; + ULONG return_len = 0; + NTSTATUS status = + NtQueryInformationProcess(process, + nt_internals::ProcessBasicInformation, + &basic_info, + sizeof(basic_info), + &return_len); + if (0 != status) + return HRESULT_FROM_NT(status); + else if (return_len != sizeof(basic_info)) + return E_UNEXPECTED; + + LPVOID gflags_location = + reinterpret_cast<char*>(basic_info.PebBaseAddress) + kGFlagsPEBOffset; + + SIZE_T num_written = 0; + if (!::WriteProcessMemory(process, + gflags_location, + &gflags, + sizeof(gflags), + &num_written) || + sizeof(gflags) != num_written) + return HRESULT_FROM_WIN32(::GetLastError()); + + return S_OK; +} + +HRESULT CreateProcessWithGFlags(LPCWSTR application_name, + LPWSTR command_line, + LPSECURITY_ATTRIBUTES process_attributes, + LPSECURITY_ATTRIBUTES thread_attributes, + BOOL inherit_handles, + DWORD creation_flags, + LPVOID environment, + LPCWSTR current_directory, + LPSTARTUPINFOW startup_info, + LPPROCESS_INFORMATION proc_info, + DWORD gflags) { + // Create the process, but make sure it's suspended at creation. + BOOL created = ::CreateProcess(application_name, + command_line, + process_attributes, + thread_attributes, + inherit_handles, + creation_flags | CREATE_SUSPENDED, + environment, + current_directory, + startup_info, + proc_info); + if (!created) + return HRESULT_FROM_WIN32(::GetLastError()); + + // If we succeeded, go through and write the new process' + // PEB with the requested gflags. + HRESULT hr = WriteProcessGFlags(proc_info->hProcess, gflags); + if (FAILED(hr)) { + // We failed to write the gflags value, clean up and bail. + ::TerminateProcess(proc_info->hProcess, 0); + ::CloseHandle(proc_info->hProcess); + ::CloseHandle(proc_info->hThread); + + return hr; + } + + // Resume the main thread unless suspended creation was requested. + if (!(creation_flags & CREATE_SUSPENDED)) + ::ResumeThread(proc_info->hThread); + + return hr; +} + +HRESULT RelaunchWithGFlags(DWORD wanted_gflags) { + DWORD current_gflags = nt_internals::RtlGetNtGlobalFlags(); + + // Do we already have all the wanted_gflags? + if (wanted_gflags == (wanted_gflags & current_gflags)) + return S_FALSE; + + wchar_t app_name[MAX_PATH] = {}; + if (!::GetModuleFileName(NULL, app_name, ARRAYSIZE(app_name))) + return HRESULT_FROM_WIN32(::GetLastError()); + + // If there is a GFlags registry entry for our executable, we don't + // want to relaunch because it's futile. The spawned process would pick + // up the registry settings in preference to ours, and we'd effectively + // have a fork bomb on our hands. + wchar_t key_name[MAX_PATH] = + L"SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\" + L"Image File Execution Options\\"; + + if (0 != wcscat_s(key_name, ::PathFindFileName(app_name))) + return E_UNEXPECTED; + + HKEY key = NULL; + LONG err = ::RegOpenKeyEx(HKEY_LOCAL_MACHINE, key_name, 0, KEY_READ, &key); + if (err == ERROR_SUCCESS) { + // There's a setting for this executable in registry, does it + // specify global flags? + DWORD data_len = 0; + err = ::RegQueryValueEx(key, L"GlobalFlag", NULL, NULL, NULL, &data_len); + + // Close the key. + ::RegCloseKey(key); + + if (err == ERROR_SUCCESS) { + // We can't relaunch because we'd run the risk of a fork bomb. + return E_FAIL; + } + } + + wchar_t current_dir[MAX_PATH] = {}; + if (!::GetCurrentDirectory(ARRAYSIZE(current_dir), current_dir)) + return HRESULT_FROM_WIN32(::GetLastError()); + + STARTUPINFO startup_info = {}; + ::GetStartupInfo(&startup_info); + PROCESS_INFORMATION proc_info = {}; + HRESULT hr; + hr = CreateProcessWithGFlags(app_name, + ::GetCommandLine(), + NULL, + NULL, + FALSE, + 0, + NULL, + current_dir, + &startup_info, + &proc_info, + current_gflags | wanted_gflags); + if (SUCCEEDED(hr)) { + ::WaitForSingleObject(proc_info.hProcess, INFINITE); + DWORD exit_code = 0; + ::GetExitCodeProcess(proc_info.hProcess, &exit_code); + ::ExitProcess(exit_code); + } + + return hr; +} + +} // namespace testing |