feat(updater): add updater service installation and update logic

yllada · yllada · commit b5795747a6f8 · 2025-12-17T15:16:05.000-05:00
diff --git a/agent/config/const.go b/agent/config/const.go
@@ -13,6 +13,10 @@ type ProtoPort struct {
 	TCP string
 }
 
+const (
+	SERVICE_UPDATER_NAME = "UTMStackUpdater"
+)
+
 var (
 	REPLACE_KEY string
 
diff --git a/agent/config/linux_amd64.go b/agent/config/linux_amd64.go
@@ -4,7 +4,7 @@
 package config
 
 var (
-	UpdaterSelf = "utmstack_updater_self%s"
 	ServiceFile = "utmstack_agent_service%s"
+	UpdaterFile = "utmstack_updater_service%s"
 	DependFiles = []string{"utmstack_agent_dependencies_linux.zip"}
 )
diff --git a/agent/config/linux_arm64.go b/agent/config/linux_arm64.go
@@ -4,7 +4,7 @@
 package config
 
 var (
-	UpdaterSelf = "utmstack_updater_self_arm64%s"
 	ServiceFile = "utmstack_agent_service_arm64%s"
+	UpdaterFile = "utmstack_updater_service%s"
 	DependFiles = []string{"utmstack_agent_dependencies_linux_arm64.zip"}
 )
diff --git a/agent/config/macos.go b/agent/config/macos.go
@@ -4,7 +4,7 @@
 package config
 
 var (
-	UpdaterSelf = "utmstack_updater_self%s"
 	ServiceFile = "utmstack_agent_service%s"
+	UpdaterFile = "utmstack_updater_service%s"
 	DependFiles = []string{}
 )
diff --git a/agent/config/windows_amd64.go b/agent/config/windows_amd64.go
@@ -4,7 +4,7 @@
 package config
 
 var (
-	UpdaterSelf = "utmstack_updater_self%s.exe"
 	ServiceFile = "utmstack_agent_service%s.exe"
+	UpdaterFile = "utmstack_updater_service%s"
 	DependFiles = []string{"utmstack_agent_dependencies_windows.zip"}
 )
diff --git a/agent/config/windows_arm64.go b/agent/config/windows_arm64.go
@@ -4,7 +4,7 @@
 package config
 
 var (
-	UpdaterSelf = "utmstack_updater_self_arm64%s.exe"
 	ServiceFile = "utmstack_agent_service_arm64%s.exe"
+	UpdaterFile = "utmstack_updater_service%s"
 	DependFiles = []string{"utmstack_agent_dependencies_windows_arm64.zip"}
 )
diff --git a/agent/logs/utmstack_agent.log b/agent/logs/utmstack_agent.log
diff --git a/agent/models/version.go b/agent/models/version.go
@@ -1,5 +1,6 @@
 package models
 
 type Version struct {
-	Version string `json:"version"`
+	Version        string `json:"version"`
+	UpdaterVersion string `json:"updater_version"`
 }
diff --git a/agent/serv/service.go b/agent/serv/service.go
@@ -2,8 +2,11 @@ package serv
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"os/signal"
+	"path/filepath"
+	"runtime"
 	"strconv"
 	"syscall"
 
@@ -46,6 +49,8 @@ func (p *program) run() {
 		utils.Logger.ErrorF("error migrating logs table: %v", err)
 	}
 
+	ensureUpdaterServiceInstalled()
+
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 	ctx = metadata.AppendToOutgoingContext(ctx, "key", cnf.AgentKey)
@@ -68,3 +73,52 @@ func (p *program) run() {
 	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
 	<-signals
 }
+
+func ensureUpdaterServiceInstalled() {
+	isInstalled, err := utils.CheckIfServiceIsInstalled("UTMStackUpdater")
+	if err != nil {
+		utils.Logger.ErrorF("error checking if updater service is installed: %v", err)
+		return
+	}
+
+	if isInstalled {
+		utils.Logger.Info("updater service is already installed")
+		return
+	}
+
+	utils.Logger.Info("updater service not found, installing...")
+
+	updaterPath := filepath.Join(utils.GetMyPath(), fmt.Sprintf(config.UpdaterFile, ""))
+
+	if !utils.CheckIfPathExist(updaterPath) {
+
+		cnf, err := config.GetCurrentConfig()
+		if err != nil {
+			utils.Logger.ErrorF("error getting config to download updater: %v", err)
+			return
+		}
+
+		updaterBinary := fmt.Sprintf(config.UpdaterFile, "")
+		if err := utils.DownloadFile(fmt.Sprintf(config.DependUrl, cnf.Server, config.DependenciesPort, updaterBinary), map[string]string{}, updaterBinary, utils.GetMyPath(), cnf.SkipCertValidation); err != nil {
+			utils.Logger.ErrorF("error downloading updater binary: %v", err)
+			return
+		}
+
+		if runtime.GOOS == "linux" || runtime.GOOS == "darwin" {
+			if err := utils.Execute("chmod", utils.GetMyPath(), "755", updaterBinary); err != nil {
+				utils.Logger.ErrorF("error setting permissions on updater: %v", err)
+				return
+			}
+		}
+
+		utils.Logger.Info("updater binary downloaded successfully")
+	}
+
+	err = utils.Execute(updaterPath, utils.GetMyPath(), "install")
+	if err != nil {
+		utils.Logger.ErrorF("error installing updater service: %v", err)
+		return
+	}
+
+	utils.Logger.Info("updater service installed successfully")
+}
diff --git a/agent/updates/dependencies.go b/agent/updates/dependencies.go
@@ -16,6 +16,17 @@ func DownloadFirstDependencies(address string, insecure bool) error {
 		return fmt.Errorf("error downloading version.json : %v", err)
 	}
 
+	updaterBinary := fmt.Sprintf(config.UpdaterFile, "")
+	if err := utils.DownloadFile(fmt.Sprintf(config.DependUrl, address, config.DependenciesPort, updaterBinary), map[string]string{}, updaterBinary, utils.GetMyPath(), insecure); err != nil {
+		return fmt.Errorf("error downloading updater binary %s: %v", updaterBinary, err)
+	}
+
+	if runtime.GOOS == "linux" || runtime.GOOS == "darwin" {
+		if err := utils.Execute("chmod", utils.GetMyPath(), "755", updaterBinary); err != nil {
+			return fmt.Errorf("error setting permissions on updater: %v", err)
+		}
+	}
+
 	dependFiles := config.DependFiles
 	for _, file := range dependFiles {
 		if err := utils.DownloadFile(fmt.Sprintf(config.DependUrl, address, config.DependenciesPort, file), map[string]string{}, file, utils.GetMyPath(), insecure); err != nil {
@@ -38,12 +49,6 @@ func handleDependenciesPostDownload(dependencies []string) error {
 				return fmt.Errorf("error unzipping dependencies: %v", err)
 			}
 
-			if runtime.GOOS == "linux" || runtime.GOOS == "darwin" {
-				if err := utils.Execute("chmod", utils.GetMyPath(), "-R", "755", fmt.Sprintf(config.UpdaterSelf, "")); err != nil {
-					return fmt.Errorf("error executing chmod on %s: %v", fmt.Sprintf(config.UpdaterSelf, ""), err)
-				}
-			}
-
 			if err := os.Remove(filepath.Join(utils.GetMyPath(), file)); err != nil {
 				return fmt.Errorf("error removing file %s: %v", file, err)
 			}
diff --git a/agent/updates/update.go b/agent/updates/update.go
@@ -40,32 +40,115 @@ func UpdateDependencies(cnf *config.Config) {
 			continue
 		}
 
-		if newVersion.Version != currentVersion.Version {
-			utils.Logger.Info("New version of agent found: %s", newVersion.Version)
-			if err := utils.DownloadFile(fmt.Sprintf(config.DependUrl, cnf.Server, config.DependenciesPort, fmt.Sprintf(config.ServiceFile, "")), map[string]string{}, fmt.Sprintf(config.ServiceFile, "_new"), utils.GetMyPath(), cnf.SkipCertValidation); err != nil {
-				utils.Logger.ErrorF("error downloading agent: %v", err)
-				continue
-			}
-
-			currentVersion = newVersion
-			err = utils.WriteJSON(config.VersionPath, &currentVersion)
-			if err != nil {
-				utils.Logger.ErrorF("error writing version file: %v", err)
+		if newVersion.UpdaterVersion != currentVersion.UpdaterVersion {
+			utils.Logger.Info("New version of updater found: %s", newVersion.UpdaterVersion)
+			if err := utils.DownloadFile(fmt.Sprintf(config.DependUrl, cnf.Server, config.DependenciesPort, fmt.Sprintf(config.UpdaterFile, "")), map[string]string{}, fmt.Sprintf(config.UpdaterFile, "_new"), utils.GetMyPath(), cnf.SkipCertValidation); err != nil {
+				utils.Logger.ErrorF("error downloading updater: %v", err)
 				continue
 			}
 
 			if runtime.GOOS == "linux" || runtime.GOOS == "darwin" {
-				if err = utils.Execute("chmod", utils.GetMyPath(), "-R", "755", filepath.Join(utils.GetMyPath(), fmt.Sprintf(config.ServiceFile, "_new"))); err != nil {
+				if err = utils.Execute("chmod", utils.GetMyPath(), "-R", "755", filepath.Join(utils.GetMyPath(), fmt.Sprintf(config.UpdaterFile, "_new"))); err != nil {
 					utils.Logger.ErrorF("error executing chmod: %v", err)
 				}
 			}
 
-			err = utils.Execute(filepath.Join(utils.GetMyPath(), fmt.Sprintf(config.UpdaterSelf, "")), utils.GetMyPath())
+			utils.Logger.Info("Starting updater update process...")
+			err = runUpdateProcess()
 			if err != nil {
-				utils.Logger.ErrorF("error executing updater: %v", err)
+				utils.Logger.ErrorF("error updating updater: %v", err)
+				os.Remove(filepath.Join(utils.GetMyPath(), "version_new.json"))
+				os.Remove(filepath.Join(utils.GetMyPath(), fmt.Sprintf(config.UpdaterFile, "_new")))
+			} else {
+				utils.Logger.Info("Updater update completed successfully")
+				if utils.CheckIfPathExist(config.VersionPath) {
+					err := utils.ReadJson(config.VersionPath, &currentVersion)
+					if err != nil {
+						utils.Logger.ErrorF("error reading updated version file: %v", err)
+					}
+				}
 			}
+		} else {
+			os.Remove(filepath.Join(utils.GetMyPath(), "version_new.json"))
+		}
+	}
+}
+
+func runUpdateProcess() error {
+	path := utils.GetMyPath()
+
+	newBin := fmt.Sprintf(config.UpdaterFile, "_new")
+	oldBin := fmt.Sprintf(config.UpdaterFile, "")
+	backupBin := fmt.Sprintf(config.UpdaterFile, ".old")
+
+	updaterNew := filepath.Join(path, newBin)
+	if _, err := os.Stat(updaterNew); err != nil {
+		return fmt.Errorf("no _new binary found to update")
+	}
+
+	if err := utils.StopService(config.SERVICE_UPDATER_NAME); err != nil {
+		return fmt.Errorf("error stopping updater: %v", err)
+	}
+
+	time.Sleep(10 * time.Second)
+
+	backupPath := filepath.Join(path, backupBin)
+	if utils.CheckIfPathExist(backupPath) {
+		utils.Logger.Info("Removing previous backup: %s", backupPath)
+		if err := os.Remove(backupPath); err != nil {
+			utils.Logger.ErrorF("could not remove old backup: %v", err)
 		}
+	}
 
-		os.Remove(filepath.Join(utils.GetMyPath(), "version_new.json"))
+	if err := os.Rename(filepath.Join(path, oldBin), backupPath); err != nil {
+		return fmt.Errorf("error backing up old binary: %v", err)
 	}
+
+	if err := os.Rename(filepath.Join(path, newBin), filepath.Join(path, oldBin)); err != nil {
+		os.Rename(backupPath, filepath.Join(path, oldBin))
+		return fmt.Errorf("error renaming new binary: %v", err)
+	}
+
+	if err := utils.StartService(config.SERVICE_UPDATER_NAME); err != nil {
+		rollbackUpdater(oldBin, backupBin, path)
+		return fmt.Errorf("error starting updater: %v", err)
+	}
+
+	time.Sleep(30 * time.Second)
+
+	isHealthy, err := utils.CheckIfServiceIsActive(config.SERVICE_UPDATER_NAME)
+	if err != nil || !isHealthy {
+		utils.Logger.Info("New version failed health check, rolling back...")
+		rollbackUpdater(oldBin, backupBin, path)
+		return fmt.Errorf("rollback completed: new version failed health check")
+	}
+
+	utils.Logger.Info("Health check passed for updater")
+
+	versionNewPath := filepath.Join(path, "version_new.json")
+	versionPath := filepath.Join(path, "version.json")
+	if utils.CheckIfPathExist(versionNewPath) {
+		if err := os.Rename(versionNewPath, versionPath); err != nil {
+			utils.Logger.ErrorF("error updating version file: %v", err)
+		} else {
+			utils.Logger.Info("Version file updated successfully")
+		}
+	}
+
+	return nil
+}
+
+func rollbackUpdater(currentBin, backupBin, path string) {
+	utils.Logger.Info("Rolling back updater to previous version...")
+
+	utils.StopService(config.SERVICE_UPDATER_NAME)
+	time.Sleep(5 * time.Second)
+
+	os.Remove(filepath.Join(path, currentBin))
+	os.Rename(filepath.Join(path, backupBin), filepath.Join(path, currentBin))
+
+	utils.StartService(config.SERVICE_UPDATER_NAME)
+	os.Remove(filepath.Join(path, "version_new.json"))
+
+	utils.Logger.Info("Rollback completed for updater")
 }
diff --git a/agent/utils/services.go b/agent/utils/services.go
@@ -4,8 +4,66 @@ import (
 	"fmt"
 	"os"
 	"runtime"
+	"strings"
 )
 
+func CheckIfServiceIsActive(serv string) (bool, error) {
+	var errB bool
+	var output string
+	path := GetMyPath()
+
+	switch runtime.GOOS {
+	case "windows":
+		output, errB = ExecuteWithResult("sc", path, "query", serv)
+	case "linux":
+		output, errB = ExecuteWithResult("systemctl", path, "is-active", serv)
+	case "darwin":
+		output, errB = ExecuteWithResult("launchctl", path, "list", serv)
+	default:
+		return false, fmt.Errorf("unknown operating system")
+	}
+
+	if errB {
+		return false, nil
+	}
+
+	serviceStatus := strings.ToLower(strings.TrimSpace(output))
+
+	switch runtime.GOOS {
+	case "windows":
+		return strings.Contains(serviceStatus, "running"), nil
+	case "linux":
+		return serviceStatus == "active", nil
+	case "darwin":
+		return true, nil
+	default:
+		return false, fmt.Errorf("unsupported operating system")
+	}
+}
+
+func StartService(name string) error {
+	path := GetMyPath()
+	switch runtime.GOOS {
+	case "windows":
+		err := Execute("sc", path, "start", name)
+		if err != nil {
+			return fmt.Errorf("error starting service: %v", err)
+		}
+	case "linux":
+		err := Execute("systemctl", path, "start", name)
+		if err != nil {
+			return fmt.Errorf("error starting service: %v", err)
+		}
+	case "darwin":
+		plistPath := fmt.Sprintf("/Library/LaunchDaemons/%s.plist", name)
+		err := Execute("launchctl", path, "load", plistPath)
+		if err != nil {
+			return fmt.Errorf("error starting macOS service: %v", err)
+		}
+	}
+	return nil
+}
+
 func StopService(name string) error {
 	path := GetMyPath()
 	switch runtime.GOOS {
diff --git a/agent/version.json b/agent/version.json
@@ -1,3 +1,4 @@
 {
-    "version": "11.0.1"
+    "version": "11.0.1",
+    "updater_version": "1.0.0"
 }

Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,10 @@ type ProtoPort struct {`
`13`	`13`	`TCP string`
`14`	`14`	`}`
`15`	`15`
	`16`	`+const (`
	`17`	`+ SERVICE_UPDATER_NAME = "UTMStackUpdater"`
	`18`	`+)`
	`19`	`+`
`16`	`20`	`var (`
`17`	`21`	`REPLACE_KEY string`
`18`	`22`
Original file line number	Diff line number	Diff line change
`@@ -4,7 +4,7 @@`
`4`	`4`	`package config`
`5`	`5`
`6`	`6`	`var (`
`7`		`- UpdaterSelf = "utmstack_updater_self%s"`
`8`	`7`	`ServiceFile = "utmstack_agent_service%s"`
	`8`	`+ UpdaterFile = "utmstack_updater_service%s"`
`9`	`9`	`DependFiles = []string{"utmstack_agent_dependencies_linux.zip"}`
`10`	`10`	`)`
Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,6 @@`
`1`	`1`	`package models`
`2`	`2`
`3`	`3`	`type Version struct {`
`4`		- Version string `json:"version"`
	`4`	+ Version string `json:"version"`
	`5`	+ UpdaterVersion string `json:"updater_version"`
`5`	`6`	`}`