Add BA messages around system restore points.

6 files changed, 342 insertions, 9 deletions

diff --git a/src/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h b/src/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
index cfbb0571..80b23686 100644
--- a/src/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
+++ b/src/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
@@ -141,6 +141,10 @@ enum BOOTSTRAPPER_APPLICATION_MESSAGE
     BOOTSTRAPPER_APPLICATION_MESSAGE_ONCOMMITMSITRANSACTIONCOMPLETE,
     BOOTSTRAPPER_APPLICATION_MESSAGE_ONROLLBACKMSITRANSACTIONBEGIN,
     BOOTSTRAPPER_APPLICATION_MESSAGE_ONROLLBACKMSITRANSACTIONCOMPLETE,
+    BOOTSTRAPPER_APPLICATION_MESSAGE_ONPAUSEAUTOMATICUPDATESBEGIN,
+    BOOTSTRAPPER_APPLICATION_MESSAGE_ONPAUSEAUTOMATICUPDATESCOMPLETE,
+    BOOTSTRAPPER_APPLICATION_MESSAGE_ONSYSTEMRESTOREPOINTBEGIN,
+    BOOTSTRAPPER_APPLICATION_MESSAGE_ONSYSTEMRESTOREPOINTCOMPLETE,
 };
 
 enum BOOTSTRAPPER_APPLYCOMPLETE_ACTION
@@ -818,6 +822,27 @@ struct BA_ONLAUNCHAPPROVEDEXECOMPLETE_RESULTS
     DWORD cbSize;
 };
 
+struct BA_ONPAUSEAUTOMATICUPDATESBEGIN_ARGS
+{
+    DWORD cbSize;
+};
+
+struct BA_ONPAUSEAUTOMATICUPDATESBEGIN_RESULTS
+{
+    DWORD cbSize;
+};
+
+struct BA_ONPAUSEAUTOMATICUPDATESCOMPLETE_ARGS
+{
+    DWORD cbSize;
+    HRESULT hrStatus;
+};
+
+struct BA_ONPAUSEAUTOMATICUPDATESCOMPLETE_RESULTS
+{
+    DWORD cbSize;
+};
+
 struct BA_ONPLANBEGIN_ARGS
 {
     DWORD cbSize;
@@ -1061,6 +1086,27 @@ struct BA_ONSTARTUP_RESULTS
     DWORD cbSize;
 };
 
