Bug 1704855 - Added automatic restart after staged --backgroundtask backgroundupdate r=nalexander,bytesized,application-update-reviewers

Differential Revision: https://phabricator.services.mozilla.com/D183041

toolkit/mozapps/update/BackgroundTask_backgroundupdate.sys.mjs

@@ -46,6 +46,11 @@ export const backgroundTaskTimeoutSec = Services.prefs.getIntPref(
   10 * 60
 );
 
+// Add 65 second minimum run time to account for restartWithSameArgs
+// having a minimum 60 second run time to properly register a restart upon
+// program exit.
+const MINIMUM_RUN_TIME_BEFORE_RESTART_MS = 65 * 1000;
+
 /**
  * Verify that pre-conditions to update this installation (both persistent and
  * transient) are fulfilled, and if they are all fulfilled, pump the update
@@ -53,7 +58,7 @@ export const backgroundTaskTimeoutSec = Services.prefs.getIntPref(
  *
  * This means checking for, downloading, and potentially applying updates.
  *
- * @returns {boolean} - `true` if an update loop was completed.
+ * @returns {any} - Returns AppUpdater status upon update loop exit.
  */
 async function _attemptBackgroundUpdate() {
   let SLUG = "_attemptBackgroundUpdate";
@@ -98,7 +103,7 @@ async function _attemptBackgroundUpdate() {
       )}'`
     );
 
