diff options
author | Rob Mensching <rob@firegiant.com> | 2017-09-03 11:22:38 -0700 |
---|---|---|
committer | Rob Mensching <rob@firegiant.com> | 2017-09-03 13:33:33 -0700 |
commit | 5d8375007754101ff2889d0e79486c8f9b7cf5ab (patch) | |
tree | a76d6fb6a38dd9f04a93ffcfd9d64e76779b3414 /src/dutil/rmutil.cpp | |
parent | 8e8da6dbc051ec884b5d439bb4f44dc027d05bbf (diff) | |
download | wix-5d8375007754101ff2889d0e79486c8f9b7cf5ab.tar.gz wix-5d8375007754101ff2889d0e79486c8f9b7cf5ab.tar.bz2 wix-5d8375007754101ff2889d0e79486c8f9b7cf5ab.zip |
Initial commit
Diffstat (limited to 'src/dutil/rmutil.cpp')
-rw-r--r-- | src/dutil/rmutil.cpp | 473 |
1 files changed, 473 insertions, 0 deletions
diff --git a/src/dutil/rmutil.cpp b/src/dutil/rmutil.cpp new file mode 100644 index 00000000..75d3e277 --- /dev/null +++ b/src/dutil/rmutil.cpp | |||
@@ -0,0 +1,473 @@ | |||
1 | // Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. | ||
2 | |||
3 | #include "precomp.h" | ||
4 | #include <restartmanager.h> | ||
5 | |||
6 | #define ARRAY_GROWTH_SIZE 5 | ||
7 | |||
8 | typedef DWORD (WINAPI *PFNRMJOINSESSION)( | ||
9 | __out DWORD *pSessionHandle, | ||
10 | __in_z const WCHAR strSessionKey[] | ||
11 | ); | ||
12 | |||
13 | typedef DWORD (WINAPI *PFNRMENDSESSION)( | ||
14 | __in DWORD dwSessionHandle | ||
15 | ); | ||
16 | |||
17 | typedef DWORD (WINAPI *PFNRMREGISTERRESOURCES)( | ||
18 | __in DWORD dwSessionHandle, | ||
19 | __in UINT nFiles, | ||
20 | __in_z_opt LPWSTR *rgsFilenames, | ||
21 | __in UINT nApplications, | ||
22 | __in_opt RM_UNIQUE_PROCESS *rgApplications, | ||
23 | __in UINT nServices, | ||
24 | __in_z_opt LPWSTR *rgsServiceNames | ||
25 | ); | ||
26 | |||
27 | typedef struct _RMU_SESSION | ||
28 | { | ||
29 | CRITICAL_SECTION cs; | ||
30 | DWORD dwSessionHandle; | ||
31 | BOOL fStartedSessionHandle; | ||
32 | BOOL fInitialized; | ||
33 | |||
34 | UINT cFilenames; | ||
35 | LPWSTR *rgsczFilenames; | ||
36 | |||
37 | UINT cApplications; | ||
38 | RM_UNIQUE_PROCESS *rgApplications; | ||
39 | |||
40 | UINT cServiceNames; | ||
41 | LPWSTR *rgsczServiceNames; | ||
42 | |||
43 | } RMU_SESSION; | ||
44 | |||
45 | static volatile LONG vcRmuInitialized = 0; | ||
46 | static HMODULE vhModule = NULL; | ||
47 | static PFNRMJOINSESSION vpfnRmJoinSession = NULL; | ||
48 | static PFNRMENDSESSION vpfnRmEndSession = NULL; | ||
49 | static PFNRMREGISTERRESOURCES vpfnRmRegisterResources = NULL; | ||
50 | |||
51 | static HRESULT RmuInitialize(); | ||
52 | static void RmuUninitialize(); | ||
53 | |||
54 | static HRESULT RmuApplicationArrayAlloc( | ||
55 | __deref_inout_ecount(*pcApplications) RM_UNIQUE_PROCESS **prgApplications, | ||
56 | __inout LPUINT pcApplications, | ||
57 | __in DWORD dwProcessId, | ||
58 | __in FILETIME ProcessStartTime | ||
59 | ); | ||
60 | |||
61 | static HRESULT RmuApplicationArrayFree( | ||
62 | __in RM_UNIQUE_PROCESS *rgApplications | ||
63 | ); | ||
64 | |||
65 | #define ReleaseNullApplicationArray(rg, c) { if (rg) { RmuApplicationArrayFree(rg); c = 0; rg = NULL; } } | ||
66 | |||
67 | /******************************************************************** | ||
68 | RmuJoinSession - Joins an existing Restart Manager session. | ||
69 | |||
70 | ********************************************************************/ | ||
71 | extern "C" HRESULT DAPI RmuJoinSession( | ||
72 | __out PRMU_SESSION *ppSession, | ||
73 | __in_z LPCWSTR wzSessionKey | ||
74 | ) | ||
75 | { | ||
76 | HRESULT hr = S_OK; | ||
77 | DWORD er = ERROR_SUCCESS; | ||
78 | PRMU_SESSION pSession = NULL; | ||
79 | |||
80 | *ppSession = NULL; | ||
81 | |||
82 | pSession = static_cast<PRMU_SESSION>(MemAlloc(sizeof(RMU_SESSION), TRUE)); | ||
83 | ExitOnNull(pSession, hr, E_OUTOFMEMORY, "Failed to allocate the RMU_SESSION structure."); | ||
84 | |||
85 | hr = RmuInitialize(); | ||
86 | ExitOnFailure(hr, "Failed to initialize Restart Manager."); | ||
87 | |||
88 | er = vpfnRmJoinSession(&pSession->dwSessionHandle, wzSessionKey); | ||
89 | ExitOnWin32Error(er, hr, "Failed to join Restart Manager session %ls.", wzSessionKey); | ||
90 | |||
91 | ::InitializeCriticalSection(&pSession->cs); | ||
92 | pSession->fInitialized = TRUE; | ||
93 | |||
94 | *ppSession = pSession; | ||
95 | |||
96 | LExit: | ||
97 | if (FAILED(hr)) | ||
98 | { | ||
99 | ReleaseNullMem(pSession); | ||
100 | } | ||
101 | |||
102 | return hr; | ||
103 | } | ||
104 | |||
105 | /******************************************************************** | ||
106 | RmuAddFile - Adds the file path to the Restart Manager session. | ||
107 | |||
108 | You should call this multiple times as necessary before calling | ||
109 | RmuRegisterResources. | ||
110 | |||
111 | ********************************************************************/ | ||
112 | extern "C" HRESULT DAPI RmuAddFile( | ||
113 | __in PRMU_SESSION pSession, | ||
114 | __in_z LPCWSTR wzPath | ||
115 | ) | ||
116 | { | ||
117 | HRESULT hr = S_OK; | ||
118 | |||
119 | ::EnterCriticalSection(&pSession->cs); | ||
120 | |||
121 | // Create or grow the jagged array. | ||
122 | hr = StrArrayAllocString(&pSession->rgsczFilenames, &pSession->cFilenames, wzPath, 0); | ||
123 | ExitOnFailure(hr, "Failed to add the filename to the array."); | ||
124 | |||
125 | LExit: | ||
126 | ::LeaveCriticalSection(&pSession->cs); | ||
127 | return hr; | ||
128 | } | ||
129 | |||
130 | /******************************************************************** | ||
131 | RmuAddProcessById - Adds the process ID to the Restart Manager sesion. | ||
132 | |||
133 | You should call this multiple times as necessary before calling | ||
134 | RmuRegisterResources. | ||
135 | |||
136 | ********************************************************************/ | ||
137 | extern "C" HRESULT DAPI RmuAddProcessById( | ||
138 | __in PRMU_SESSION pSession, | ||
139 | __in DWORD dwProcessId | ||
140 | ) | ||
141 | { | ||
142 | HRESULT hr = S_OK; | ||
143 | HANDLE hProcess = NULL; | ||
144 | FILETIME CreationTime = {}; | ||
145 | FILETIME ExitTime = {}; | ||
146 | FILETIME KernelTime = {}; | ||
147 | FILETIME UserTime = {}; | ||
148 | BOOL fLocked = FALSE; | ||
149 | |||
150 | HANDLE hToken = NULL; | ||
151 | TOKEN_PRIVILEGES priv = { 0 }; | ||
152 | TOKEN_PRIVILEGES* pPrevPriv = NULL; | ||
153 | DWORD cbPrevPriv = 0; | ||
154 | DWORD er = ERROR_SUCCESS; | ||
155 | BOOL fAdjustedPrivileges = FALSE; | ||
156 | BOOL fElevated = FALSE; | ||
157 | ProcElevated(::GetCurrentProcess(), &fElevated); | ||
158 | |||
159 | // Must be elevated to adjust process privileges | ||
160 | if (fElevated) { | ||
161 | // Adding SeDebugPrivilege in the event that the process targeted by ::OpenProcess() is in a another user context. | ||
162 | if (!::OpenProcessToken(::GetCurrentProcess(), TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, &hToken)) | ||
163 | { | ||
164 | ExitWithLastError(hr, "Failed to get process token."); | ||
165 | } | ||
166 | |||
167 | priv.PrivilegeCount = 1; | ||
168 | priv.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED; | ||
169 | if (!::LookupPrivilegeValueW(NULL, L"SeDebugPrivilege", &priv.Privileges[0].Luid)) | ||
170 | { | ||
171 | ExitWithLastError(hr, "Failed to get debug privilege LUID."); | ||
172 | } | ||
173 | |||
174 | cbPrevPriv = sizeof(TOKEN_PRIVILEGES); | ||
175 | pPrevPriv = static_cast<TOKEN_PRIVILEGES*>(MemAlloc(cbPrevPriv, TRUE)); | ||
176 | ExitOnNull(pPrevPriv, hr, E_OUTOFMEMORY, "Failed to allocate memory for empty previous privileges."); | ||
177 | |||
178 | if (!::AdjustTokenPrivileges(hToken, FALSE, &priv, cbPrevPriv, pPrevPriv, &cbPrevPriv)) | ||
179 | { | ||
180 | LPVOID pv = MemReAlloc(pPrevPriv, cbPrevPriv, TRUE); | ||
181 | ExitOnNull(pv, hr, E_OUTOFMEMORY, "Failed to allocate memory for previous privileges."); | ||
182 | pPrevPriv = static_cast<TOKEN_PRIVILEGES*>(pv); | ||
183 | |||
184 | if (!::AdjustTokenPrivileges(hToken, FALSE, &priv, cbPrevPriv, pPrevPriv, &cbPrevPriv)) | ||
185 | { | ||
186 | ExitWithLastError(hr, "Failed to get debug privilege LUID."); | ||
187 | } | ||
188 | } | ||
189 | |||
190 | fAdjustedPrivileges = TRUE; | ||
191 | } | ||
192 | |||
193 | hProcess = ::OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, dwProcessId); | ||
194 | if (hProcess) | ||
195 | { | ||
196 | if (!::GetProcessTimes(hProcess, &CreationTime, &ExitTime, &KernelTime, &UserTime)) | ||
197 | { | ||
198 | ExitWithLastError(hr, "Failed to get the process times for process ID %d.", dwProcessId); | ||
199 | } | ||
200 | |||
201 | ::EnterCriticalSection(&pSession->cs); | ||
202 | fLocked = TRUE; | ||
203 | hr = RmuApplicationArrayAlloc(&pSession->rgApplications, &pSession->cApplications, dwProcessId, CreationTime); | ||
204 | ExitOnFailure(hr, "Failed to add the application to the array."); | ||
205 | } | ||
206 | else | ||
207 | { | ||
208 | er = ::GetLastError(); | ||
209 | if (ERROR_ACCESS_DENIED == er) | ||
210 | { | ||
211 | // OpenProcess will fail when not elevated and the target process is in another user context. Let the caller log and continue. | ||
212 | hr = E_NOTFOUND; | ||
213 | } | ||
214 | else | ||
215 | { | ||
216 | ExitOnWin32Error(er, hr, "Failed to open the process ID %d.", dwProcessId); | ||
217 | } | ||
218 | } | ||
219 | |||
220 | LExit: | ||
221 | if (hProcess) | ||
222 | { | ||
223 | ::CloseHandle(hProcess); | ||
224 | } | ||
225 | |||
226 | if (fAdjustedPrivileges) | ||
227 | { | ||
228 | ::AdjustTokenPrivileges(hToken, FALSE, pPrevPriv, 0, NULL, NULL); | ||
229 | } | ||
230 | |||
231 | ReleaseMem(pPrevPriv); | ||
232 | ReleaseHandle(hToken); | ||
233 | |||
234 | if (fLocked) | ||
235 | { | ||
236 | ::LeaveCriticalSection(&pSession->cs); | ||
237 | } | ||
238 | |||
239 | return hr; | ||
240 | } | ||
241 | |||
242 | /******************************************************************** | ||
243 | RmuAddProcessesByName - Adds all processes by the given process name | ||
244 | to the Restart Manager Session. | ||
245 | |||
246 | You should call this multiple times as necessary before calling | ||
247 | RmuRegisterResources. | ||
248 | |||
249 | ********************************************************************/ | ||
250 | extern "C" HRESULT DAPI RmuAddProcessesByName( | ||
251 | __in PRMU_SESSION pSession, | ||
252 | __in_z LPCWSTR wzProcessName | ||
253 | ) | ||
254 | { | ||
255 | HRESULT hr = S_OK; | ||
256 | DWORD *pdwProcessIds = NULL; | ||
257 | DWORD cProcessIds = 0; | ||
258 | BOOL fNotFound = FALSE; | ||
259 | |||
260 | hr = ProcFindAllIdsFromExeName(wzProcessName, &pdwProcessIds, &cProcessIds); | ||
261 | ExitOnFailure(hr, "Failed to enumerate all the processes by name %ls.", wzProcessName); | ||
262 | |||
263 | for (DWORD i = 0; i < cProcessIds; ++i) | ||
264 | { | ||
265 | hr = RmuAddProcessById(pSession, pdwProcessIds[i]); | ||
266 | if (E_NOTFOUND == hr) | ||
267 | { | ||
268 | // RmuAddProcessById returns E_NOTFOUND when this setup is not elevated and OpenProcess returned access denied (target process running under another user account). | ||
269 | fNotFound = TRUE; | ||
270 | } | ||
271 | else | ||
272 | { | ||
273 | ExitOnFailure(hr, "Failed to add process %ls (%d) to the Restart Manager session.", wzProcessName, pdwProcessIds[i]); | ||
274 | } | ||
275 | } | ||
276 | |||
277 | // If one or more calls to RmuAddProcessById returned E_NOTFOUND, then return E_NOTFOUND even if other calls succeeded, so that caller can log the issue. | ||
278 | if (fNotFound) | ||
279 | { | ||
280 | hr = E_NOTFOUND; | ||
281 | } | ||
282 | |||
283 | LExit: | ||
284 | ReleaseMem(pdwProcessIds); | ||
285 | |||
286 | return hr; | ||
287 | } | ||
288 | |||
289 | /******************************************************************** | ||
290 | RmuAddService - Adds the service name to the Restart Manager session. | ||
291 | |||
292 | You should call this multiple times as necessary before calling | ||
293 | RmuRegisterResources. | ||
294 | |||
295 | ********************************************************************/ | ||
296 | extern "C" HRESULT DAPI RmuAddService( | ||
297 | __in PRMU_SESSION pSession, | ||
298 | __in_z LPCWSTR wzServiceName | ||
299 | ) | ||
300 | { | ||
301 | HRESULT hr = S_OK; | ||
302 | |||
303 | ::EnterCriticalSection(&pSession->cs); | ||
304 | |||
305 | hr = StrArrayAllocString(&pSession->rgsczServiceNames, &pSession->cServiceNames, wzServiceName, 0); | ||
306 | ExitOnFailure(hr, "Failed to add the service name to the array."); | ||
307 | |||
308 | LExit: | ||
309 | ::LeaveCriticalSection(&pSession->cs); | ||
310 | return hr; | ||
311 | } | ||
312 | |||
313 | /******************************************************************** | ||
314 | RmuRegisterResources - Registers resources for the Restart Manager. | ||
315 | |||
316 | This should be called rarely because it is expensive to run. Call | ||
317 | functions like RmuAddFile for multiple resources then commit them | ||
318 | as a batch of updates to RmuRegisterResources. | ||
319 | |||
320 | Duplicate resources appear to be handled by Restart Manager. | ||
321 | Only one WM_QUERYENDSESSION is being sent for each top-level window. | ||
322 | |||
323 | ********************************************************************/ | ||
324 | extern "C" HRESULT DAPI RmuRegisterResources( | ||
325 | __in PRMU_SESSION pSession | ||
326 | ) | ||
327 | { | ||
328 | HRESULT hr = S_OK; | ||
329 | DWORD er = ERROR_SUCCESS; | ||
330 | |||
331 | AssertSz(vcRmuInitialized, "Restart Manager was not properly initialized."); | ||
332 | |||
333 | ::EnterCriticalSection(&pSession->cs); | ||
334 | |||
335 | er = vpfnRmRegisterResources( | ||
336 | pSession->dwSessionHandle, | ||
337 | pSession->cFilenames, | ||
338 | pSession->rgsczFilenames, | ||
339 | pSession->cApplications, | ||
340 | pSession->rgApplications, | ||
341 | pSession->cServiceNames, | ||
342 | pSession->rgsczServiceNames | ||
343 | ); | ||
344 | ExitOnWin32Error(er, hr, "Failed to register the resources with the Restart Manager session."); | ||
345 | |||
346 | // Empty the arrays if registered in case additional resources are added later. | ||
347 | ReleaseNullStrArray(pSession->rgsczFilenames, pSession->cFilenames); | ||
348 | ReleaseNullApplicationArray(pSession->rgApplications, pSession->cApplications); | ||
349 | ReleaseNullStrArray(pSession->rgsczServiceNames, pSession->cServiceNames); | ||
350 | |||
351 | LExit: | ||
352 | ::LeaveCriticalSection(&pSession->cs); | ||
353 | return hr; | ||
354 | } | ||
355 | |||
356 | /******************************************************************** | ||
357 | RmuEndSession - Ends the session. | ||
358 | |||
359 | If the session was joined by RmuJoinSession, any remaining resources | ||
360 | are registered before the session is ended (left). | ||
361 | |||
362 | ********************************************************************/ | ||
363 | extern "C" HRESULT DAPI RmuEndSession( | ||
364 | __in PRMU_SESSION pSession | ||
365 | ) | ||
366 | { | ||
367 | HRESULT hr = S_OK; | ||
368 | DWORD er = ERROR_SUCCESS; | ||
369 | |||
370 | AssertSz(vcRmuInitialized, "Restart Manager was not properly initialized."); | ||
371 | |||
372 | // Make sure all resources are registered if we joined the session. | ||
373 | if (!pSession->fStartedSessionHandle) | ||
374 | { | ||
375 | hr = RmuRegisterResources(pSession); | ||
376 | ExitOnFailure(hr, "Failed to register remaining resources."); | ||
377 | } | ||
378 | |||
379 | er = vpfnRmEndSession(pSession->dwSessionHandle); | ||
380 | ExitOnWin32Error(er, hr, "Failed to end the Restart Manager session."); | ||
381 | |||
382 | LExit: | ||
383 | if (pSession->fInitialized) | ||
384 | { | ||
385 | ::DeleteCriticalSection(&pSession->cs); | ||
386 | } | ||
387 | |||
388 | ReleaseNullStrArray(pSession->rgsczFilenames, pSession->cFilenames); | ||
389 | ReleaseNullApplicationArray(pSession->rgApplications, pSession->cApplications); | ||
390 | ReleaseNullStrArray(pSession->rgsczServiceNames, pSession->cServiceNames); | ||
391 | ReleaseNullMem(pSession); | ||
392 | |||
393 | RmuUninitialize(); | ||
394 | |||
395 | return hr; | ||
396 | } | ||
397 | |||
398 | static HRESULT RmuInitialize() | ||
399 | { | ||
400 | HRESULT hr = S_OK; | ||
401 | HMODULE hModule = NULL; | ||
402 | |||
403 | LONG iRef = ::InterlockedIncrement(&vcRmuInitialized); | ||
404 | if (1 == iRef && !vhModule) | ||
405 | { | ||
406 | hr = LoadSystemLibrary(L"rstrtmgr.dll", &hModule); | ||
407 | ExitOnFailure(hr, "Failed to load the rstrtmgr.dll module."); | ||
408 | |||
409 | vpfnRmJoinSession = reinterpret_cast<PFNRMJOINSESSION>(::GetProcAddress(hModule, "RmJoinSession")); | ||
410 | ExitOnNullWithLastError(vpfnRmJoinSession, hr, "Failed to get the RmJoinSession procedure from rstrtmgr.dll."); | ||
411 | |||
412 | vpfnRmRegisterResources = reinterpret_cast<PFNRMREGISTERRESOURCES>(::GetProcAddress(hModule, "RmRegisterResources")); | ||
413 | ExitOnNullWithLastError(vpfnRmRegisterResources, hr, "Failed to get the RmRegisterResources procedure from rstrtmgr.dll."); | ||
414 | |||
415 | vpfnRmEndSession = reinterpret_cast<PFNRMENDSESSION>(::GetProcAddress(hModule, "RmEndSession")); | ||
416 | ExitOnNullWithLastError(vpfnRmEndSession, hr, "Failed to get the RmEndSession procedure from rstrtmgr.dll."); | ||
417 | |||
418 | vhModule = hModule; | ||
419 | } | ||
420 | |||
421 | LExit: | ||
422 | return hr; | ||
423 | } | ||
424 | |||
425 | static void RmuUninitialize() | ||
426 | { | ||
427 | LONG iRef = ::InterlockedDecrement(&vcRmuInitialized); | ||
428 | if (0 == iRef && vhModule) | ||
429 | { | ||
430 | vpfnRmJoinSession = NULL; | ||
431 | vpfnRmEndSession = NULL; | ||
432 | vpfnRmRegisterResources = NULL; | ||
433 | |||
434 | ::FreeLibrary(vhModule); | ||
435 | vhModule = NULL; | ||
436 | } | ||
437 | } | ||
438 | |||
439 | static HRESULT RmuApplicationArrayAlloc( | ||
440 | __deref_inout_ecount(*pcApplications) RM_UNIQUE_PROCESS **prgApplications, | ||
441 | __inout LPUINT pcApplications, | ||
442 | __in DWORD dwProcessId, | ||
443 | __in FILETIME ProcessStartTime | ||
444 | ) | ||
445 | { | ||
446 | HRESULT hr = S_OK; | ||
447 | RM_UNIQUE_PROCESS *pApplication = NULL; | ||
448 | |||
449 | hr = MemEnsureArraySize(reinterpret_cast<LPVOID*>(prgApplications), *pcApplications + 1, sizeof(RM_UNIQUE_PROCESS), ARRAY_GROWTH_SIZE); | ||
450 | ExitOnFailure(hr, "Failed to allocate memory for the application array."); | ||
451 | |||
452 | pApplication = static_cast<RM_UNIQUE_PROCESS*>(&(*prgApplications)[*pcApplications]); | ||
453 | pApplication->dwProcessId = dwProcessId; | ||
454 | pApplication->ProcessStartTime = ProcessStartTime; | ||
455 | |||
456 | ++(*pcApplications); | ||
457 | |||
458 | LExit: | ||
459 | return hr; | ||
460 | } | ||
461 | |||
462 | static HRESULT RmuApplicationArrayFree( | ||
463 | __in RM_UNIQUE_PROCESS *rgApplications | ||
464 | ) | ||
465 | { | ||
466 | HRESULT hr = S_OK; | ||
467 | |||
468 | hr = MemFree(rgApplications); | ||
469 | ExitOnFailure(hr, "Failed to free memory for the application array."); | ||
470 | |||
471 | LExit: | ||
472 | return hr; | ||
473 | } | ||