From 4f01e214fd1eb4870e90a1f99480526ef3dd2299 Mon Sep 17 00:00:00 2001
From: wh1te909 <dcparsi@gmail.com>
Date: Sat, 29 Apr 2023 16:56:15 -0700
Subject: [PATCH] fix crash when update fails to download fixes #33

---
 agent/agent_unix.go    | 16 +++++++++-------
 agent/agent_windows.go |  9 ++++-----
 agent/rpc.go           |  6 +++++-
 3 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/agent/agent_unix.go b/agent/agent_unix.go
index 42f51ee..aa633b3 100644
--- a/agent/agent_unix.go
+++ b/agent/agent_unix.go
@@ -16,6 +16,7 @@ package agent
 
 import (
 	"bufio"
+	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -218,12 +219,12 @@ func (a *Agent) seEnforcing() bool {
 	return out.Status.Exit == 0 && strings.Contains(out.Stdout, "Enforcing")
 }
 
-func (a *Agent) AgentUpdate(url, inno, version string) {
+func (a *Agent) AgentUpdate(url, inno, version string) error {
 
 	self, err := os.Executable()
 	if err != nil {
 		a.Logger.Errorln("AgentUpdate() os.Executable():", err)
-		return
+		return err
 	}
 
 	// more reliable method to get current working directory than os.Getwd()
@@ -233,7 +234,7 @@ func (a *Agent) AgentUpdate(url, inno, version string) {
 	f, err := os.CreateTemp(cwd, "trmm")
 	if err != nil {
 		a.Logger.Errorln("AgentUpdate() os.CreateTemp:", err)
-		return
+		return err
 	}
 	defer os.Remove(f.Name())
 
@@ -252,12 +253,12 @@ func (a *Agent) AgentUpdate(url, inno, version string) {
 	if err != nil {
 		a.Logger.Errorln("AgentUpdate() download:", err)
 		f.Close()
-		return
+		return err
 	}
 	if r.IsError() {
 		a.Logger.Errorln("AgentUpdate() status code:", r.StatusCode())
 		f.Close()
-		return
+		return errors.New("err")
 	}
 
 	f.Close()
@@ -265,7 +266,7 @@ func (a *Agent) AgentUpdate(url, inno, version string) {
 	err = os.Rename(f.Name(), self)
 	if err != nil {
 		a.Logger.Errorln("AgentUpdate() os.Rename():", err)
-		return
+		return err
 	}
 
 	if runtime.GOOS == "linux" && a.seEnforcing() {
@@ -283,10 +284,11 @@ func (a *Agent) AgentUpdate(url, inno, version string) {
 	case "darwin":
 		opts.Command = "launchctl kickstart -k system/tacticalagent"
 	default:
-		return
+		return nil
 	}
 
 	a.CmdV2(opts)
+	return nil
 }
 
 func (a *Agent) AgentUninstall(code string) {
diff --git a/agent/agent_windows.go b/agent/agent_windows.go
index 5e8ee3b..fd6fb19 100644
--- a/agent/agent_windows.go
+++ b/agent/agent_windows.go
@@ -601,7 +601,7 @@ func (a *Agent) UninstallCleanup() {
 	os.RemoveAll(a.WinRunAsUserTmpDir)
 }
 
-func (a *Agent) AgentUpdate(url, inno, version string) {
+func (a *Agent) AgentUpdate(url, inno, version string) error {
 	time.Sleep(time.Duration(randRange(1, 15)) * time.Second)
 	a.KillHungUpdates()
 	time.Sleep(1 * time.Second)
@@ -620,13 +620,11 @@ func (a *Agent) AgentUpdate(url, inno, version string) {
 	r, err := rClient.R().SetOutput(updater).Get(url)
 	if err != nil {
 		a.Logger.Errorln(err)
-		CMD("net", []string{"start", winSvcName}, 10, false)
-		return
+		return err
 	}
 	if r.IsError() {
 		a.Logger.Errorln("Download failed with status code", r.StatusCode())
-		CMD("net", []string{"start", winSvcName}, 10, false)
-		return
+		return err
 	}
 
 	innoLogFile := filepath.Join(a.WinTmpDir, fmt.Sprintf("tacticalagent_update_v%s.txt", version))
@@ -638,6 +636,7 @@ func (a *Agent) AgentUpdate(url, inno, version string) {
 	}
 	cmd.Start()
 	time.Sleep(1 * time.Second)
+	return nil
 }
 
 func (a *Agent) osString() string {
diff --git a/agent/rpc.go b/agent/rpc.go
index 5eaf456..ed975d9 100644
--- a/agent/rpc.go
+++ b/agent/rpc.go
@@ -481,7 +481,11 @@ func (a *Agent) RunRPC() {
 				} else {
 					ret.Encode("ok")
 					msg.Respond(resp)
-					a.AgentUpdate(p.Data["url"], p.Data["inno"], p.Data["version"])
+					err := a.AgentUpdate(p.Data["url"], p.Data["inno"], p.Data["version"])
+					if err != nil {
+						atomic.StoreUint32(&agentUpdateLocker, 0)
+						return
+					}
 					atomic.StoreUint32(&agentUpdateLocker, 0)
 					nc.Flush()
 					nc.Close()