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