aboutsummaryrefslogtreecommitdiff
path: root/src/dutil/rmutil.cpp
diff options
context:
space:
mode:
authorRob Mensching <rob@firegiant.com>2017-09-03 11:22:38 -0700
committerRob Mensching <rob@firegiant.com>2017-09-03 13:33:33 -0700
commit5d8375007754101ff2889d0e79486c8f9b7cf5ab (patch)
treea76d6fb6a38dd9f04a93ffcfd9d64e76779b3414 /src/dutil/rmutil.cpp
parent8e8da6dbc051ec884b5d439bb4f44dc027d05bbf (diff)
downloadwix-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.cpp473
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
8typedef DWORD (WINAPI *PFNRMJOINSESSION)(
9 __out DWORD *pSessionHandle,
10 __in_z const WCHAR strSessionKey[]
11 );
12
13typedef DWORD (WINAPI *PFNRMENDSESSION)(
14 __in DWORD dwSessionHandle
15 );
16
17typedef 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
27typedef 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
45static volatile LONG vcRmuInitialized = 0;
46static HMODULE vhModule = NULL;
47static PFNRMJOINSESSION vpfnRmJoinSession = NULL;
48static PFNRMENDSESSION vpfnRmEndSession = NULL;
49static PFNRMREGISTERRESOURCES vpfnRmRegisterResources = NULL;
50
51static HRESULT RmuInitialize();
52static void RmuUninitialize();
53
54static 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
61static 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/********************************************************************
68RmuJoinSession - Joins an existing Restart Manager session.
69
70********************************************************************/
71extern "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
96LExit:
97 if (FAILED(hr))
98 {
99 ReleaseNullMem(pSession);
100 }
101
102 return hr;
103}
104
105/********************************************************************
106RmuAddFile - Adds the file path to the Restart Manager session.
107
108You should call this multiple times as necessary before calling
109RmuRegisterResources.
110
111********************************************************************/
112extern "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
125LExit:
126 ::LeaveCriticalSection(&pSession->cs);
127 return hr;
128}
129
130/********************************************************************
131RmuAddProcessById - Adds the process ID to the Restart Manager sesion.
132
133You should call this multiple times as necessary before calling
134RmuRegisterResources.
135
136********************************************************************/
137extern "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
220LExit:
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/********************************************************************
243RmuAddProcessesByName - Adds all processes by the given process name
244 to the Restart Manager Session.
245
246You should call this multiple times as necessary before calling
247RmuRegisterResources.
248
249********************************************************************/
250extern "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
283LExit:
284 ReleaseMem(pdwProcessIds);
285
286 return hr;
287}
288
289/********************************************************************
290RmuAddService - Adds the service name to the Restart Manager session.
291
292You should call this multiple times as necessary before calling
293RmuRegisterResources.
294
295********************************************************************/
296extern "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
308LExit:
309 ::LeaveCriticalSection(&pSession->cs);
310 return hr;
311}
312
313/********************************************************************
314RmuRegisterResources - Registers resources for the Restart Manager.
315
316This should be called rarely because it is expensive to run. Call
317functions like RmuAddFile for multiple resources then commit them
318as a batch of updates to RmuRegisterResources.
319
320Duplicate resources appear to be handled by Restart Manager.
321Only one WM_QUERYENDSESSION is being sent for each top-level window.
322
323********************************************************************/
324extern "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
351LExit:
352 ::LeaveCriticalSection(&pSession->cs);
353 return hr;
354}
355
356/********************************************************************
357RmuEndSession - Ends the session.
358
359If the session was joined by RmuJoinSession, any remaining resources
360are registered before the session is ended (left).
361
362********************************************************************/
363extern "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
382LExit:
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
398static 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
421LExit:
422 return hr;
423}
424
425static 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
439static 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
458LExit:
459 return hr;
460}
461
462static 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
471LExit:
472 return hr;
473}