+struct BA_ONSYSTEMRESTOREPOINTBEGIN_ARGS
+{
+    DWORD cbSize;
+};
+
+struct BA_ONSYSTEMRESTOREPOINTBEGIN_RESULTS
+{
+    DWORD cbSize;
+};
+
+struct BA_ONSYSTEMRESTOREPOINTCOMPLETE_ARGS
+{
+    DWORD cbSize;
+    HRESULT hrStatus;
+};
+
+struct BA_ONSYSTEMRESTOREPOINTCOMPLETE_RESULTS
+{
+    DWORD cbSize;
+};
+
 struct BA_ONSYSTEMSHUTDOWN_ARGS
 {
     DWORD cbSize;
diff --git a/src/engine/core.cpp b/src/engine/core.cpp
index aeae6bea..028dc1cc 100644
--- a/src/engine/core.cpp
+++ b/src/engine/core.cpp
@@ -626,7 +626,7 @@ extern "C" HRESULT CoreApply(
         hr = CoreElevate(pEngineState, pEngineState->userExperience.hwndApply);
         ExitOnFailure(hr, "Failed to elevate.");
 
-        hr = ElevationApplyInitialize(pEngineState->companionConnection.hPipe, &pEngineState->variables, pEngineState->plan.action, pEngineState->automaticUpdates, !pEngineState->fDisableSystemRestore);
+        hr = ElevationApplyInitialize(pEngineState->companionConnection.hPipe, &pEngineState->userExperience, &pEngineState->variables, pEngineState->plan.action, pEngineState->automaticUpdates, !pEngineState->fDisableSystemRestore);
         ExitOnFailure(hr, "Another per-machine setup is already executing.");
 
         fElevated = TRUE;
diff --git a/src/engine/elevation.cpp b/src/engine/elevation.cpp
index 1d0e1f1d..af5610dc 100644
--- a/src/engine/elevation.cpp
+++ b/src/engine/elevation.cpp
@@ -32,6 +32,10 @@ typedef enum _BURN_ELEVATION_MESSAGE_TYPE
     BURN_ELEVATION_MESSAGE_TYPE_COMMIT_MSI_TRANSACTION,
     BURN_ELEVATION_MESSAGE_TYPE_ROLLBACK_MSI_TRANSACTION,
 
+    BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE_PAUSE_AU_BEGIN,
+    BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE_PAUSE_AU_COMPLETE,
+    BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE_SYSTEM_RESTORE_POINT_BEGIN,
+    BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE_SYSTEM_RESTORE_POINT_COMPLETE,
     BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_PROGRESS,
     BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_ERROR,
     BURN_ELEVATION_MESSAGE_TYPE_EXECUTE_MSI_MESSAGE,
@@ -42,6 +46,13 @@ typedef enum _BURN_ELEVATION_MESSAGE_TYPE
 
 // struct
 
+typedef struct _BURN_ELEVATION_APPLY_INITIALIZE_MESSAGE_CONTEXT
+{
+    BURN_USER_EXPERIENCE* pBA;
+    BOOL fPauseCompleteNeeded;
+    BOOL fSrpCompleteNeeded;
+} BURN_ELEVATION_APPLY_INITIALIZE_MESSAGE_CONTEXT;
+
 typedef struct _BURN_ELEVATION_GENERIC_MESSAGE_CONTEXT
 {
     PFN_GENERICMESSAGEHANDLER pfnGenericMessageHandler;
@@ -89,6 +100,11 @@ static HRESULT OnLoadCompatiblePackage(
     __in BYTE* pbData,
     __in DWORD cbData
     );
+static HRESULT ProcessApplyInitializeMessages(
+    __in BURN_PIPE_MESSAGE* pMsg,
+    __in_opt LPVOID pvContext,
+    __out DWORD* pdwResult
+    );
 static HRESULT ProcessGenericExecuteMessages(
     __in BURN_PIPE_MESSAGE* pMsg,
     __in_opt LPVOID pvContext,
@@ -119,6 +135,7 @@ static HRESULT ProcessResult(
     __out BOOTSTRAPPER_APPLY_RESTART* pRestart
     );
 static HRESULT OnApplyInitialize(
+    __in HANDLE hPipe,
     __in BURN_VARIABLES* pVariables,
     __in BURN_REGISTRATION* pRegistration,
     __in HANDLE* phLock,
@@ -245,7 +262,20 @@ static HRESULT OnMsiRollbackTransaction(
     __in BYTE* pbData,
     __in DWORD cbData
     );
-
+static HRESULT ElevatedOnPauseAUBegin(
+    __in HANDLE hPipe
+    );
+static HRESULT ElevatedOnPauseAUComplete(
+    __in HANDLE hPipe,
+    __in HRESULT hrStatus
+    );
+static HRESULT ElevatedOnSystemRestorePointBegin(
+    __in HANDLE hPipe
+    );
+static HRESULT ElevatedOnSystemRestorePointComplete(
+    __in HANDLE hPipe,
+    __in HRESULT hrStatus
+    );
 
 
 // function definitions
@@ -321,6 +351,7 @@ LExit:
 
 extern "C" HRESULT ElevationApplyInitialize(
     __in HANDLE hPipe,
+    __in BURN_USER_EXPERIENCE* pBA,
     __in BURN_VARIABLES* pVariables,
     __in BOOTSTRAPPER_ACTION action,
     __in BURN_AU_PAUSE_ACTION auAction,
@@ -331,6 +362,9 @@ extern "C" HRESULT ElevationApplyInitialize(
     BYTE* pbData = NULL;
     SIZE_T cbData = 0;
     DWORD dwResult = 0;
+    BURN_ELEVATION_APPLY_INITIALIZE_MESSAGE_CONTEXT context = { };
+
+    context.pBA = pBA;
 
     // serialize message data
     hr = BuffWriteNumber(&pbData, &cbData, (DWORD)action);
@@ -346,11 +380,21 @@ extern "C" HRESULT ElevationApplyInitialize(
     ExitOnFailure(hr, "Failed to write variables.");
 
     // send message
-    hr = PipeSendMessage(hPipe, BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE, pbData, cbData, NULL, NULL, &dwResult);
+    hr = PipeSendMessage(hPipe, BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE, pbData, cbData, ProcessApplyInitializeMessages, &context, &dwResult);
     ExitOnFailure(hr, "Failed to send message to per-machine process.");
 
     hr = (HRESULT)dwResult;
 
+    // Best effort to keep the sequence of BA events sane.
+    if (context.fPauseCompleteNeeded)
+    {
+        UserExperienceOnPauseAUComplete(pBA, hr);
+    }
+    if (context.fSrpCompleteNeeded)
+    {
+        UserExperienceOnSystemRestorePointComplete(pBA, hr);
+    }
+
 LExit:
     ReleaseBuffer(pbData);
 
@@ -1317,6 +1361,68 @@ LExit:
     return hr;
 }
 
+static HRESULT ProcessApplyInitializeMessages(
+    __in BURN_PIPE_MESSAGE* pMsg,
+    __in_opt LPVOID pvContext,
+    __out DWORD* pdwResult
+    )
+{
+    HRESULT hr = S_OK;
+    BURN_ELEVATION_APPLY_INITIALIZE_MESSAGE_CONTEXT* pContext = static_cast<BURN_ELEVATION_APPLY_INITIALIZE_MESSAGE_CONTEXT*>(pvContext);
+    BYTE* pbData = (BYTE*)pMsg->pvData;
+    SIZE_T iData = 0;
+    HRESULT hrStatus = S_OK;
+    HRESULT hrBA = S_OK;
+
+    // Process the message.
+    switch (pMsg->dwMessage)
+    {
+    case BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE_PAUSE_AU_BEGIN:
+        pContext->fPauseCompleteNeeded = TRUE;
+        hrBA = UserExperienceOnPauseAUBegin(pContext->pBA);
+        break;
+
+    case BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE_PAUSE_AU_COMPLETE:
+        // read hrStatus
+        hr = BuffReadNumber(pbData, pMsg->cbData, &iData, reinterpret_cast<DWORD*>(&hrStatus));
+        ExitOnFailure(hr, "Failed to read pause AU hrStatus.");
+
+        pContext->fPauseCompleteNeeded = FALSE;
+        hrBA = UserExperienceOnPauseAUComplete(pContext->pBA, hrStatus);
+        break;
+
+    case BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE_SYSTEM_RESTORE_POINT_BEGIN:
+        if (pContext->fPauseCompleteNeeded)
+        {
+            pContext->fPauseCompleteNeeded = FALSE;
+            hrBA = UserExperienceOnPauseAUComplete(pContext->pBA, E_INVALIDSTATE);
+        }
+
+        pContext->fSrpCompleteNeeded = TRUE;
+        hrBA = UserExperienceOnSystemRestorePointBegin(pContext->pBA);
+        break;
+
+    case BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE_SYSTEM_RESTORE_POINT_COMPLETE:
+        // read hrStatus
+        hr = BuffReadNumber(pbData, pMsg->cbData, &iData, reinterpret_cast<DWORD*>(&hrStatus));
+        ExitOnFailure(hr, "Failed to read system restore point hrStatus.");
+
+        pContext->fSrpCompleteNeeded = FALSE;
+        hrBA = UserExperienceOnSystemRestorePointComplete(pContext->pBA, hrStatus);
+        break;
+
+    default:
+        hr = E_INVALIDARG;
+        ExitOnRootFailure(hr, "Invalid apply initialize message.");
+        break;
+    }
+
+    *pdwResult = static_cast<DWORD>(hrBA);
+
+LExit:
+    return hr;
+}
+
 static HRESULT ProcessGenericExecuteMessages(
     __in BURN_PIPE_MESSAGE* pMsg,
     __in_opt LPVOID pvContext,
@@ -1385,7 +1491,7 @@ static HRESULT ProcessGenericExecuteMessages(
     }
 
     // send message
-    *pdwResult =  (DWORD)pContext->pfnGenericMessageHandler(&message, pContext->pvContext);;
+    *pdwResult =  (DWORD)pContext->pfnGenericMessageHandler(&message, pContext->pvContext);
 
 LExit:
     ReleaseStr(sczMessage);
@@ -1562,7 +1668,7 @@ static HRESULT ProcessElevatedChildMessage(
         break;
 
     case BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE:
-        hrResult = OnApplyInitialize(pContext->pVariables, pContext->pRegistration, pContext->phLock, pContext->pfDisabledAutomaticUpdates, (BYTE*)pMsg->pvData, pMsg->cbData);
+        hrResult = OnApplyInitialize(pContext->hPipe, pContext->pVariables, pContext->pRegistration, pContext->phLock, pContext->pfDisabledAutomaticUpdates, (BYTE*)pMsg->pvData, pMsg->cbData);
         break;
 
     case BURN_ELEVATION_MESSAGE_TYPE_APPLY_UNINITIALIZE:
@@ -1697,6 +1803,7 @@ static HRESULT ProcessResult(
 }
 
 static HRESULT OnApplyInitialize(
+    __in HANDLE hPipe,
     __in BURN_VARIABLES* pVariables,
     __in BURN_REGISTRATION* pRegistration,
     __in HANDLE* phLock,
@@ -1711,6 +1818,7 @@ static HRESULT OnApplyInitialize(
     DWORD dwAUAction = 0;
     DWORD dwTakeSystemRestorePoint = 0;
     LPWSTR sczBundleName = NULL;
+    HRESULT hrStatus = S_OK;
 
     // Deserialize message data.
     hr = BuffReadNumber(pbData, cbData, &iData, &dwAction);
@@ -1738,9 +1846,12 @@ static HRESULT OnApplyInitialize(
     // Attempt to pause AU with best effort.
     if (BURN_AU_PAUSE_ACTION_IFELEVATED == dwAUAction || BURN_AU_PAUSE_ACTION_IFELEVATED_NORESUME == dwAUAction)
     {
+        hr = ElevatedOnPauseAUBegin(hPipe);
+        ExitOnFailure(hr, "ElevatedOnPauseAUBegin failed.");
+
         LogId(REPORT_STANDARD, MSG_PAUSE_AU_STARTING);
 
-        hr = WuaPauseAutomaticUpdates();
+        hrStatus = hr = WuaPauseAutomaticUpdates();
         if (FAILED(hr))
         {
             LogId(REPORT_STANDARD, MSG_FAILED_PAUSE_AU, hr);
@@ -1754,6 +1865,9 @@ static HRESULT OnApplyInitialize(
                 *pfDisabledWindowsUpdate = TRUE;
             }
         }
+
+        hr = ElevatedOnPauseAUComplete(hPipe, hrStatus);
+        ExitOnFailure(hr, "ElevatedOnPauseAUComplete failed.");
     }
 
     if (dwTakeSystemRestorePoint)
@@ -1765,11 +1879,14 @@ static HRESULT OnApplyInitialize(
             ExitFunction();
         }
 
+        hr = ElevatedOnSystemRestorePointBegin(hPipe);
+        ExitOnFailure(hr, "ElevatedOnSystemRestorePointBegin failed.");
+
         LogId(REPORT_STANDARD, MSG_SYSTEM_RESTORE_POINT_STARTING);
 
         BOOTSTRAPPER_ACTION action = static_cast<BOOTSTRAPPER_ACTION>(dwAction);
         SRP_ACTION restoreAction = (BOOTSTRAPPER_ACTION_INSTALL == action) ? SRP_ACTION_INSTALL : (BOOTSTRAPPER_ACTION_UNINSTALL == action) ? SRP_ACTION_UNINSTALL : SRP_ACTION_MODIFY;
-        hr = SrpCreateRestorePoint(sczBundleName, restoreAction);
+        hrStatus = hr = SrpCreateRestorePoint(sczBundleName, restoreAction);
         if (SUCCEEDED(hr))
         {
             LogId(REPORT_STANDARD, MSG_SYSTEM_RESTORE_POINT_SUCCEEDED);
@@ -1779,11 +1896,14 @@ static HRESULT OnApplyInitialize(
             LogId(REPORT_STANDARD, MSG_SYSTEM_RESTORE_POINT_DISABLED);
             hr = S_OK;
         }
-        else if (FAILED(hr))
+        else
         {
             LogId(REPORT_STANDARD, MSG_SYSTEM_RESTORE_POINT_FAILED, hr);
             hr = S_OK;
         }
+
+        hr = ElevatedOnSystemRestorePointComplete(hPipe, hrStatus);
+        ExitOnFailure(hr, "ElevatedOnSystemRestorePointComplete failed.");
     }
 
 LExit:
@@ -2699,7 +2819,7 @@ static int MsiExecuteMessageHandler(
 
     // send message
     hr = PipeSendMessage(hPipe, dwMessage, pbData, cbData, NULL, NULL, (DWORD*)&nResult);
-    ExitOnFailure(hr, "Failed to send message to per-machine process.");
+    ExitOnFailure(hr, "Failed to send msi message to per-user process.");
 
 LExit:
     ReleaseBuffer(pbData);
@@ -2867,3 +2987,75 @@ LExit:
 
     return hr;
 }
+
+static HRESULT ElevatedOnPauseAUBegin(
+    __in HANDLE hPipe
+    )
+{
+    HRESULT hr = S_OK;
+    DWORD dwResult = 0;
+
+    hr = PipeSendMessage(hPipe, BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE_PAUSE_AU_BEGIN, NULL, 0, NULL, NULL, &dwResult);
+    ExitOnFailure(hr, "Failed to send BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE_PAUSE_AU_BEGIN message to per-user process.");
+
+LExit:
+    return hr;
+}
+
+static HRESULT ElevatedOnPauseAUComplete(
+    __in HANDLE hPipe,
+    __in HRESULT hrStatus
+    )
+{
+    HRESULT hr = S_OK;
+    BYTE* pbSendData = NULL;
+    SIZE_T cbSendData = 0;
+    DWORD dwResult = 0;
+
+    hr = BuffWriteNumber(&pbSendData, &cbSendData, hrStatus);
+    ExitOnFailure(hr, "Failed to write the pause au status to message buffer.");
+
+    hr = PipeSendMessage(hPipe, BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE_PAUSE_AU_COMPLETE, pbSendData, cbSendData, NULL, NULL, &dwResult);
+    ExitOnFailure(hr, "Failed to send BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE_PAUSE_AU_COMPLETE message to per-user process.");
+
+LExit:
+    ReleaseBuffer(pbSendData);
+
+    return hr;
+}
+
+static HRESULT ElevatedOnSystemRestorePointBegin(
+    __in HANDLE hPipe
+    )
+{
+    HRESULT hr = S_OK;
+    DWORD dwResult = 0;
+
+    hr = PipeSendMessage(hPipe, BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE_SYSTEM_RESTORE_POINT_BEGIN, NULL, 0, NULL, NULL, &dwResult);
+    ExitOnFailure(hr, "Failed to send BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE_SYSTEM_RESTORE_POINT_BEGIN message to per-user process.");
+
+LExit:
+    return hr;
+}
+
+static HRESULT ElevatedOnSystemRestorePointComplete(
+    __in HANDLE hPipe,
+    __in HRESULT hrStatus
+    )
+{
+    HRESULT hr = S_OK;
+    BYTE* pbSendData = NULL;
+    SIZE_T cbSendData = 0;
+    DWORD dwResult = 0;
+
+    hr = BuffWriteNumber(&pbSendData, &cbSendData, hrStatus);
+    ExitOnFailure(hr, "Failed to write the system restore point status to message buffer.");
+
+    hr = PipeSendMessage(hPipe, BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE_SYSTEM_RESTORE_POINT_COMPLETE, pbSendData, cbSendData, NULL, NULL, &dwResult);
+    ExitOnFailure(hr, "Failed to send BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE_SYSTEM_RESTORE_POINT_COMPLETE message to per-user process.");
+
+LExit:
+    ReleaseBuffer(pbSendData);
+
+    return hr;
+}
diff --git a/src/engine/elevation.h b/src/engine/elevation.h
index 96111ae2..05fecdf6 100644
--- a/src/engine/elevation.h
+++ b/src/engine/elevation.h
@@ -14,6 +14,7 @@ HRESULT ElevationElevate(
     );
 HRESULT ElevationApplyInitialize(
     __in HANDLE hPipe,
+    __in BURN_USER_EXPERIENCE* pBA,
     __in BURN_VARIABLES* pVariables,
     __in BOOTSTRAPPER_ACTION action,
     __in BURN_AU_PAUSE_ACTION auAction,
diff --git a/src/engine/userexperience.cpp b/src/engine/userexperience.cpp
index 08451f93..b5bdc623 100644
--- a/src/engine/userexperience.cpp
+++ b/src/engine/userexperience.cpp
@@ -1482,6 +1482,46 @@ LExit:
     return hr;
 }
 
+EXTERN_C BAAPI UserExperienceOnPauseAUBegin(
+    __in BURN_USER_EXPERIENCE* pUserExperience
+    )
+{
+    HRESULT hr = S_OK;
+    BA_ONPAUSEAUTOMATICUPDATESBEGIN_ARGS args = { };
+    BA_ONPAUSEAUTOMATICUPDATESBEGIN_RESULTS results = { };
+
+    args.cbSize = sizeof(args);
+
+    results.cbSize = sizeof(results);
+    
+    hr = pUserExperience->pfnBAProc(BOOTSTRAPPER_APPLICATION_MESSAGE_ONPAUSEAUTOMATICUPDATESBEGIN, &args, &results, pUserExperience->pvBAProcContext);
+    ExitOnFailure(hr, "BA OnPauseAUBegin failed.");
+
+LExit:
+    return hr;
+}
+
+EXTERN_C BAAPI UserExperienceOnPauseAUComplete(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __in HRESULT hrStatus
+    )
+{
+    HRESULT hr = S_OK;
+    BA_ONPAUSEAUTOMATICUPDATESCOMPLETE_ARGS args = { };
+    BA_ONPAUSEAUTOMATICUPDATESCOMPLETE_RESULTS results = { };
+
+    args.cbSize = sizeof(args);
+    args.hrStatus = hrStatus;
+
+    results.cbSize = sizeof(results);
+
+    hr = pUserExperience->pfnBAProc(BOOTSTRAPPER_APPLICATION_MESSAGE_ONPAUSEAUTOMATICUPDATESCOMPLETE, &args, &results, pUserExperience->pvBAProcContext);
+    ExitOnFailure(hr, "BA OnPauseAUComplete failed.");
+
+LExit:
+    return hr;
+}
+
 EXTERN_C BAAPI UserExperienceOnPlanBegin(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __in DWORD cPackages
@@ -1983,6 +2023,46 @@ LExit:
     return hr;
 }
 
+EXTERN_C BAAPI UserExperienceOnSystemRestorePointBegin(
+    __in BURN_USER_EXPERIENCE* pUserExperience
+    )
+{
+    HRESULT hr = S_OK;
+    BA_ONSYSTEMRESTOREPOINTBEGIN_ARGS args = { };
+    BA_ONSYSTEMRESTOREPOINTBEGIN_RESULTS results = { };
+
+    args.cbSize = sizeof(args);
+
+    results.cbSize = sizeof(results);
+    
+    hr = pUserExperience->pfnBAProc(BOOTSTRAPPER_APPLICATION_MESSAGE_ONSYSTEMRESTOREPOINTBEGIN, &args, &results, pUserExperience->pvBAProcContext);
+    ExitOnFailure(hr, "BA OnSystemRestorePointBegin failed.");
+
+LExit:
+    return hr;
+}
+
+EXTERN_C BAAPI UserExperienceOnSystemRestorePointComplete(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __in HRESULT hrStatus
+    )
+{
+    HRESULT hr = S_OK;
+    BA_ONSYSTEMRESTOREPOINTCOMPLETE_ARGS args = { };
+    BA_ONSYSTEMRESTOREPOINTCOMPLETE_RESULTS results = { };
+
+    args.cbSize = sizeof(args);
+    args.hrStatus = hrStatus;
+
+    results.cbSize = sizeof(results);
+
+    hr = pUserExperience->pfnBAProc(BOOTSTRAPPER_APPLICATION_MESSAGE_ONSYSTEMRESTOREPOINTCOMPLETE, &args, &results, pUserExperience->pvBAProcContext);
+    ExitOnFailure(hr, "BA OnSystemRestorePointComplete failed.");
+
+LExit:
+    return hr;
+}
+
 EXTERN_C BAAPI UserExperienceOnSystemShutdown(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __in DWORD dwEndSession,
diff --git a/src/engine/userexperience.h b/src/engine/userexperience.h
index 387ed62b..a02d968c 100644
--- a/src/engine/userexperience.h
+++ b/src/engine/userexperience.h
@@ -343,6 +343,13 @@ BAAPI UserExperienceOnLaunchApprovedExeComplete(
     __in HRESULT hrStatus,
     __in DWORD dwProcessId
     );
+BAAPI UserExperienceOnPauseAUBegin(
+    __in BURN_USER_EXPERIENCE* pUserExperience
+    );
+BAAPI UserExperienceOnPauseAUComplete(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __in HRESULT hrStatus
+    );
 BAAPI UserExperienceOnPlanBegin(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __in DWORD cPackages
@@ -445,6 +452,13 @@ BAAPI UserExperienceOnShutdown(
 BAAPI UserExperienceOnStartup(
     __in BURN_USER_EXPERIENCE* pUserExperience
     );
+BAAPI UserExperienceOnSystemRestorePointBegin(
+    __in BURN_USER_EXPERIENCE* pUserExperience
+    );
+BAAPI UserExperienceOnSystemRestorePointComplete(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __in HRESULT hrStatus
+    );
 BAAPI UserExperienceOnSystemShutdown(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __in DWORD dwEndSession,