Delay closing Burn's UI thread so that it can log the bundle's restart.
Sean Hall committed
Sep 2, 2022 at 16:11 UTC
eea6121f388197435529922b3cb13d3631afb9a8
5 files changed
+159
-40
src/burn/engine/core.cpp
+15
@@ -1963,6 +1963,21 @@ LExit:
1963
return hr;
1964
}
1965
1966
+extern "C" void CoreUpdateRestartState(
1967
+ __in BURN_ENGINE_STATE* pEngineState,
1968
+ __in BURN_RESTART_STATE restartState
1969
+ )
1970
+{
1971
+ ::EnterCriticalSection(&pEngineState->csRestartState);
1972
+
1973
+ if (pEngineState->fRestarting && restartState > pEngineState->restartState)
1974
+ {
1975
+ pEngineState->restartState = restartState;
1976
+ }
1977
+
1978
+ ::LeaveCriticalSection(&pEngineState->csRestartState);
1979
+}
1980
+
1981
extern "C" void CoreFunctionOverride(
1982
__in_opt PFN_CREATEPROCESSW pfnCreateProcessW,
1983
__in_opt PFN_PROCWAITFORCOMPLETION pfnProcWaitForCompletion
src/burn/engine/core.h
+18
@@ -82,6 +82,16 @@ enum BURN_AU_PAUSE_ACTION
82
BURN_AU_PAUSE_ACTION_IFELEVATED_NORESUME,
83
};
84
85
+enum BURN_RESTART_STATE
86
+{
87
+ BURN_RESTART_STATE_NONE,
88
+ BURN_RESTART_STATE_REQUESTING,
89
+ BURN_RESTART_STATE_REQUESTED,
90
+ BURN_RESTART_STATE_INITIATING,
91
+ BURN_RESTART_STATE_INITIATED,
92
+ BURN_RESTART_STATE_BLOCKED,
93
+};
94
+
95
96
// structs
97
@@ -160,6 +170,10 @@ typedef struct _BURN_ENGINE_STATE
170
BURN_PIPE_CONNECTION companionConnection;
171
BURN_PIPE_CONNECTION embeddedConnection;
172
173
+ CRITICAL_SECTION csRestartState;
174
+ BOOL fRestarting;
175
+ BURN_RESTART_STATE restartState;
176
+
177
BOOL fCriticalShutdownInitiated;
178
BURN_RESUME_MODE resumeMode;
179
LPCWSTR wzRestartInitiatedPackageId;
@@ -300,6 +314,10 @@ HRESULT CoreParseCommandLine(
314
__inout HANDLE* phSectionFile,
315
__inout HANDLE* phSourceEngineFile
316
);
317
+void CoreUpdateRestartState(
318
+ __in BURN_ENGINE_STATE* pEngineState,
319
+ __in BURN_RESTART_STATE restartState
320
+ );
321
void CoreFunctionOverride(
322
__in_opt PFN_CREATEPROCESSW pfnCreateProcessW,
323
__in_opt PFN_PROCWAITFORCOMPLETION pfnProcWaitForCompletion
src/burn/engine/engine.cpp
+62
-23
@@ -72,7 +72,9 @@ static HRESULT WaitForElevatedLoggingThread(
72
static HRESULT WaitForUnelevatedLoggingThread(
73
__in HANDLE hUnelevatedLoggingThread
74
);
75
-static HRESULT Restart();
75
+static HRESULT Restart(
76
+ __in BURN_ENGINE_STATE* pEngineState
77
+ );
78
static void CALLBACK BurnTraceError(
79
__in_z LPCSTR szFile,
80
__in int iLine,
@@ -131,7 +133,6 @@ extern "C" HRESULT EngineRun(
133
BOOL fRunNormal = FALSE;
134
BOOL fRunElevated = FALSE;
135
BOOL fRunRunOnce = FALSE;
134
- BOOL fRestart = FALSE;
136
137
BURN_ENGINE_STATE engineState = { };
138
engineState.command.cbSize = sizeof(BOOTSTRAPPER_COMMAND);
@@ -269,9 +270,7 @@ extern "C" HRESULT EngineRun(
270
ExitOnFailure(hr, "Invalid run mode.");
271
}
272
272
- // set exit code and remember if we are supposed to restart.
273
*pdwExitCode = engineState.userExperience.dwExitCode;
274
- fRestart = engineState.fRestart;
274
275
LExit:
276
ReleaseStr(sczExePath);
@@ -289,17 +288,17 @@ LExit:
288
CacheUninitialize(&engineState.cache);
289
290
// If this is a related bundle (but not an update) suppress restart and return the standard restart error code.
292
- if (fRestart && BOOTSTRAPPER_RELATION_NONE != engineState.command.relationType && BOOTSTRAPPER_RELATION_UPDATE != engineState.command.relationType)
291
+ if (engineState.fRestart && BOOTSTRAPPER_RELATION_NONE != engineState.command.relationType && BOOTSTRAPPER_RELATION_UPDATE != engineState.command.relationType)
292
{
293
LogId(REPORT_STANDARD, MSG_RESTART_ABORTED, LoggingRelationTypeToString(engineState.command.relationType));
294
296
- fRestart = FALSE;
295
+ engineState.fRestart = FALSE;
296
hr = SUCCEEDED(hr) ? HRESULT_FROM_WIN32(ERROR_SUCCESS_REBOOT_REQUIRED) : HRESULT_FROM_WIN32(ERROR_FAIL_REBOOT_REQUIRED);
297
}
298
299
if (fRunNormal)
300
{
302
- LogId(REPORT_STANDARD, MSG_EXITING, FAILED(hr) ? (int)hr : *pdwExitCode, LoggingBoolToString(fRestart));
301
+ LogId(REPORT_STANDARD, MSG_EXITING, FAILED(hr) ? (int)hr : *pdwExitCode, LoggingBoolToString(engineState.fRestart));
302
}
303
else if (fRunUntrusted)
304
{
@@ -321,17 +320,20 @@ LExit:
320
LogFlush();
321
}
322
324
- if (fRestart)
323
+ if (engineState.fRestart)
324
{
325
LogId(REPORT_STANDARD, MSG_RESTARTING);
326
328
- HRESULT hrRestart = Restart();
327
+ HRESULT hrRestart = Restart(&engineState);
328
if (FAILED(hrRestart))
329
{
330
LogErrorId(hrRestart, MSG_RESTART_FAILED);
331
}
332
}
333
334
+ // If the message window is still around, close it.
335
+ UiCloseMessageWindow(&engineState);
336
+
337
UninitializeEngineState(&engineState);
338
339
if (fXmlInitialized)
@@ -385,6 +387,8 @@ static HRESULT InitializeEngineState(
387
HANDLE hSectionFile = hEngineFile;
388
HANDLE hSourceEngineFile = INVALID_HANDLE_VALUE;
389
390
+ ::InitializeCriticalSection(&pEngineState->csRestartState);
391
+
392
pEngineState->internalCommand.automaticUpdates = BURN_AU_PAUSE_ACTION_IFELEVATED;
393
::InitializeCriticalSection(&pEngineState->userExperience.csEngineActive);
394
PipeConnectionInitialize(&pEngineState->companionConnection);
@@ -459,6 +463,8 @@ static void UninitializeEngineState(
463
ReleaseStr(pEngineState->log.sczPath);
464
ReleaseStr(pEngineState->log.sczPathVariable);
465
466
+ ::DeleteCriticalSection(&pEngineState->csRestartState);
467
+
468
// clear struct
469
memset(pEngineState, 0, sizeof(BURN_ENGINE_STATE));
470
}
@@ -635,9 +641,6 @@ LExit:
641
642
BurnExtensionUnload(&pEngineState->extensions);
643
638
- // If the message window is still around, close it.
639
- UiCloseMessageWindow(pEngineState);
640
-
644
VariablesDump(&pEngineState->variables);
645
646
// end per-machine process if running
@@ -720,9 +723,6 @@ LExit:
723
724
LogRedirect(NULL, NULL); // we're done talking to the child so always reset logging now.
725
723
- // If the message window is still around, close it.
724
- UiCloseMessageWindow(pEngineState);
725
-
726
if (fDeleteLoggingCs)
727
{
728
::DeleteCriticalSection(&loggingContext.csBuffer);
@@ -1112,7 +1112,9 @@ LExit:
1112
return hr;
1113
}
1114
1115
-static HRESULT Restart()
1115
+static HRESULT Restart(
1116
+ __in BURN_ENGINE_STATE* pEngineState
1117
+ )
1118
{
1119
HRESULT hr = S_OK;
1120
HANDLE hProcessToken = NULL;
@@ -1136,17 +1138,19 @@ static HRESULT Restart()
1138
ExitWithLastError(hr, "Failed to adjust token to add shutdown privileges.");
1139
}
1140
1141
+ pEngineState->fRestarting = TRUE;
1142
+ CoreUpdateRestartState(pEngineState, BURN_RESTART_STATE_REQUESTING);
1143
+
1144
do
1145
{
1146
hr = S_OK;
1147
1143
- // Wait a second to let the companion process (assuming we did an elevated install) to get to the
1144
- // point where it too is thinking about restarting the computer. Only one will schedule the restart
1145
- // but both will have their log files closed and otherwise be ready to exit.
1146
- //
1147
- // On retry, we'll also wait a second to let the OS try to get to a place where the restart can
1148
- // be initiated.
1149
- ::Sleep(1000);
1148
+ if (dwRetries)
1149
+ {
1150
+ // On retry, wait a second to let the OS try to get to a place where the restart can
1151
+ // be initiated.
1152
+ ::Sleep(1000);
1153
+ }
1154
1155
if (!vpfnInitiateSystemShutdownExW(NULL, NULL, 0, FALSE, TRUE, SHTDN_REASON_MAJOR_APPLICATION | SHTDN_REASON_MINOR_INSTALLATION | SHTDN_REASON_FLAG_PLANNED))
1156
{
@@ -1155,6 +1159,41 @@ static HRESULT Restart()
1159
} while (dwRetries++ < RESTART_RETRIES && (HRESULT_FROM_WIN32(ERROR_MACHINE_LOCKED) == hr || HRESULT_FROM_WIN32(ERROR_NOT_READY) == hr));
1160
ExitOnRootFailure(hr, "Failed to schedule restart.");
1161
1162
+ CoreUpdateRestartState(pEngineState, BURN_RESTART_STATE_REQUESTED);
1163
+
1164
+ // Give the UI thread approximately 15 seconds to get the WM_QUERYENDSESSION message.
1165
+ for (DWORD i = 0; i < 60; ++i)
1166
+ {
1167
+ if (!::IsWindow(pEngineState->hMessageWindow))
1168
+ {
1169
+ ExitFunction();
1170
+ }
1171
+
1172
+ if (BURN_RESTART_STATE_REQUESTED < pEngineState->restartState)
1173
+ {
1174
+ break;
1175
+ }
1176
+
1177
+ ::Sleep(250);
1178
+ }
1179
+
1180
+ if (BURN_RESTART_STATE_INITIATING > pEngineState->restartState)
1181
+ {
1182
+ LogId(REPORT_WARNING, MSG_RESTART_BLOCKED);
1183
+ ExitFunction();
1184
+ }
1185
+
1186
+ // Give the UI thread the chance to process the WM_ENDSESSION message.
1187
+ for (;;)
1188
+ {
1189
+ if (!::IsWindow(pEngineState->hMessageWindow) || BURN_RESTART_STATE_INITIATING < pEngineState->restartState)
1190
+ {
1191
+ break;
1192
+ }
1193
+
1194
+ ::Sleep(250);
1195
+ }
1196
+
1197
LExit:
1198
ReleaseHandle(hProcessToken);
1199
return hr;
src/burn/engine/engine.mc
+16
-3
@@ -62,7 +62,6 @@ Severity=Warning
62
SymbolicName=MSG_RESTARTING
63
Language=English
64
Restarting computer...
65
-=======================================
65
.
66
67
MessageId=6
@@ -191,6 +190,13 @@ Language=English
190
The restart request failed, error: %1!ls!. The machine will need to be manually restarted.
191
.
192
193
+MessageId=24
194
+Severity=Error
195
+SymbolicName=MSG_RESTART_BLOCKED
196
+Language=English
197
+The restart request was successful, but no system restart messages have been received. This may be caused by another application blocking the restart or taking too long to respond. The machine might need to be manually restarted.
198
+.
199
+
200
MessageId=51
201
Severity=Error
202
SymbolicName=MSG_FAILED_PARSE_CONDITION
@@ -1138,9 +1144,16 @@ Apply complete, result: 0x%1!x!, restart: %2!hs!, ba requested restart: %3!hs!
1144
1145
MessageId=400
1146
Severity=Success
1141
-SymbolicName=MSG_SYSTEM_SHUTDOWN
1147
+SymbolicName=MSG_SYSTEM_SHUTDOWN_REQUEST
1148
+Language=English
1149
+Received system request to shut down the process: allowed: %1!hs!, elevated: %2!hs!, critical: %3!hs!, logoff: %4!hs!, close app: %5!hs!
1150
+.
1151
+
1152
+MessageId=401
1153
+Severity=Success
1154
+SymbolicName=MSG_SYSTEM_SHUTDOWN_RESULT
1155
Language=English
1143
-Received system request to shut down the process: critical: %1!hs!, elevated: %2!hs!, allowed: %3!hs!
1156
+Received result of system request to shut down the process: closing: %1!hs!, elevated: %2!hs!, critical: %3!hs!, logoff: %4!hs!, close app: %5!hs!
1157
.
1158
1159
MessageId=410
src/burn/engine/uithread.cpp
+48
-14
@@ -46,6 +46,11 @@ HRESULT UiCreateMessageWindow(
46
HANDLE rgWaitHandles[2] = { };
47
UITHREAD_CONTEXT context = { };
48
49
+ // Try to make this process the first one to receive WM_QUERYENDSESSION.
50
+ // When blocking shutdown during Apply, this prevents other applications from being closed even though the restart will be blocked.
51
+ // When initiating a restart, this makes it reasonable to assume WM_QUERYENDSESSION will be received quickly because otherwise other applications could delay indefinitely.
52
+ ::SetProcessShutdownParameters(0x3FF, 0);
53
+
54
// Create event to signal after the UI thread / window is initialized.
55
rgWaitHandles[0] = ::CreateEventW(NULL, TRUE, FALSE, NULL);
56
ExitOnNullWithLastError(rgWaitHandles[0], hr, "Failed to create initialization event.");
@@ -79,10 +84,6 @@ void UiCloseMessageWindow(
84
if (::IsWindow(pEngineState->hMessageWindow))
85
{
86
::PostMessageW(pEngineState->hMessageWindow, WM_CLOSE, 0, 0);
82
-
83
- // Give the window 15 seconds to close because if it stays open it can prevent
84
- // the engine from starting a reboot (should a reboot actually be necessary).
85
- ::WaitForSingleObject(pEngineState->hMessageWindowThread, 15 * 1000);
87
}
88
}
89
@@ -180,23 +181,56 @@ static LRESULT CALLBACK WndProc(
181
182
case WM_QUERYENDSESSION:
183
{
183
- DWORD dwEndSession = static_cast<DWORD>(lParam);
184
- BOOL fCritical = ENDSESSION_CRITICAL & dwEndSession;
185
- BOOL fCancel = FALSE;
186
- BOOL fRet = FALSE;
184
+ BOOL fCritical = ENDSESSION_CRITICAL & lParam;
185
+ BOOL fAllowed = FALSE;
186
188
- // Always block shutdown during apply.
187
UITHREAD_INFO* pInfo = reinterpret_cast<UITHREAD_INFO*>(::GetWindowLongPtrW(hWnd, GWLP_USERDATA));
190
- if (pInfo->pEngineState->plan.fApplying)
188
+ if (!pInfo->pEngineState->plan.fApplying && // always block shutdown during apply.
189
+ !fCritical) // always block critical shutdowns to receive the WM_ENDSESSION message.
190
{
192
- fCancel = TRUE;
191
+ fAllowed = TRUE;
192
}
193
194
+ CoreUpdateRestartState(pInfo->pEngineState, BURN_RESTART_STATE_INITIATING);
195
pInfo->pEngineState->fCriticalShutdownInitiated |= fCritical;
196
197
- fRet = !fCancel;
198
- LogId(REPORT_STANDARD, MSG_SYSTEM_SHUTDOWN, LoggingBoolToString(fCritical), LoggingBoolToString(pInfo->fElevatedEngine), LoggingBoolToString(fRet));
199
- return fRet;
197
+ LogId(REPORT_STANDARD, MSG_SYSTEM_SHUTDOWN_REQUEST, LoggingBoolToString(fAllowed), LoggingBoolToString(pInfo->fElevatedEngine), LoggingBoolToString(fCritical), LoggingBoolToString(lParam & ENDSESSION_LOGOFF), LoggingBoolToString(lParam & ENDSESSION_CLOSEAPP));
198
+ LogFlush();
199
+ return fAllowed;
200
+ }
201
+
202
+ case WM_ENDSESSION:
203
+ {
204
+ UITHREAD_INFO* pInfo = reinterpret_cast<UITHREAD_INFO*>(::GetWindowLongPtrW(hWnd, GWLP_USERDATA));
205
+ BOOL fAllowed = 0 != wParam;
206
+
207
+ LogId(REPORT_STANDARD, MSG_SYSTEM_SHUTDOWN_RESULT, LoggingBoolToString(fAllowed), LoggingBoolToString(pInfo->fElevatedEngine), LoggingBoolToString(lParam & ENDSESSION_CRITICAL), LoggingBoolToString(lParam & ENDSESSION_LOGOFF), LoggingBoolToString(lParam & ENDSESSION_CLOSEAPP));
208
+
209
+ if (fAllowed)
210
+ {
211
+ // Windows will shutdown the process as soon as we return from this message.
212
+ // https://docs.microsoft.com/en-us/previous-versions/windows/desktop/ms700677(v=vs.85)
213
+
214
+ // Give Apply approximately 20 seconds to complete.
215
+ for (DWORD i = 0; i < 80; ++i)
216
+ {
217
+ if (!pInfo->pEngineState->plan.fApplying)
218
+ {
219
+ break;
220
+ }
221
+
222
+ ::Sleep(250);
223
+ }
224
+
225
+ LogStringWorkRaw("=======================================\r\n");
226
+
227
+ // Close the log to try to make sure everything is flushed to disk.
228
+ LogClose(FALSE);
229
+ }
230
+
231
+ CoreUpdateRestartState(pInfo->pEngineState, fAllowed ? BURN_RESTART_STATE_INITIATED : BURN_RESTART_STATE_BLOCKED);
232
+
233
+ return 0;
234
}
235
236
case WM_DESTROY: