1 files changed, 185 insertions, 0 deletions

diff --git a/src/ext/Util/ca/RestartManager.cpp b/src/ext/Util/ca/RestartManager.cpp
new file mode 100644
index 00000000..c31819c1
--- /dev/null
+++ b/src/ext/Util/ca/RestartManager.cpp
@@ -0,0 +1,185 @@
+// 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"
+#include <restartmanager.h>
+
+// Include space for the terminating null.
+#define CCH_SESSION_KEY CCH_RM_SESSION_KEY + 1
+
+enum eRmuResourceType
+{
+    etInvalid,
+    etFilename,
+    etApplication,
+    etServiceName,
+
+    // Mask types from Attributes.
+    etTypeMask = 0xf,
+};
+
+LPCWSTR vcsRestartResourceQuery =
+    L"SELECT `Wix4RestartResource`.`Wix4RestartResource`, `Wix4RestartResource`.`Component_`, `Wix4RestartResource`.`Resource`, `Wix4RestartResource`.`Attributes` "
+    L"FROM `Wix4RestartResource`";
+enum eRestartResourceQuery { rrqRestartResource = 1, rrqComponent, rrqResource, rrqAttributes };
+
+/********************************************************************
+WixRegisterRestartResources - Immediate CA to register resources with RM.
+
+Enumerates components before InstallValidate and registers resources
+to be restarted by Restart Manager if the component action
+is anything other than None.
+
+Do not disable file system redirection.
+
+********************************************************************/
+extern "C" UINT __stdcall WixRegisterRestartResources(
+    __in MSIHANDLE hInstall
+    )
+{
+    HRESULT hr = S_OK;
+    UINT er = ERROR_SUCCESS;
+
+    PMSIHANDLE hView = NULL;
+    PMSIHANDLE hRec = NULL;
+
+    LPWSTR wzSessionKey = NULL;
+    size_t cchSessionKey = 0;
+    PRMU_SESSION pSession = NULL;
+
+    LPWSTR wzRestartResource = NULL;
+    LPWSTR wzComponent = NULL;
+    LPWSTR wzResource = NULL;
+    int iAttributes = NULL;
+    BOOL fIsComponentNull = FALSE;
+    WCA_TODO todo = WCA_TODO_UNKNOWN;
+    int iType = etInvalid;
+
+    hr = WcaInitialize(hInstall, "WixRegisterRestartResources");
+    ExitOnFailure(hr, "Failed to initialize.");
+
+    // Skip if the table doesn't exist.
+    if (S_OK != WcaTableExists(L"Wix4RestartResource"))
+    {
+        WcaLog(LOGMSG_STANDARD, "The Wix4RestartResource table does not exist; there are no resources to register with Restart Manager.");
+        ExitFunction();
+    }
+
+    // Get the existing Restart Manager session if available.
+    hr = WcaGetProperty(L"MsiRestartManagerSessionKey", &wzSessionKey);
+    ExitOnFailure(hr, "Failed to get the MsiRestartManagerSessionKey property.");
+
+    hr = ::StringCchLengthW(wzSessionKey, CCH_SESSION_KEY, &cchSessionKey);
+    ExitOnFailure(hr, "Failed to get the MsiRestartManagerSessionKey string length.");
+
+    // Skip if the property doesn't exist.
+    if (0 == cchSessionKey)
+    {
+        WcaLog(LOGMSG_STANDARD, "The MsiRestartManagerSessionKey property is not available to join.");
+        ExitFunction();
+    }
+
+    // Join the existing Restart Manager session if supported.
+    hr = RmuJoinSession(&pSession, wzSessionKey);
+    if (E_MODNOTFOUND == hr)
+    {
+        WcaLog(LOGMSG_STANDARD, "The Restart Manager is not supported on this platform. Skipping.");
+        ExitFunction1(hr = S_OK);
+    }
+    else if (FAILED(hr))
+    {
+        WcaLog(LOGMSG_STANDARD, "Failed to join the existing Restart Manager session %ls.", wzSessionKey);
+        ExitFunction1(hr = S_OK);
+    }
+
+    // Loop through each record in the table.
+    hr = WcaOpenExecuteView(vcsRestartResourceQuery, &hView);
+    ExitOnFailure(hr, "Failed to open a view on the RestartResource table.");
+
+    while (S_OK == (hr = WcaFetchRecord(hView, &hRec)))
+    {
+        hr = WcaGetRecordString(hRec, rrqRestartResource, &wzRestartResource);
+        ExitOnFailure(hr, "Failed to get the RestartResource field value.");
+
+        hr = WcaGetRecordString(hRec, rrqComponent, &wzComponent);
+        ExitOnFailure(hr, "Failed to get the Component_ field value.");
+
+        hr = WcaGetRecordFormattedString(hRec, rrqResource, &wzResource);
+        ExitOnFailure(hr, "Failed to get the Resource formatted field value.");
+
+        hr = WcaGetRecordInteger(hRec, rrqAttributes, &iAttributes);
+        ExitOnFailure(hr, "Failed to get the Attributes field value.");
+
+        fIsComponentNull = ::MsiRecordIsNull(hRec, rrqComponent);
+        todo = WcaGetComponentToDo(wzComponent);
+
+        // Only register resources for components that are null, or being installed, reinstalled, or uninstalled.
+        if (!fIsComponentNull && WCA_TODO_UNKNOWN == todo)
+        {
+            WcaLog(LOGMSG_VERBOSE, "Skipping resource %ls.", wzRestartResource);
+            continue;
+        }
+
+        // Get the type from Attributes and add to the Restart Manager.
+        iType = iAttributes & etTypeMask;
+        switch (iType)
+        {
+        case etFilename:
+            WcaLog(LOGMSG_VERBOSE, "Registering file name %ls with the Restart Manager.", wzResource);
+            hr = RmuAddFile(pSession, wzResource);
+            ExitOnFailure(hr, "Failed to register the file name with the Restart Manager session.");
+            break;
+
+        case etApplication:
+            WcaLog(LOGMSG_VERBOSE, "Registering process name %ls with the Restart Manager.", wzResource);
+            hr = RmuAddProcessesByName(pSession, wzResource);
+            if (E_NOTFOUND == hr)
+            {
+                // ERROR_ACCESS_DENIED was returned when trying to register this process.
+                // Since other instances may have been registered, log a message and continue the setup rather than failing.
+                WcaLog(LOGMSG_STANDARD, "The process, %ls, could not be registered with the Restart Manager (probably because the setup is not elevated and the process is in another user context). A reboot may be requested later.", wzResource);
+                hr = S_OK;
+            }
+            else
+            {
+                ExitOnFailure(hr, "Failed to register the process name with the Restart Manager session.");
+            }
+            break;
+
+        case etServiceName:
+            WcaLog(LOGMSG_VERBOSE, "Registering service name %ls with the Restart Manager.", wzResource);
+            hr = RmuAddService(pSession, wzResource);
+            ExitOnFailure(hr, "Failed to register the service name with the Restart Manager session.");
+            break;
+
+        default:
+            WcaLog(LOGMSG_VERBOSE, "The resource type %d for %ls is not supported and will not be registered.", iType, wzRestartResource);
+            break;
+        }
+    }
+
+    if (E_NOMOREITEMS == hr)
+    {
+        hr = S_OK;
+    }
+    ExitOnFailure(hr, "Failed while looping through all rows to register resources.");
+
+    // Register the resources and unjoin the session.
+    hr = RmuEndSession(pSession);
+    if (FAILED(hr))
+    {
+        WcaLog(LOGMSG_VERBOSE, "Failed to register the resources with the Restart Manager.");
+        ExitFunction1(hr = S_OK);
+    }
+
+LExit:
+    ReleaseStr(wzRestartResource);
+    ReleaseStr(wzComponent);
+    ReleaseStr(wzResource);
+
+    if (FAILED(hr))
+    {
+        er = ERROR_INSTALL_FAILURE;
+    }
+
+    return WcaFinalize(er);
+}