Abort restarts if updated version is rolled back

updates/main.go

@@ -65,6 +65,8 @@ const (
 
 	updateFailed  = "updates:failed"
 	updateSuccess = "updates:success"
+
+	updateTaskRepeatDuration = 1 * time.Hour
 )
 
 func init() {
@@ -157,7 +159,7 @@ func start() error {
 
 	if !disableTaskSchedule {
 		updateTask.
-			Repeat(1 * time.Hour).
+			Repeat(updateTaskRepeatDuration).
 			MaxDelay(30 * time.Minute)
 	}
 

updates/restart.go

@@ -2,6 +2,7 @@ package updates
 
 import (
 	"context"
+	"sync"
 	"time"
 
 	"github.com/tevino/abool"
@@ -19,11 +20,26 @@ var (
 	restartTask      *modules.Task
 	restartPending   = abool.New()
 	restartTriggered = abool.New()
+
+	restartTime     time.Time
+	restartTimeLock sync.Mutex
 )
 
-// IsRestarting returns whether a restart is pending or currently in progress.
+// IsRestarting returns whether a restart has been triggered.
 func IsRestarting() bool {
-	return restartPending.IsSet()
+	return restartTriggered.IsSet()
+}
+
+// RestartIsPending returns whether a restart is pending.
+func RestartIsPending() (pending bool, restartAt time.Time) {
+	if restartPending.IsNotSet() {
+		return false, time.Time{}
+	}
+
+	restartTimeLock.Lock()
+	defer restartTimeLock.Unlock()
+
+	return true, restartTime
 }
 
 // DelayedRestart triggers a restart of the application by shutting down the
@@ -31,36 +47,57 @@ func IsRestarting() bool {
 // may be further delayed by up to 10 minutes by the internal task scheduling
 // system. This only works if the process is managed by portmaster-start.
 func DelayedRestart(delay time.Duration) {
-	log.Warningf("updates: restart triggered, will execute in %s", delay)
+	// Check if restart is already pending.
-
+	if !restartPending.SetToIf(false, true) {
-	// This enables TriggerRestartIfPending.
+		return
-	// Subsequent calls to TriggerRestart should be able to set a new delay.
+	}
-	restartPending.Set()
 
 	// Schedule the restart task.
-	restartTask.Schedule(time.Now().Add(delay))
+	log.Warningf("updates: restart triggered, will execute in %s", delay)
+	restartAt := time.Now().Add(delay)
+	restartTask.Schedule(restartAt)
+
+	// Set restartTime.
+	restartTimeLock.Lock()
+	defer restartTimeLock.Unlock()
+	restartTime = restartAt
+}
+
+// AbortRestart aborts a (delayed) restart.
+func AbortRestart() {
+	if restartPending.SetToIf(true, false) {
+		log.Warningf("updates: restart aborted")
+
+		// Cancel schedule.
+		restartTask.Schedule(time.Time{})
+	}
 }
 
 // TriggerRestartIfPending triggers an automatic restart, if one is pending.
 // This can be used to prepone a scheduled restart if the conditions are preferable.
 func TriggerRestartIfPending() {
 	if restartPending.IsSet() {
-		_ = automaticRestart(module.Ctx, nil)
+		restartTask.StartASAP()
 	}
 }
 
 // RestartNow immediately executes a restart.
 // This only works if the process is managed by portmaster-start.
 func RestartNow() {
-	_ = automaticRestart(module.Ctx, nil)
+	restartPending.Set()
+	restartTask.StartASAP()
 }
 
 func automaticRestart(_ context.Context, _ *modules.Task) error {
-	if restartTriggered.SetToIf(false, true) {
+	// Check if the restart is still scheduled.
-		log.Info("updates: initiating (automatic) restart")
+	if restartPending.IsNotSet() {
+		return nil
+	}
+
+	// Trigger restart.
+	if restartTriggered.SetToIf(false, true) {
+		log.Warning("updates: initiating (automatic) restart")
 
-		// Set restart pending to ensure IsRestarting() returns true.
-		restartPending.Set()
 		// Set restart exit code.
 		modules.SetExitStatusCode(RestartExitCode)
 		// Do not use a worker, as this would block itself here.

updates/upgrader.go

@@ -51,18 +51,19 @@ func initUpgrader() error {
 }
 
 func upgrader(_ context.Context, _ interface{}) error {
-	// like a lock, but discard additional runs
+	// Lock runs, but discard additional runs.
 	if !upgraderActive.SetToIf(false, true) {
 		return nil
 	}
 	defer upgraderActive.SetTo(false)
 
-	// upgrade portmaster-start
+	// Upgrade portmaster-start.
 	err := upgradePortmasterStart()
 	if err != nil {
 		log.Warningf("updates: failed to upgrade portmaster-start: %s", err)
 	}
 
+	// Upgrade based on binary.
 	binBaseName := strings.Split(filepath.Base(os.Args[0]), "_")[0]
 	switch binBaseName {
 	case "portmaster-core":
@@ -149,36 +150,45 @@ func upgradeCoreNotifyActionHandler(_ context.Context, n *notifications.Notifica
 }
 
 func upgradeHub() error {
-	if hubUpgradeStarted {
-		return nil
-	}
 	if spnHubUpdate != nil && !spnHubUpdate.UpgradeAvailable() {
 		return nil
 	}
 
-	// make identifier
+	// Make identifier for getting file from updater.
 	identifier := "hub/spn-hub" // identifier, use forward slash!
 	if onWindows {
 		identifier += exeExt
 	}
 
-	// get newest spn-hub
+	// Get newest spn-hub file.
 	newFile, err := GetPlatformFile(identifier)
 	if err != nil {
 		return err
 	}
 	spnHubUpdate = newFile
 
-	// check for new version
+	// Check if the new version is different.
 	if info.GetInfo().Version != spnHubUpdate.Version() {
-		// get random delay with up to three hours
+		// Get random delay with up to three hours.
 		delayMinutes, err := rng.Number(3 * 60)
 		if err != nil {
 			return err
 		}
 
-		DelayedRestart(time.Duration(delayMinutes) * time.Minute)
+		// Delay restart for at least one hour for preparations.
-		hubUpgradeStarted = true
+		DelayedRestart(time.Duration(delayMinutes+60) * time.Minute)
+
+		// Increase update checks in order to detect aborts better.
+		if !disableTaskSchedule {
+			updateTask.Repeat(10 * time.Minute)
+		}
+	} else {
+		AbortRestart()
+
+		// Set update task schedule back to normal.
+		if !disableTaskSchedule {
+			updateTask.Repeat(updateTaskRepeatDuration)
+		}
 	}
 
 	return nil