1 files changed, 992 insertions, 0 deletions

diff --git a/src/burn/engine/engine.cpp b/src/burn/engine/engine.cpp
new file mode 100644
index 00000000..8f024e98
--- /dev/null
+++ b/src/burn/engine/engine.cpp
@@ -0,0 +1,992 @@
+// 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.
+
+#include "precomp.h"
+
+
+// constants
+
+const DWORD RESTART_RETRIES = 10;
+
+// internal function declarations
+
+static HRESULT InitializeEngineState(
+    __in BURN_ENGINE_STATE* pEngineState,
+    __in HANDLE hEngineFile
+    );
+static void UninitializeEngineState(
+    __in BURN_ENGINE_STATE* pEngineState
+    );
+static HRESULT RunUntrusted(
+    __in LPCWSTR wzCommandLine,
+    __in BURN_ENGINE_STATE* pEngineState
+    );
+static HRESULT RunNormal(
+    __in HINSTANCE hInstance,
+    __in BURN_ENGINE_STATE* pEngineState
+    );
+static HRESULT RunElevated(
+    __in HINSTANCE hInstance,
+    __in LPCWSTR wzCommandLine,
+    __in BURN_ENGINE_STATE* pEngineState
+    );
+static HRESULT RunEmbedded(
+    __in HINSTANCE hInstance,
+    __in BURN_ENGINE_STATE* pEngineState
+    );
+static HRESULT RunRunOnce(
+    __in const BURN_REGISTRATION* pRegistration,
+    __in int nCmdShow
+    );
+static HRESULT RunApplication(
+    __in BURN_ENGINE_STATE* pEngineState,
+    __out BOOL* pfReloadApp,
+    __out BOOL* pfSkipCleanup
+    );
+static HRESULT ProcessMessage(
+    __in BURN_ENGINE_STATE* pEngineState,
+    __in const MSG* pmsg
+    );
+static HRESULT DAPI RedirectLoggingOverPipe(
+    __in_z LPCSTR szString,
+    __in_opt LPVOID pvContext
+    );
+static HRESULT Restart();
+
+
+// function definitions
+
+extern "C" BOOL EngineInCleanRoom(
+    __in_z_opt LPCWSTR wzCommandLine
+    )
+{
+    // Be very careful with the functions you call from here.
+    // This function will be called before ::SetDefaultDllDirectories()
+    // has been called so dependencies outside of kernel32.dll are
+    // very likely to introduce DLL hijacking opportunities.
+
+    static DWORD cchCleanRoomSwitch = lstrlenW(BURN_COMMANDLINE_SWITCH_CLEAN_ROOM);
+
+    // This check is wholly dependent on the clean room command line switch being
+    // present at the beginning of the command line. Since Burn is the only thing
+    // that should be setting this command line option, that is in our control.
+    BOOL fInCleanRoom = (wzCommandLine &&
+        (wzCommandLine[0] == L'-' || wzCommandLine[0] == L'/') &&
+        CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, NORM_IGNORECASE, wzCommandLine + 1, cchCleanRoomSwitch, BURN_COMMANDLINE_SWITCH_CLEAN_ROOM, cchCleanRoomSwitch) &&
+        wzCommandLine[1 + cchCleanRoomSwitch] == L'='
+    );
+
+    return fInCleanRoom;
+}
+
+extern "C" HRESULT EngineRun(
+    __in HINSTANCE hInstance,
+    __in HANDLE hEngineFile,
+    __in_z_opt LPCWSTR wzCommandLine,
+    __in int nCmdShow,
+    __out DWORD* pdwExitCode
+    )
+{
+    HRESULT hr = S_OK;
+    BOOL fComInitialized = FALSE;
+    BOOL fLogInitialized = FALSE;
+    BOOL fCrypInitialized = FALSE;
+    BOOL fDpiuInitialized = FALSE;
+    BOOL fRegInitialized = FALSE;
+    BOOL fWiuInitialized = FALSE;
+    BOOL fXmlInitialized = FALSE;
+    SYSTEM_INFO si = { };
+    RTL_OSVERSIONINFOEXW ovix = { };
+    LPWSTR sczExePath = NULL;
+    BOOL fRunNormal = FALSE;
+    BOOL fRestart = FALSE;
+
+    BURN_ENGINE_STATE engineState = { };
+    engineState.command.cbSize = sizeof(BOOTSTRAPPER_COMMAND);
+
+    // Always initialize logging first
+    LogInitialize(::GetModuleHandleW(NULL));
+    fLogInitialized = TRUE;
+
+    // Ensure that log contains approriate level of information
+#ifdef _DEBUG
+    LogSetLevel(REPORT_DEBUG, FALSE);
+#else
+    LogSetLevel(REPORT_VERBOSE, FALSE); // FALSE means don't write an additional text line to the log saying the level changed
+#endif
+
+    hr = AppParseCommandLine(wzCommandLine, &engineState.argc, &engineState.argv);
+    ExitOnFailure(hr, "Failed to parse command line.");
+
+    hr = InitializeEngineState(&engineState, hEngineFile);
+    ExitOnFailure(hr, "Failed to initialize engine state.");
+
+    engineState.command.nCmdShow = nCmdShow;
+
+    // initialize platform layer
+    PlatformInitialize();
+
+    // initialize COM
+    hr = ::CoInitializeEx(NULL, COINIT_MULTITHREADED);
+    ExitOnFailure(hr, "Failed to initialize COM.");
+    fComInitialized = TRUE;
+
+    // Initialize dutil.
+    hr = CrypInitialize();
+    ExitOnFailure(hr, "Failed to initialize Cryputil.");
+    fCrypInitialized = TRUE;
+
+    DpiuInitialize();
+    fDpiuInitialized = TRUE;
+
+    hr = RegInitialize();
+    ExitOnFailure(hr, "Failed to initialize Regutil.");
+    fRegInitialized = TRUE;
+
+    hr = WiuInitialize();
+    ExitOnFailure(hr, "Failed to initialize Wiutil.");
+    fWiuInitialized = TRUE;
+
+    hr = XmlInitialize();
+    ExitOnFailure(hr, "Failed to initialize XML util.");
+    fXmlInitialized = TRUE;
+
+    hr = OsRtlGetVersion(&ovix);
+    ExitOnFailure(hr, "Failed to get OS info.");
+
+#if defined(_M_ARM64)
+    LPCSTR szBurnPlatform = "ARM64";
+#elif defined(_M_AMD64)
+    LPCSTR szBurnPlatform = "x64";
+#else
+    LPCSTR szBurnPlatform = "x86";
+#endif
+
+    LPCSTR szMachinePlatform = "unknown architecture";
+    ::GetNativeSystemInfo(&si);
+    switch (si.wProcessorArchitecture)
+    {
+    case PROCESSOR_ARCHITECTURE_AMD64:
+        szMachinePlatform = "x64";
+        break;
+    case PROCESSOR_ARCHITECTURE_ARM:
+        szMachinePlatform = "ARM";
+        break;
+    case PROCESSOR_ARCHITECTURE_ARM64:
+        szMachinePlatform = "ARM64";
+        break;
+    case PROCESSOR_ARCHITECTURE_INTEL:
+        szMachinePlatform = "x86";
+        break;
+    }
+
+    PathForCurrentProcess(&sczExePath, NULL); // Ignore failure.
+    LogId(REPORT_STANDARD, MSG_BURN_INFO, szVerMajorMinorBuild, ovix.dwMajorVersion, ovix.dwMinorVersion, ovix.dwBuildNumber, ovix.wServicePackMajor, sczExePath, szBurnPlatform, szMachinePlatform);
+    ReleaseNullStr(sczExePath);
+
+    // initialize core
+    hr = CoreInitialize(&engineState);
+    ExitOnFailure(hr, "Failed to initialize core.");
+
+    // Select run mode.
+    switch (engineState.mode)
+    {
+    case BURN_MODE_UNTRUSTED:
+        hr = RunUntrusted(wzCommandLine, &engineState);
+        ExitOnFailure(hr, "Failed to run untrusted mode.");
+        break;
+
+    case BURN_MODE_NORMAL:
+        fRunNormal = TRUE;
+
+        hr = RunNormal(hInstance, &engineState);
+        ExitOnFailure(hr, "Failed to run per-user mode.");
+        break;
+
+    case BURN_MODE_ELEVATED:
+        hr = RunElevated(hInstance, wzCommandLine, &engineState);
+        ExitOnFailure(hr, "Failed to run per-machine mode.");
+        break;
+
+    case BURN_MODE_EMBEDDED:
+        fRunNormal = TRUE;
+
+        hr = RunEmbedded(hInstance, &engineState);
+        ExitOnFailure(hr, "Failed to run embedded mode.");
+        break;
+
+    case BURN_MODE_RUNONCE:
+        hr = RunRunOnce(&engineState.registration, nCmdShow);
+        ExitOnFailure(hr, "Failed to run RunOnce mode.");
+        break;
+
+    default:
+        hr = E_UNEXPECTED;
+        ExitOnFailure(hr, "Invalid run mode.");
+    }
+
+    // set exit code and remember if we are supposed to restart.
+    *pdwExitCode = engineState.userExperience.dwExitCode;
+    fRestart = engineState.fRestart;
+
+LExit:
+    ReleaseStr(sczExePath);
+
+    // If anything went wrong but the log was never open, try to open a "failure" log
+    // and that will dump anything captured in the log memory buffer to the log.
+    if (FAILED(hr) && BURN_LOGGING_STATE_CLOSED == engineState.log.state)
+    {
+        LoggingOpenFailed();
+    }
+
+    UserExperienceRemove(&engineState.userExperience);
+
+    CacheRemoveWorkingFolder(engineState.registration.sczId);
+    CacheUninitialize();
+
+    // If this is a related bundle (but not an update) suppress restart and return the standard restart error code.
+    if (fRestart && BOOTSTRAPPER_RELATION_NONE != engineState.command.relationType && BOOTSTRAPPER_RELATION_UPDATE != engineState.command.relationType)
+    {
+        LogId(REPORT_STANDARD, MSG_RESTART_ABORTED, LoggingRelationTypeToString(engineState.command.relationType));
+
+        fRestart = FALSE;
+        hr = HRESULT_FROM_WIN32(ERROR_SUCCESS_REBOOT_REQUIRED);
+    }
+
+    UninitializeEngineState(&engineState);
+
+    if (fXmlInitialized)
+    {
+        XmlUninitialize();
+    }
+
+    if (fWiuInitialized)
+    {
+        WiuUninitialize();
+    }
+
+    if (fRegInitialized)
+    {
+        RegUninitialize();
+    }
+
+    if (fDpiuInitialized)
+    {
+        DpiuUninitialize();
+    }
+
+    if (fCrypInitialized)
+    {
+        CrypUninitialize();
+    }
+
+    if (fComInitialized)
+    {
+        ::CoUninitialize();
+    }
+
+    if (fRunNormal)
+    {
+        LogId(REPORT_STANDARD, MSG_EXITING, FAILED(hr) ? (int)hr : *pdwExitCode, LoggingBoolToString(fRestart));
+
+        if (fRestart)
+        {
+            LogId(REPORT_STANDARD, MSG_RESTARTING);
+        }
+    }
+
+    if (fLogInitialized)
+    {
+        LogClose(FALSE);
+    }
+
+    if (fRestart)
+    {
+        Restart();
+    }
+
+    if (fLogInitialized)
+    {
+        LogUninitialize(FALSE);
+    }
+
+    return hr;
+}
+
+
+// internal function definitions
+
+static HRESULT InitializeEngineState(
+    __in BURN_ENGINE_STATE* pEngineState,
+    __in HANDLE hEngineFile
+    )
+{
+    HRESULT hr = S_OK;
+    LPCWSTR wzParam = NULL;
+    HANDLE hSectionFile = hEngineFile;
+    HANDLE hSourceEngineFile = INVALID_HANDLE_VALUE;
+    DWORD64 qw = 0;
+
+    pEngineState->automaticUpdates = BURN_AU_PAUSE_ACTION_IFELEVATED;
+    pEngineState->dwElevatedLoggingTlsId = TLS_OUT_OF_INDEXES;
+    ::InitializeCriticalSection(&pEngineState->userExperience.csEngineActive);
+    PipeConnectionInitialize(&pEngineState->companionConnection);
+    PipeConnectionInitialize(&pEngineState->embeddedConnection);
+
+    for (int i = 0; i < pEngineState->argc; ++i)
+    {
+        if (pEngineState->argv[i][0] == L'-')
+        {
+            if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, NORM_IGNORECASE, &pEngineState->argv[i][1], lstrlenW(BURN_COMMANDLINE_SWITCH_FILEHANDLE_ATTACHED), BURN_COMMANDLINE_SWITCH_FILEHANDLE_ATTACHED, lstrlenW(BURN_COMMANDLINE_SWITCH_FILEHANDLE_ATTACHED)))
+            {
+                wzParam = &pEngineState->argv[i][2 + lstrlenW(BURN_COMMANDLINE_SWITCH_FILEHANDLE_ATTACHED)];
+                if (L'=' != wzParam[-1] || L'\0' == wzParam[0])
+                {
+                    ExitOnRootFailure(hr = E_INVALIDARG, "Missing required parameter for switch: %ls", BURN_COMMANDLINE_SWITCH_FILEHANDLE_ATTACHED);
+                }
+
+                hr = StrStringToUInt64(wzParam, 0, &qw);
+                ExitOnFailure(hr, "Failed to parse file handle: '%ls'", (wzParam));
+
+                hSourceEngineFile = (HANDLE)qw;
+            }
+            if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, NORM_IGNORECASE, &pEngineState->argv[i][1], lstrlenW(BURN_COMMANDLINE_SWITCH_FILEHANDLE_SELF), BURN_COMMANDLINE_SWITCH_FILEHANDLE_SELF, lstrlenW(BURN_COMMANDLINE_SWITCH_FILEHANDLE_SELF)))
+            {
+                wzParam = &pEngineState->argv[i][2 + lstrlenW(BURN_COMMANDLINE_SWITCH_FILEHANDLE_SELF)];
+                if (L'=' != wzParam[-1] || L'\0' == wzParam[0])
+                {
+                    ExitOnRootFailure(hr = E_INVALIDARG, "Missing required parameter for switch: %ls", BURN_COMMANDLINE_SWITCH_FILEHANDLE_SELF);
+                }
+
+                hr = StrStringToUInt64(wzParam, 0, &qw);
+                ExitOnFailure(hr, "Failed to parse file handle: '%ls'", (wzParam));
+
+                hSectionFile = (HANDLE)qw;
+            }
+        }
+    }
+
+    hr = SectionInitialize(&pEngineState->section, hSectionFile, hSourceEngineFile);
+    ExitOnFailure(hr, "Failed to initialize engine section.");
+
+LExit:
+    return hr;
+}
+
+static void UninitializeEngineState(
+    __in BURN_ENGINE_STATE* pEngineState
+    )
+{
+    if (pEngineState->argv)
+    {
+        AppFreeCommandLineArgs(pEngineState->argv);
+    }
+
+    ReleaseStr(pEngineState->sczIgnoreDependencies);
+
+    PipeConnectionUninitialize(&pEngineState->embeddedConnection);
+    PipeConnectionUninitialize(&pEngineState->companionConnection);
+    ReleaseStr(pEngineState->sczBundleEngineWorkingPath)
+
+    ReleaseHandle(pEngineState->hMessageWindowThread);
+
+    BurnExtensionUninitialize(&pEngineState->extensions);
+
+    ::DeleteCriticalSection(&pEngineState->userExperience.csEngineActive);
+    UserExperienceUninitialize(&pEngineState->userExperience);
+
+    ApprovedExesUninitialize(&pEngineState->approvedExes);
+    UpdateUninitialize(&pEngineState->update);
+    VariablesUninitialize(&pEngineState->variables);
+    SearchesUninitialize(&pEngineState->searches);
+    RegistrationUninitialize(&pEngineState->registration);
+    PayloadsUninitialize(&pEngineState->payloads);
+    PackagesUninitialize(&pEngineState->packages);
+    SectionUninitialize(&pEngineState->section);
+    ContainersUninitialize(&pEngineState->containers);
+
+    ReleaseStr(pEngineState->command.wzBootstrapperApplicationDataPath);
+    ReleaseStr(pEngineState->command.wzBootstrapperWorkingFolder);
+    ReleaseStr(pEngineState->command.wzLayoutDirectory);
+    ReleaseStr(pEngineState->command.wzCommandLine);
+
+    ReleaseStr(pEngineState->log.sczExtension);
+    ReleaseStr(pEngineState->log.sczPrefix);
+    ReleaseStr(pEngineState->log.sczPath);
+    ReleaseStr(pEngineState->log.sczPathVariable);
+
+    if (TLS_OUT_OF_INDEXES != pEngineState->dwElevatedLoggingTlsId)
+    {
+        ::TlsFree(pEngineState->dwElevatedLoggingTlsId);
+    }
+
+    // clear struct
+    memset(pEngineState, 0, sizeof(BURN_ENGINE_STATE));
+}
+
+static HRESULT RunUntrusted(
+    __in LPCWSTR wzCommandLine,
+    __in BURN_ENGINE_STATE* pEngineState
+    )
+{
+    HRESULT hr = S_OK;
+    LPWSTR sczCurrentProcessPath = NULL;
+    LPWSTR wzCleanRoomBundlePath = NULL;
+    LPWSTR sczCachedCleanRoomBundlePath = NULL;
+    LPWSTR sczParameters = NULL;
+    LPWSTR sczFullCommandLine = NULL;
+    STARTUPINFOW si = { };
+    PROCESS_INFORMATION pi = { };
+    HANDLE hFileAttached = NULL;
+    HANDLE hFileSelf = NULL;
+    HANDLE hProcess = NULL;
+
+    hr = PathForCurrentProcess(&sczCurrentProcessPath, NULL);
+    ExitOnFailure(hr, "Failed to get path for current process.");
+
+    BOOL fRunningFromCache = CacheBundleRunningFromCache();
+
+    // If we're running from the package cache, we're in a secure
+    // folder (DLLs cannot be inserted here for hijacking purposes)
+    // so just launch the current process's path as the clean room
+    // process. Technically speaking, we'd be able to skip creating
+    // a clean room process at all (since we're already running from
+    // a secure folder) but it makes the code that only wants to run
+    // in clean room more complicated if we don't launch an explicit
+    // clean room process.
+    if (fRunningFromCache)
+    {
+        wzCleanRoomBundlePath = sczCurrentProcessPath;
+    }
+    else
+    {
+        hr = CacheBundleToCleanRoom(&pEngineState->section, &sczCachedCleanRoomBundlePath);
+        ExitOnFailure(hr, "Failed to cache to clean room.");
+
+        wzCleanRoomBundlePath = sczCachedCleanRoomBundlePath;
+    }
+
+    // The clean room switch must always be at the front of the command line so
+    // the EngineInCleanRoom function will operate correctly.
+    hr = StrAllocFormatted(&sczParameters, L"-%ls=\"%ls\"", BURN_COMMANDLINE_SWITCH_CLEAN_ROOM, sczCurrentProcessPath);
+    ExitOnFailure(hr, "Failed to allocate parameters for unelevated process.");
+
+    // Send a file handle for the child Burn process to access the attached container.
+    hr = CoreAppendFileHandleAttachedToCommandLine(pEngineState->section.hEngineFile, &hFileAttached, &sczParameters);
+    ExitOnFailure(hr, "Failed to append %ls", BURN_COMMANDLINE_SWITCH_FILEHANDLE_ATTACHED);
+
+    // Grab a file handle for the child Burn process.
+    hr = CoreAppendFileHandleSelfToCommandLine(wzCleanRoomBundlePath, &hFileSelf, &sczParameters, NULL);
+    ExitOnFailure(hr, "Failed to append %ls", BURN_COMMANDLINE_SWITCH_FILEHANDLE_SELF);
+
+    hr = StrAllocFormattedSecure(&sczParameters, L"%ls %ls", sczParameters, wzCommandLine);
+    ExitOnFailure(hr, "Failed to append original command line.");
+
+#ifdef ENABLE_UNELEVATE
+    // TODO: Pass file handle to unelevated process if this ever gets reenabled.
+    if (!pEngineState->fDisableUnelevate)
+    {
+        // Try to launch unelevated and if that fails for any reason, we'll launch our process normally (even though that may make it elevated).
+        hr = ProcExecuteAsInteractiveUser(wzCleanRoomBundlePath, sczParameters, &hProcess);
+    }
+#endif
+
+    if (!hProcess)
+    {
+        hr = StrAllocFormattedSecure(&sczFullCommandLine, L"\"%ls\" %ls", wzCleanRoomBundlePath, sczParameters);
+        ExitOnFailure(hr, "Failed to allocate full command-line.");
+
+        si.cb = sizeof(si);
+        si.wShowWindow = static_cast<WORD>(pEngineState->command.nCmdShow);
+        if (!::CreateProcessW(wzCleanRoomBundlePath, sczFullCommandLine, NULL, NULL, TRUE, 0, 0, NULL, &si, &pi))
+        {
+            ExitWithLastError(hr, "Failed to launch clean room process: %ls", sczFullCommandLine);
+        }
+
+        hProcess = pi.hProcess;
+        pi.hProcess = NULL;
+    }
+
+    hr = ProcWaitForCompletion(hProcess, INFINITE, &pEngineState->userExperience.dwExitCode);
+    ExitOnFailure(hr, "Failed to wait for clean room process: %ls", wzCleanRoomBundlePath);
+
+LExit:
+    ReleaseHandle(pi.hThread);
+    ReleaseFileHandle(hFileSelf);
+    ReleaseFileHandle(hFileAttached);
+    ReleaseHandle(hProcess);
+    StrSecureZeroFreeString(sczFullCommandLine);
+    StrSecureZeroFreeString(sczParameters);
+    ReleaseStr(sczCachedCleanRoomBundlePath);
+    ReleaseStr(sczCurrentProcessPath);
+
+    return hr;
+}
+
+static HRESULT RunNormal(
+    __in HINSTANCE hInstance,
+    __in BURN_ENGINE_STATE* pEngineState
+    )
+{
+    HRESULT hr = S_OK;
+    LPWSTR sczOriginalSource = NULL;
+    LPWSTR sczCopiedOriginalSource = NULL;
+    BOOL fContinueExecution = TRUE;
+    BOOL fReloadApp = FALSE;
+    BOOL fSkipCleanup = FALSE;
+    BURN_EXTENSION_ENGINE_CONTEXT extensionEngineContext = { };
+
+    // Initialize logging.
+    hr = LoggingOpen(&pEngineState->log, &pEngineState->variables, pEngineState->command.display, pEngineState->registration.sczDisplayName);
+    ExitOnFailure(hr, "Failed to open log.");
+
+    // Ensure we're on a supported operating system.
+    hr = ConditionGlobalCheck(&pEngineState->variables, &pEngineState->condition, pEngineState->command.display, pEngineState->registration.sczDisplayName, &pEngineState->userExperience.dwExitCode, &fContinueExecution);
+    ExitOnFailure(hr, "Failed to check global conditions");
+
+    if (!fContinueExecution)
+    {
+        LogId(REPORT_STANDARD, MSG_FAILED_CONDITION_CHECK);
+
+        // If the block told us to abort, abort!
+        ExitFunction1(hr = S_OK);
+    }
+
+    if (pEngineState->userExperience.fSplashScreen && BOOTSTRAPPER_DISPLAY_NONE < pEngineState->command.display)
+    {
+        SplashScreenCreate(hInstance, NULL, &pEngineState->command.hwndSplashScreen);
+    }
+
+    // Create a top-level window to handle system messages.
+    hr = UiCreateMessageWindow(hInstance, pEngineState);
+    ExitOnFailure(hr, "Failed to create the message window.");
+
+    // Query registration state.
+    hr = CoreQueryRegistration(pEngineState);
+    ExitOnFailure(hr, "Failed to query registration.");
+
+    // Best effort to set the source of attached containers to BURN_BUNDLE_ORIGINAL_SOURCE.
+    hr = VariableGetString(&pEngineState->variables, BURN_BUNDLE_ORIGINAL_SOURCE, &sczOriginalSource);
+    if (SUCCEEDED(hr))
+    {
+        for (DWORD i = 0; i < pEngineState->containers.cContainers; ++i)
+        {
+            BURN_CONTAINER* pContainer = pEngineState->containers.rgContainers + i;
+            if (pContainer->fAttached)
+            {
+                hr = StrAllocString(&sczCopiedOriginalSource, sczOriginalSource, 0);
+                if (SUCCEEDED(hr))
+                {
+                    ReleaseNullStr(pContainer->sczSourcePath);
+                    pContainer->sczSourcePath = sczCopiedOriginalSource;
+                    sczCopiedOriginalSource = NULL;
+                }
+            }
+        }
+    }
+    hr = S_OK;
+
+    // Set some built-in variables before loading the BA.
+    hr = PlanSetVariables(pEngineState->command.action, &pEngineState->variables);
+    ExitOnFailure(hr, "Failed to set action variables.");
+
+    hr = RegistrationSetVariables(&pEngineState->registration, &pEngineState->variables);
+    ExitOnFailure(hr, "Failed to set registration variables.");
+
+    // If a layout directory was specified on the command-line, set it as a well-known variable.
+    if (pEngineState->command.wzLayoutDirectory && *pEngineState->command.wzLayoutDirectory)
+    {
+        hr = VariableSetString(&pEngineState->variables, BURN_BUNDLE_LAYOUT_DIRECTORY, pEngineState->command.wzLayoutDirectory, FALSE, FALSE);
+        ExitOnFailure(hr, "Failed to set layout directory variable to value provided from command-line.");
+    }
+
+    // Setup the extension engine.
+    extensionEngineContext.pEngineState = pEngineState;
+
+    // Load the extensions.
+    hr = BurnExtensionLoad(&pEngineState->extensions, &extensionEngineContext);
+    ExitOnFailure(hr, "Failed to load BundleExtensions.");
+
+    do
+    {
+        fReloadApp = FALSE;
+        pEngineState->fQuit = FALSE;
+
+        hr = RunApplication(pEngineState, &fReloadApp, &fSkipCleanup);
+        ExitOnFailure(hr, "Failed while running ");
+    } while (fReloadApp);
+
+LExit:
+    if (!fSkipCleanup)
+    {
+        CoreCleanup(pEngineState);
+    }
+
+    BurnExtensionUnload(&pEngineState->extensions);
+
+    // If the message window is still around, close it.
+    UiCloseMessageWindow(pEngineState);
+
+    VariablesDump(&pEngineState->variables);
+
+    // end per-machine process if running
+    if (INVALID_HANDLE_VALUE != pEngineState->companionConnection.hPipe)
+    {
+        PipeTerminateChildProcess(&pEngineState->companionConnection, pEngineState->userExperience.dwExitCode, FALSE);
+    }
+
+    // If the splash screen is still around, close it.
+    if (::IsWindow(pEngineState->command.hwndSplashScreen))
+    {
+        ::PostMessageW(pEngineState->command.hwndSplashScreen, WM_CLOSE, 0, 0);
+    }
+
+    ReleaseStr(sczOriginalSource);
+    ReleaseStr(sczCopiedOriginalSource);
+
+    return hr;
+}
+
+static HRESULT RunElevated(
+    __in HINSTANCE hInstance,
+    __in LPCWSTR /*wzCommandLine*/,
+    __in BURN_ENGINE_STATE* pEngineState
+    )
+{
+    HRESULT hr = S_OK;
+    HANDLE hLock = NULL;
+    BOOL fDisabledAutomaticUpdates = FALSE;
+
+    // connect to per-user process
+    hr = PipeChildConnect(&pEngineState->companionConnection, TRUE);
+    ExitOnFailure(hr, "Failed to connect to unelevated process.");
+
+    // Set up the thread local storage to store the correct pipe to communicate logging then
+    // override logging to write over the pipe.
+    pEngineState->dwElevatedLoggingTlsId = ::TlsAlloc();
+    if (TLS_OUT_OF_INDEXES == pEngineState->dwElevatedLoggingTlsId)
+    {
+        ExitWithLastError(hr, "Failed to allocate thread local storage for logging.");
+    }
+
+    if (!::TlsSetValue(pEngineState->dwElevatedLoggingTlsId, pEngineState->companionConnection.hPipe))
+    {
+        ExitWithLastError(hr, "Failed to set elevated pipe into thread local storage for logging.");
+    }
+
+    LogRedirect(RedirectLoggingOverPipe, pEngineState);
+
+    // Create a top-level window to prevent shutting down the elevated process.
+    hr = UiCreateMessageWindow(hInstance, pEngineState);
+    ExitOnFailure(hr, "Failed to create the message window.");
+
+    SrpInitialize(TRUE);
+
+    // Pump messages from parent process.
+    hr = ElevationChildPumpMessages(pEngineState->dwElevatedLoggingTlsId, pEngineState->companionConnection.hPipe, pEngineState->companionConnection.hCachePipe, &pEngineState->approvedExes, &pEngineState->containers, &pEngineState->packages, &pEngineState->payloads, &pEngineState->variables, &pEngineState->registration, &pEngineState->userExperience, &hLock, &fDisabledAutomaticUpdates, &pEngineState->userExperience.dwExitCode, &pEngineState->fRestart);
+    LogRedirect(NULL, NULL); // reset logging so the next failure gets written to "log buffer" for the failure log.
+    ExitOnFailure(hr, "Failed to pump messages from parent process.");
+
+LExit:
+    LogRedirect(NULL, NULL); // we're done talking to the child so always reset logging now.
+
+    // If the message window is still around, close it.
+    UiCloseMessageWindow(pEngineState);
+
+    if (fDisabledAutomaticUpdates)
+    {
+        ElevationChildResumeAutomaticUpdates();
+    }
+
+    if (hLock)
+    {
+        ::ReleaseMutex(hLock);
+        ::CloseHandle(hLock);
+    }
+
+    return hr;
+}
+
+static HRESULT RunEmbedded(
+    __in HINSTANCE hInstance,
+    __in BURN_ENGINE_STATE* pEngineState
+    )
+{
+    HRESULT hr = S_OK;
+
+    // Disable system restore since the parent bundle may have done it.
+    pEngineState->fDisableSystemRestore = TRUE;
+
+    // Connect to parent process.
+    hr = PipeChildConnect(&pEngineState->embeddedConnection, FALSE);
+    ExitOnFailure(hr, "Failed to connect to parent of embedded process.");
+
+    // Do not register the bundle to automatically restart if embedded.
+    if (BOOTSTRAPPER_DISPLAY_EMBEDDED == pEngineState->command.display)
+    {
+        pEngineState->registration.fDisableResume = TRUE;
+    }
+
+    // Now run the application like normal.
+    hr = RunNormal(hInstance, pEngineState);
+    ExitOnFailure(hr, "Failed to run bootstrapper application embedded.");
+
+LExit:
+    return hr;
+}
+
+static HRESULT RunRunOnce(
+    __in const BURN_REGISTRATION* pRegistration,
+    __in int nCmdShow
+    )
+{
+    HRESULT hr = S_OK;
+    LPWSTR sczNewCommandLine = NULL;
+    LPWSTR sczBurnPath = NULL;
+    HANDLE hProcess = NULL;
+
+    hr = RegistrationGetResumeCommandLine(pRegistration, &sczNewCommandLine);
+    ExitOnFailure(hr, "Unable to get resume command line from the registry");
+
+    // and re-launch
+    hr = PathForCurrentProcess(&sczBurnPath, NULL);
+    ExitOnFailure(hr, "Failed to get current process path.");
+
+    hr = ProcExec(sczBurnPath, 0 < sczNewCommandLine ? sczNewCommandLine : L"", nCmdShow, &hProcess);
+    ExitOnFailure(hr, "Failed to re-launch bundle process after RunOnce: %ls", sczBurnPath);
+
+LExit:
+    ReleaseHandle(hProcess);
+    ReleaseStr(sczNewCommandLine);
+    ReleaseStr(sczBurnPath);
+
+    return hr;
+}
+
+static HRESULT RunApplication(
+    __in BURN_ENGINE_STATE* pEngineState,
+    __out BOOL* pfReloadApp,
+    __out BOOL* pfSkipCleanup
+    )
+{
+    HRESULT hr = S_OK;
+    BOOTSTRAPPER_ENGINE_CONTEXT engineContext = { };
+    BOOL fStartupCalled = FALSE;
+    BOOL fRet = FALSE;
+    MSG msg = { };
+    BOOTSTRAPPER_SHUTDOWN_ACTION shutdownAction = BOOTSTRAPPER_SHUTDOWN_ACTION_NONE;
+
+    ::PeekMessageW(&msg, NULL, WM_USER, WM_USER, PM_NOREMOVE);
+
+    // Setup the bootstrapper engine.
+    engineContext.dwThreadId = ::GetCurrentThreadId();
+    engineContext.pEngineState = pEngineState;
+
+    // Load the bootstrapper application.
+    hr = UserExperienceLoad(&pEngineState->userExperience, &engineContext, &pEngineState->command);
+    ExitOnFailure(hr, "Failed to load BA.");
+
+    fStartupCalled = TRUE;
+    hr = UserExperienceOnStartup(&pEngineState->userExperience);
+    ExitOnFailure(hr, "Failed to start bootstrapper application.");
+
+    // Enter the message pump.
+    while (0 != (fRet = ::GetMessageW(&msg, NULL, 0, 0)))
+    {
+        if (-1 == fRet)
+        {
+            hr = E_UNEXPECTED;
+            ExitOnRootFailure(hr, "Unexpected return value from message pump.");
+        }
+        else
+        {
+            // When the BA makes a request from its own thread, it's common for the PostThreadMessage in externalengine.cpp
+            // to block until this thread waits on something. It's also common for Detect and Plan to never wait on something.
+            // In the extreme case, the engine could be elevating in Apply before the Detect call returned to the BA.
+            // This helps to avoid that situation, which could be blocking a UI thread.
+            ::Sleep(0);
+
+            ProcessMessage(pEngineState, &msg);
+        }
+    }
+
+    // Get exit code.
+    pEngineState->userExperience.dwExitCode = (DWORD)msg.wParam;
+
+LExit:
+    if (fStartupCalled)
+    {
+        UserExperienceOnShutdown(&pEngineState->userExperience, &shutdownAction);
+        if (BOOTSTRAPPER_SHUTDOWN_ACTION_RESTART == shutdownAction)
+        {
+            LogId(REPORT_STANDARD, MSG_BA_REQUESTED_RESTART, LoggingBoolToString(pEngineState->fRestart));
+            pEngineState->fRestart = TRUE;
+        }
+        else if (BOOTSTRAPPER_SHUTDOWN_ACTION_RELOAD_BOOTSTRAPPER == shutdownAction)
+        {
+            LogId(REPORT_STANDARD, MSG_BA_REQUESTED_RELOAD);
+            *pfReloadApp = TRUE;
+        }
+        else if (BOOTSTRAPPER_SHUTDOWN_ACTION_SKIP_CLEANUP == shutdownAction)
+        {
+            LogId(REPORT_STANDARD, MSG_BA_REQUESTED_SKIP_CLEANUP);
+            *pfSkipCleanup = TRUE;
+        }
+    }
+
+    // Unload BA.
+    UserExperienceUnload(&pEngineState->userExperience);
+
+    return hr;
+}
+
+static HRESULT ProcessMessage(
+    __in BURN_ENGINE_STATE* pEngineState,
+    __in const MSG* pmsg
+    )
+{
+    HRESULT hr = S_OK;
+
+    UserExperienceActivateEngine(&pEngineState->userExperience);
+
+    if (pEngineState->fQuit)
+    {
+        LogId(REPORT_WARNING, MSG_IGNORE_OPERATION_AFTER_QUIT, LoggingBurnMessageToString(pmsg->message));
+        ExitFunction1(hr = E_INVALIDSTATE);
+    }
+
+    switch (pmsg->message)
+    {
+    case WM_BURN_DETECT:
+        hr = CoreDetect(pEngineState, reinterpret_cast<HWND>(pmsg->lParam));
+        break;
+
+    case WM_BURN_PLAN:
+        hr = CorePlan(pEngineState, static_cast<BOOTSTRAPPER_ACTION>(pmsg->lParam));
+        break;
+
+    case WM_BURN_ELEVATE:
+        hr = CoreElevate(pEngineState, reinterpret_cast<HWND>(pmsg->lParam));
+        break;
+
+    case WM_BURN_APPLY:
+        hr = CoreApply(pEngineState, reinterpret_cast<HWND>(pmsg->lParam));
+        break;
+
+    case WM_BURN_LAUNCH_APPROVED_EXE:
+        hr = CoreLaunchApprovedExe(pEngineState, reinterpret_cast<BURN_LAUNCH_APPROVED_EXE*>(pmsg->lParam));
+        break;
+
+    case WM_BURN_QUIT:
+        hr = CoreQuit(pEngineState, static_cast<int>(pmsg->wParam));
+        break;
+    }
+
+LExit:
+    UserExperienceDeactivateEngine(&pEngineState->userExperience);
+
+    return hr;
+}
+
+static HRESULT DAPI RedirectLoggingOverPipe(
+    __in_z LPCSTR szString,
+    __in_opt LPVOID pvContext
+    )
+{
+    static BOOL s_fCurrentlyLoggingToPipe = FALSE;
+
+    HRESULT hr = S_OK;
+    BURN_ENGINE_STATE* pEngineState = static_cast<BURN_ENGINE_STATE*>(pvContext);
+    BOOL fStartedLogging = FALSE;
+    HANDLE hPipe = INVALID_HANDLE_VALUE;
+    BYTE* pbData = NULL;
+    SIZE_T cbData = 0;
+    DWORD dwResult = 0;
+
+    // Prevent this function from being called recursively.
+    if (s_fCurrentlyLoggingToPipe)
+    {
+        ExitFunction();
+    }
+
+    s_fCurrentlyLoggingToPipe = TRUE;
+    fStartedLogging = TRUE;
+
+    // Make sure the current thread set the pipe in TLS.
+    hPipe = ::TlsGetValue(pEngineState->dwElevatedLoggingTlsId);
+    if (!hPipe || INVALID_HANDLE_VALUE == hPipe)
+    {
+        hr = HRESULT_FROM_WIN32(ERROR_PIPE_NOT_CONNECTED);
+        ExitFunction();
+    }
+
+    // Do not log or use ExitOnFailure() macro here because they will be discarded
+    // by the recursive block at the top of this function.
+    hr = BuffWriteStringAnsi(&pbData, &cbData, szString);
+    if (SUCCEEDED(hr))
+    {
+        hr = PipeSendMessage(hPipe, static_cast<DWORD>(BURN_PIPE_MESSAGE_TYPE_LOG), pbData, cbData, NULL, NULL, &dwResult);
+        if (SUCCEEDED(hr))
+        {
+            hr = (HRESULT)dwResult;
+        }
+    }
+
+LExit:
+    ReleaseBuffer(pbData);
+
+    // We started logging so remember to say we are no longer logging.
+    if (fStartedLogging)
+    {
+        s_fCurrentlyLoggingToPipe = FALSE;
+    }
+
+    return hr;
+}
+
+static HRESULT Restart()
+{
+    HRESULT hr = S_OK;
+    HANDLE hProcessToken = NULL;
+    TOKEN_PRIVILEGES priv = { };
+    DWORD dwRetries = 0;
+
+    if (!::OpenProcessToken(::GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hProcessToken))
+    {
+        ExitWithLastError(hr, "Failed to get process token.");
+    }
+
+    priv.PrivilegeCount = 1;
+    priv.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
+    if (!::LookupPrivilegeValueW(NULL, L"SeShutdownPrivilege", &priv.Privileges[0].Luid))
+    {
+        ExitWithLastError(hr, "Failed to get shutdown privilege LUID.");
+    }
+
+    if (!::AdjustTokenPrivileges(hProcessToken, FALSE, &priv, sizeof(TOKEN_PRIVILEGES), NULL, 0))
+    {
+        ExitWithLastError(hr, "Failed to adjust token to add shutdown privileges.");
+    }
+
+    do
+    {
+        hr = S_OK;
+
+        // Wait a second to let the companion process (assuming we did an elevated install) to get to the
+        // point where it too is thinking about restarting the computer. Only one will schedule the restart
+        // but both will have their log files closed and otherwise be ready to exit.
+        //
+        // On retry, we'll also wait a second to let the OS try to get to a place where the restart can
+        // be initiated.
+        ::Sleep(1000);
+
+        if (!vpfnInitiateSystemShutdownExW(NULL, NULL, 0, FALSE, TRUE, SHTDN_REASON_MAJOR_APPLICATION | SHTDN_REASON_MINOR_INSTALLATION | SHTDN_REASON_FLAG_PLANNED))
+        {
+            hr = HRESULT_FROM_WIN32(::GetLastError());
+        }
+    } while (dwRetries++ < RESTART_RETRIES && (HRESULT_FROM_WIN32(ERROR_MACHINE_LOCKED) == hr || HRESULT_FROM_WIN32(ERROR_NOT_READY) == hr));
+    ExitOnRootFailure(hr, "Failed to schedule restart.");
+
+LExit:
+    ReleaseHandle(hProcessToken);
+    return hr;
+}