-    return false;
+    return lazy.AppUpdater.STATUS.NEVER_CHECKED;
   }
 
   let result = new Promise(resolve => {
@@ -115,7 +120,7 @@ async function _attemptBackgroundUpdate() {
         );
         appUpdater.removeListener(_appUpdaterListener);
         appUpdater.stop();
-        resolve(true);
+        resolve(status);
       } else if (status == lazy.AppUpdater.STATUS.CHECKING) {
         // The usual initial flow for the Background Update Task is to kick off
         // the update download and immediately exit. For consistency, we are
@@ -152,7 +157,7 @@ async function _attemptBackgroundUpdate() {
 
           appUpdater.removeListener(_appUpdaterListener);
           appUpdater.stop();
-          resolve(true);
+          resolve(status);
         } else {
           lazy.log.debug(`${SLUG}: Download has completed!`);
         }
@@ -194,6 +199,9 @@ export async function maybeSubmitBackgroundUpdatePing() {
 export async function runBackgroundTask(commandLine) {
   let SLUG = "runBackgroundTask";
   lazy.log.error(`${SLUG}: backgroundupdate`);
+  const taskStartTime = new Date().getTime();
+  let registeredRestartFound =
+    -1 !== commandLine.findFlag("registered-restart", false);
 
   // Help debugging.  This is a pared down version of
   // `dataProviders.application` in `Troubleshoot.sys.mjs`.  When adding to this
@@ -372,8 +380,11 @@ export async function runBackgroundTask(commandLine) {
   Glean.backgroundUpdate.states.add(stringStatus);
   Glean.backgroundUpdate.finalState.set(stringStatus);
 
+  let updateStatus = lazy.AppUpdater.STATUS.NEVER_CHECKED;
   try {
-    await _attemptBackgroundUpdate();
+    // Return AppUpdater status from _attemptBackgroundUpdate() to
+    // check if the status is STATUS.READY_FOR_RESTART.
+    updateStatus = await _attemptBackgroundUpdate();
 
     lazy.log.info(`${SLUG}: attempted background update`);
     Glean.backgroundUpdate.exitCodeSuccess.set(true);
@@ -419,5 +430,42 @@ export async function runBackgroundTask(commandLine) {
   // TODO: ensure that Glean's upload mechanism is aware of Gecko shutdown.  Bug 1703572.
   await lazy.ExtensionUtils.promiseTimeout(500);
 
+  // If we're in a staged background update, we need to restart Firefox to complete the update.
+  lazy.log.debug(
+    `${SLUG}: Checking if staged background update is ready for restart`
+  );
+  // If a restart loop is occurring then registeredRestartFound will be true.
+  if (
+    updateStatus === lazy.AppUpdater.STATUS.READY_FOR_RESTART &&
+    !registeredRestartFound
+  ) {
+    lazy.log.debug(
+      `${SLUG}: Registering Firefox restart after staged background update, waiting for program to have a run time greater than 65 seconds`
+    );
+
+    // We need to restart Firefox with the same arguments to ensure
+    // the background update continues from where it was before the restart.
+    // Wait for at least 65 seconds to ensure that the application
+    // has been open for long enough to correctly register a restart.
+    const taskRunTimeMs = new Date().getTime() - taskStartTime;
+    if (taskRunTimeMs < MINIMUM_RUN_TIME_BEFORE_RESTART_MS) {
+      await lazy.ExtensionUtils.promiseTimeout(
+        MINIMUM_RUN_TIME_BEFORE_RESTART_MS - taskRunTimeMs
+      );
+    }
+
+    try {
+      Cc["@mozilla.org/updates/update-processor;1"]
+        .createInstance(Ci.nsIUpdateProcessor)
+        .registerApplicationRestartWithLaunchArgs(["-registered-restart"]);
+      lazy.log.debug(`${SLUG}: register application restart succeeded`);
+    } catch (e) {
+      lazy.log.error(
+        `${SLUG}: caught exception; failed to register application restart`,
+        e
+      );
+    }
+  }
+
   return result;
 }

toolkit/mozapps/update/nsIUpdateService.idl

@@ -655,6 +655,27 @@ interface nsIUpdateProcessor : nsISupports
    *          If this is called on a non-Windows platform.
    */
   bool getServiceRegKeyExists();
+
+  /**
+  * Registers an application restart upon program exit with the same
+  * arguments it was started with, while accepting additional arguments.
+  * The application must have been running for a minimum of 60 seconds
+  * before invoking this function due to limitations in
+  * RegisterApplicationRestart.
+  *
+  * This function should only be called on Windows.
+  *
+  * @param  argvExtra
+  *         An array of strings to be passed to the application upon
+  *         restart as additional arguments.
+  * @throws NS_ERROR_ABORT
+  *         If the application is in a restart loop.
+  * @throws NS_ERROR_NOT_IMPLEMENTED
+  *         If this is called on a non-Windows platform.
+  * @throws NS_ERROR_NOT_AVAILABLE
+  *         If the command line cannot be read.
+  */
+  void registerApplicationRestartWithLaunchArgs(in Array<AString> argvExtra);
 };
 
 /**

toolkit/xre/nsUpdateDriver.cpp

@@ -33,6 +33,7 @@
 #include "mozilla/ScopeExit.h"
 #include "mozilla/Services.h"
 #include "mozilla/dom/Promise.h"
+#include "mozilla/CmdLineAndEnvUtils.h"
 
 #ifdef XP_MACOSX
 #  include "nsILocalFileMac.h"
@@ -49,6 +50,7 @@
 #  include <windows.h>
 #  include <shlwapi.h>
 #  include <strsafe.h>
+#  include <shellapi.h>
 #  include "commonupdatedir.h"
 #  include "nsWindowsHelpers.h"
 #  include "pathhash.h"
@@ -953,3 +955,49 @@ nsUpdateProcessor::GetServiceRegKeyExists(bool* aResult) {
   return NS_ERROR_NOT_AVAILABLE;
 #endif  // #ifdef XP_WIN
 }
+
+NS_IMETHODIMP
+nsUpdateProcessor::RegisterApplicationRestartWithLaunchArgs(
+    const nsTArray<nsString>& argvExtra) {
+#ifndef XP_WIN
+  return NS_ERROR_NOT_IMPLEMENTED;
+#else
+  // Retrieve current command line arguments for restart
+  // GetCommandLineW() returns a read only pointer to
+  // the arguments the process was launched with.
+  LPWSTR currentCommandLine = GetCommandLineW();
+
+  // Register a restart flag for the application based on the current
+  // command line. The program will then automatically restart
+  // upon termination.
+  // The application must have been running for a minimum of 60
+  // seconds for a restart to be correctly registered.
+  if (currentCommandLine) {
+    // Append additional command line arguments to current command line for
+    // restart
+    nsTArray<const wchar_t*> additionalArgv(argvExtra.Length());
+    for (const nsString& arg : argvExtra) {
+      additionalArgv.AppendElement(static_cast<const wchar_t*>(arg.get()));
+    }
+
+    int currentArgc = 0;
+    LPWSTR* currentCommandLineArgv =
+        CommandLineToArgvW(currentCommandLine, &currentArgc);
+    UniquePtr<LPWSTR, LocalFreeDeleter> uniqueCurrentArgv(
+        currentCommandLineArgv);
+    mozilla::UniquePtr<wchar_t[]> restartCommandLine = mozilla::MakeCommandLine(
+        currentArgc, uniqueCurrentArgv.get(), additionalArgv.Length(),
+        additionalArgv.Elements());
+    ::RegisterApplicationRestart(restartCommandLine.get(),
+                                 RESTART_NO_CRASH | RESTART_NO_HANG);
+
+    MOZ_LOG(sUpdateLog, mozilla::LogLevel::Debug,
+            ("register application restart succeeded"));
+  } else {
+    MOZ_LOG(sUpdateLog, mozilla::LogLevel::Error,
+            ("could not register application restart"));
+    return NS_ERROR_NOT_AVAILABLE;
+  }
+  return NS_OK;
+#endif  // #ifndef XP_WIN
+}