@samitouri / QOSAMI-WSL
Code Commits Issues Pulls Sessions CI/CD
master
master
cpp 279 lines 7.91 KB
Raw
1 /*++
2
3 Copyright (c) Microsoft. All rights reserved.
4
5 Module Name:
6
7 SubProcess.cpp
8
9 Abstract:
10
11 This file contains the subprocess helper class implementation.
12
13 --*/
14
15 #include "precomp.h"
16
17 #include "SubProcess.h"
18
19 using namespace wsl::windows::common::io;
20 using wsl::windows::common::SubProcess;
21
22 SubProcess::SubProcess(LPCWSTR ApplicationName, LPCWSTR CommandLine, DWORD Flags, DWORD StartupFlags) :
23 m_applicationName(ApplicationName), m_commandLine(CommandLine), m_flags(Flags), m_startupFlags(StartupFlags)
24 {
25 }
26
27 void SubProcess::SetStdHandles(HANDLE Stdin, HANDLE Stdout, HANDLE Stderr)
28 {
29 m_stdIn = Stdin;
30 m_stdOut = Stdout;
31 m_stdErr = Stderr;
32 }
33
34 void SubProcess::InheritHandle(HANDLE Handle)
35 {
36 // N.B. Trying to inherit the same handle twice will cause CreateProcess to fail with INVALID_ARG.
37 if (std::find(m_inheritHandles.begin(), m_inheritHandles.end(), Handle) == m_inheritHandles.end())
38 {
39 m_inheritHandles.emplace_back(Handle);
40 }
41 }
42
43 void SubProcess::SetPseudoConsole(HPCON Console)
44 {
45 m_pseudoConsole = Console;
46 }
47
48 void SubProcess::SetDesktopAppPolicy(DWORD Policy)
49 {
50 m_desktopAppPolicy = Policy;
51 }
52
53 void SubProcess::SetEnvironment(LPVOID Environment)
54 {
55 m_environment = Environment;
56 }
57
58 void SubProcess::SetWorkingDirectory(LPCWSTR Directory)
59 {
60 m_workingDirectory = Directory;
61 }
62
63 void SubProcess::SetFlags(DWORD Flag)
64 {
65 WI_SetAllFlags(m_flags, Flag);
66 }
67
68 void SubProcess::SetToken(HANDLE Token)
69 {
70 m_token = Token;
71 }
72
73 void SubProcess::SetShowWindow(WORD ShowWindow)
74 {
75 m_showWindow = ShowWindow;
76 }
77
78 void SubProcess::SetJobObject(HANDLE JobObject)
79 {
80 m_jobObject = JobObject;
81 }
82
83 wsl::windows::common::helpers::unique_proc_attribute_list SubProcess::BuildProcessAttributes()
84 {
85 DWORD attributes = 0;
86 if (!m_inheritHandles.empty())
87 {
88 attributes++;
89 }
90
91 if (m_desktopAppPolicy.has_value())
92 {
93 attributes++;
94 }
95
96 if (m_pseudoConsole != nullptr)
97 {
98 attributes++;
99 }
100
101 if (m_jobObject != nullptr)
102 {
103 attributes++;
104 }
105
106 if (attributes == 0)
107 {
108 return {};
109 }
110
111 auto list = helpers::CreateProcThreadAttributeList(attributes);
112
113 // Handles to inherit
114 // N.B. Pseudoconsoles can't be passed to PROC_THREAD_ATTRIBUTE_HANDLE_LIST
115 // so if a pseudoconsole is passed, all handles need to be inherited.
116 if (!m_inheritHandles.empty())
117 {
118 THROW_IF_WIN32_BOOL_FALSE(UpdateProcThreadAttribute(
119 list.get(), 0, PROC_THREAD_ATTRIBUTE_HANDLE_LIST, m_inheritHandles.data(), m_inheritHandles.size() * sizeof(HANDLE), nullptr, nullptr));
120 }
121
122 // Desktop app policy
123 if (m_desktopAppPolicy.has_value())
124 {
125 THROW_IF_WIN32_BOOL_FALSE(UpdateProcThreadAttribute(
126 list.get(), 0, PROC_THREAD_ATTRIBUTE_DESKTOP_APP_POLICY, &m_desktopAppPolicy.value(), sizeof(m_desktopAppPolicy.value()), nullptr, nullptr));
127 }
128
129 // Pseudoconsole
130 if (m_pseudoConsole != nullptr)
131 {
132 THROW_IF_WIN32_BOOL_FALSE(UpdateProcThreadAttribute(
133 list.get(), 0, PROC_THREAD_ATTRIBUTE_PSEUDOCONSOLE, m_pseudoConsole, sizeof(m_pseudoConsole), nullptr, nullptr));
134 }
135
136 // Job object
137 if (m_jobObject != nullptr)
138 {
139 THROW_IF_WIN32_BOOL_FALSE(UpdateProcThreadAttribute(
140 list.get(), 0, PROC_THREAD_ATTRIBUTE_JOB_LIST, &m_jobObject, sizeof(m_jobObject), nullptr, nullptr));
141 }
142
143 return list;
144 }
145
146 wil::unique_handle SubProcess::Start()
147 {
148 WI_SetFlag(m_flags, EXTENDED_STARTUPINFO_PRESENT);
149
150 STARTUPINFOEX StartupInfo{};
151 StartupInfo.StartupInfo.cb = sizeof(StartupInfo);
152 StartupInfo.StartupInfo.dwFlags = STARTF_USESTDHANDLES | m_startupFlags;
153
154 // N.B. Passing a pseudoconsole requires all standard handles to be null
155 if (m_pseudoConsole == nullptr)
156 {
157 StartupInfo.StartupInfo.hStdInput = ARGUMENT_PRESENT(m_stdIn) ? m_stdIn : GetStdHandle(STD_INPUT_HANDLE);
158 StartupInfo.StartupInfo.hStdOutput = ARGUMENT_PRESENT(m_stdOut) ? m_stdOut : GetStdHandle(STD_OUTPUT_HANDLE);
159 StartupInfo.StartupInfo.hStdError = ARGUMENT_PRESENT(m_stdErr) ? m_stdErr : GetStdHandle(STD_ERROR_HANDLE);
160
161 if (StartupInfo.StartupInfo.hStdInput != nullptr)
162 {
163 InheritHandle(StartupInfo.StartupInfo.hStdInput);
164 }
165
166 if (StartupInfo.StartupInfo.hStdOutput != nullptr)
167 {
168 InheritHandle(StartupInfo.StartupInfo.hStdOutput);
169 }
170
171 if (StartupInfo.StartupInfo.hStdError != nullptr)
172 {
173 InheritHandle(StartupInfo.StartupInfo.hStdError);
174 }
175 }
176
177 StartupInfo.StartupInfo.lpDesktop = const_cast<LPWSTR>(m_desktop);
178
179 if (m_showWindow.has_value())
180 {
181 WI_SetFlag(StartupInfo.StartupInfo.dwFlags, STARTF_USESHOWWINDOW);
182 StartupInfo.StartupInfo.wShowWindow = m_showWindow.value();
183 }
184
185 const auto attributes = BuildProcessAttributes();
186 StartupInfo.lpAttributeList = attributes.get();
187
188 wil::unique_process_information processInfo;
189 THROW_IF_WIN32_BOOL_FALSE_MSG(
190 CreateProcessAsUserW(
191 m_token,
192 m_applicationName,
193 m_commandLine.data(),
194 nullptr,
195 nullptr,
196 !m_inheritHandles.empty(),
197 m_flags,
198 m_environment,
199 m_workingDirectory,
200 &StartupInfo.StartupInfo,
201 &processInfo),
202 "ApplicationName: %ls, CommandLine: %ls, WorkingDirectory: %ls",
203 m_applicationName,
204 m_commandLine.c_str(),
205 m_workingDirectory != nullptr ? m_workingDirectory : L"<null>");
206
207 wil::unique_handle createdProcess{processInfo.hProcess};
208
209 // Make sure that the process handle doesn't get closed on return
210 processInfo.hProcess = nullptr;
211
212 return createdProcess;
213 }
214
215 DWORD SubProcess::GetExitCode(HANDLE Process, DWORD Timeout)
216 {
217 const auto status = WaitForSingleObject(Process, Timeout);
218 THROW_HR_IF(HRESULT_FROM_NT(ERROR_TIMEOUT), status == WAIT_TIMEOUT);
219 THROW_LAST_ERROR_IF(status != WAIT_OBJECT_0);
220
221 DWORD exitCode{};
222 THROW_IF_WIN32_BOOL_FALSE(GetExitCodeProcess(Process, &exitCode));
223 return exitCode;
224 }
225
226 DWORD SubProcess::Run(DWORD Timeout)
227 {
228 return GetExitCode(Start().get(), Timeout);
229 }
230
231 SubProcess::ProcessOutput SubProcess::RunAndCaptureOutput(DWORD Timeout, HANDLE StdErr)
232 {
233 auto cleanup = wil::scope_exit_log(WI_DIAGNOSTICS_INFO, [&]() {
234 // Clear out references to stdout and stderr pipes.
235 m_stdOut = nullptr;
236 m_stdErr = nullptr;
237 });
238
239 auto [stdoutRead, stdoutWrite] = wsl::windows::common::wslutil::OpenAnonymousPipe(0, true, false);
240 THROW_IF_WIN32_BOOL_FALSE(SetHandleInformation(stdoutWrite.get(), HANDLE_FLAG_INHERIT, HANDLE_FLAG_INHERIT));
241
242 m_stdOut = stdoutWrite.get();
243
244 io::MultiHandleWait io;
245 std::string stdoutNative;
246 std::string stderrNative;
247
248 io.AddHandle(std::make_unique<io::ReadHandle>(
249 std::move(stdoutRead), [&](const gsl::span<char>& buffer) { stdoutNative.append(buffer.data(), buffer.size()); }));
250
251 wil::unique_hfile stderrWrite;
252 if (StdErr == nullptr)
253 {
254 wil::unique_hfile stderrRead;
255 std::tie(stderrRead, stderrWrite) = wsl::windows::common::wslutil::OpenAnonymousPipe(0, true, false);
256 THROW_IF_WIN32_BOOL_FALSE(SetHandleInformation(stderrWrite.get(), HANDLE_FLAG_INHERIT, HANDLE_FLAG_INHERIT));
257
258 m_stdErr = stderrWrite.get();
259
260 io.AddHandle(std::make_unique<io::ReadHandle>(
261 std::move(stderrRead), [&](const gsl::span<char>& buffer) { stderrNative.append(buffer.data(), buffer.size()); }));
262 }
263 else
264 {
265 m_stdErr = StdErr;
266 }
267
268 auto process = Start();
269 stdoutWrite.reset();
270 stderrWrite.reset();
271
272 io.Run(std::chrono::milliseconds{Timeout});
273
274 // Reusing the same timeout since the std handles have been fully read at that point.
275 const DWORD ExitCode = GetExitCode(process.get(), Timeout);
276 ProcessOutput output{ExitCode, shared::string::MultiByteToWide(stdoutNative), shared::string::MultiByteToWide(stderrNative)};
277
278 return output;
279 }