1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
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
|
