From 092da058a5bb1517acba86273db32d0b8be517b7 Mon Sep 17 00:00:00 2001
From: Daniel <dhaavi@users.noreply.github.com>
Date: Tue, 11 Aug 2020 07:55:26 +0200
Subject: [PATCH] Add automatic upgrades for spn-hub

---
 core/events.go      |  4 +--
 updates/restart.go  | 48 ++++++++++++++++++++++++++++
 updates/upgrader.go | 77 ++++++++++++++++++++++++++++++++++++---------
 3 files changed, 112 insertions(+), 17 deletions(-)
 create mode 100644 updates/restart.go

diff --git a/core/events.go b/core/events.go
index 8f1fd446..5bc36504 100644
--- a/core/events.go
+++ b/core/events.go
@@ -5,12 +5,12 @@ import (
 
 	"github.com/safing/portbase/log"
 	"github.com/safing/portbase/modules"
+	"github.com/safing/portmaster/updates"
 )
 
 const (
 	eventShutdown = "shutdown"
 	eventRestart  = "restart"
-	restartCode   = 23
 )
 
 func registerEvents() {
@@ -43,7 +43,7 @@ func shutdown(ctx context.Context, _ interface{}) error {
 // restart restarts the Portmaster.
 func restart(ctx context.Context, data interface{}) error {
 	log.Info("core: user requested restart")
-	modules.SetExitStatusCode(restartCode)
+	modules.SetExitStatusCode(updates.RestartExitCode)
 	// Do not use a worker, as this would block itself here.
 	go modules.Shutdown() //nolint:errcheck
 	return nil
diff --git a/updates/restart.go b/updates/restart.go
new file mode 100644
index 00000000..551191bf
--- /dev/null
+++ b/updates/restart.go
@@ -0,0 +1,48 @@
+package updates
+
+import (
+	"context"
+	"time"
+
+	"github.com/safing/portbase/log"
+	"github.com/safing/portbase/modules"
+	"github.com/tevino/abool"
+)
+
+const (
+	// RestartExitCode will instruct portmaster-start to restart the process immediately, potentially with a new version.
+	RestartExitCode = 23
+)
+
+var (
+	restartTask      *modules.Task
+	restartPending   *abool.AtomicBool
+	restartTriggered *abool.AtomicBool
+)
+
+func init() {
+	restartTask = module.NewTask("automatic restart", automaticRestart)
+}
+
+func triggerRestart(delay time.Duration) {
+	restartPending.Set()
+	restartTask.Schedule(time.Now().Add(delay))
+}
+
+// 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)
+	}
+}
+
+func automaticRestart(_ context.Context, _ *modules.Task) error {
+	if restartTriggered.SetToIf(false, true) {
+		log.Info("updates: initiating automatic restart")
+		modules.SetExitStatusCode(RestartExitCode)
+		// Do not use a worker, as this would block itself here.
+		go modules.Shutdown() //nolint:errcheck
+	}
+
+	return nil
+}
diff --git a/updates/upgrader.go b/updates/upgrader.go
index e8862528..3c353b7f 100644
--- a/updates/upgrader.go
+++ b/updates/upgrader.go
@@ -11,16 +11,15 @@ import (
 	"strings"
 	"time"
 
-	"github.com/tevino/abool"
-
 	"github.com/google/renameio"
+	processInfo "github.com/shirou/gopsutil/process"
+	"github.com/tevino/abool"
 
 	"github.com/safing/portbase/info"
 	"github.com/safing/portbase/log"
 	"github.com/safing/portbase/notifications"
+	"github.com/safing/portbase/rng"
 	"github.com/safing/portbase/updater"
-
-	processInfo "github.com/shirou/gopsutil/process"
 )
 
 const (
@@ -29,12 +28,13 @@ const (
 )
 
 var (
-	// UpgradeCore specifies if portmaster-core should be upgraded.
-	UpgradeCore = true
-
 	upgraderActive = abool.NewBool(false)
-	pmCtrlUpdate   *updater.File
-	pmCoreUpdate   *updater.File
+
+	pmCtrlUpdate *updater.File
+	pmCoreUpdate *updater.File
+
+	spnHubUpdate      *updater.File
+	hubUpgradeStarted bool
 
 	rawVersionRegex = regexp.MustCompile(`^[0-9]+\.[0-9]+\.[0-9]+b?\*?$`)
 )
@@ -61,24 +61,35 @@ func upgrader(_ context.Context, _ interface{}) error {
 		log.Warningf("updates: failed to upgrade portmaster-start: %s", err)
 	}
 
-	if UpgradeCore {
+	binName := strings.TrimSuffix(
+		filepath.Base(os.Args[0]),
+		".exe",
+	)
+	switch binName {
+	case "portmaster-core":
 		err = upgradeCoreNotify()
 		if err != nil {
 			log.Warningf("updates: failed to notify about core upgrade: %s", err)
 		}
+
+	case "spn-hub":
+		err = upgradeHub()
+		if err != nil {
+			log.Warningf("updates: failed to initiate hub upgrade: %s", err)
+		}
 	}
 
 	return nil
 }
 
 func upgradeCoreNotify() error {
-	identifier := "core/portmaster-core" // identifier, use forward slash!
-	if onWindows {
-		identifier += exeExt
-	}
-
 	// check if we can upgrade
 	if pmCoreUpdate == nil || pmCoreUpdate.UpgradeAvailable() {
+		identifier := "core/portmaster-core" // identifier, use forward slash!
+		if onWindows {
+			identifier += exeExt
+		}
+
 		// get newest portmaster-core
 		new, err := GetPlatformFile(identifier)
 		if err != nil {
@@ -128,6 +139,42 @@ func upgradeCoreNotifyActionHandler(n *notifications.Notification) {
 	}
 }
 
+func upgradeHub() error {
+	if hubUpgradeStarted {
+		return nil
+	}
+
+	// check if we can upgrade
+	if spnHubUpdate == nil || spnHubUpdate.UpgradeAvailable() {
+		identifier := "hub/spn-hub" // identifier, use forward slash!
+		if onWindows {
+			identifier += exeExt
+		}
+
+		// get newest spn-hub
+		new, err := GetPlatformFile(identifier)
+		if err != nil {
+			return err
+		}
+		spnHubUpdate = new
+	} else {
+		return nil
+	}
+
+	if info.GetInfo().Version != spnHubUpdate.Version() {
+		// get random delay with up to three hours
+		delayMinutes, err := rng.Number(3 * 60)
+		if err != nil {
+			return err
+		}
+
+		triggerRestart(time.Duration(delayMinutes) * time.Minute)
+		hubUpgradeStarted = true
+	}
+
+	return nil
+}
+
 func upgradePortmasterStart() error {
 	filename := "portmaster-start"
 	if onWindows {