imrpove updating clients
This commit is contained in:
@@ -2,6 +2,8 @@ package handlers
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"crypto/sha256"
|
||||||
|
"encoding/hex"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"log"
|
"log"
|
||||||
@@ -15,7 +17,7 @@ import (
|
|||||||
"overlord-client/cmd/agent/wire"
|
"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)
|
sourcePath = strings.TrimSpace(sourcePath)
|
||||||
if sourcePath == "" {
|
if sourcePath == "" {
|
||||||
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: "missing update path"})
|
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()})
|
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 {
|
if err := wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: true}); err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
@@ -139,3 +145,46 @@ func samePath(left string, right string) bool {
|
|||||||
}
|
}
|
||||||
return samePathByOS(left, right)
|
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 (
|
import (
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"log"
|
||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
|
|
||||||
@@ -43,9 +44,14 @@ func runAgentUpdate(sourcePath string, enablePersistence bool) error {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if currentErr == nil && !samePath(currentExe, restartPath) {
|
if currentErr == nil {
|
||||||
if err := copyExecutableAtomic(sourcePath, currentExe); err != nil {
|
if !(startupEnabled && samePath(currentExe, startupPath)) {
|
||||||
return err
|
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)
|
_ = os.Remove(sourcePath)
|
||||||
|
|||||||
@@ -46,6 +46,9 @@ func runAgentUpdate(sourcePath string, enablePersistence bool) error {
|
|||||||
if startupEnabled {
|
if startupEnabled {
|
||||||
// Keep startup binary in sync when persistence/startup is enabled.
|
// Keep startup binary in sync when persistence/startup is enabled.
|
||||||
if !samePath(startupPath, currentExe) {
|
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 {
|
if err := copyExecutableAtomic(sourcePath, startupPath); err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
@@ -73,6 +76,9 @@ func runAgentUpdate(sourcePath string, enablePersistence bool) error {
|
|||||||
return launchDeferredUpdateScript(sourcePath, restartPath)
|
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 {
|
if err := copyExecutableAtomic(sourcePath, restartPath); err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1630,7 +1630,8 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in
|
|||||||
case "agent_update":
|
case "agent_update":
|
||||||
payload, _ := envelope["payload"].(map[string]interface{})
|
payload, _ := envelope["payload"].(map[string]interface{})
|
||||||
path, _ := payload["path"].(string)
|
path, _ := payload["path"].(string)
|
||||||
return HandleAgentUpdate(ctx, env, cmdID, path)
|
hash, _ := payload["hash"].(string)
|
||||||
|
return HandleAgentUpdate(ctx, env, cmdID, path, hash)
|
||||||
case "process_list":
|
case "process_list":
|
||||||
return HandleProcessList(ctx, env, cmdID)
|
return HandleProcessList(ctx, env, cmdID)
|
||||||
case "process_kill":
|
case "process_kill":
|
||||||
|
|||||||
@@ -81,7 +81,8 @@ export type CommandType =
|
|||||||
| "plugin_load_init"
|
| "plugin_load_init"
|
||||||
| "plugin_load_chunk"
|
| "plugin_load_chunk"
|
||||||
| "plugin_load_finish"
|
| "plugin_load_finish"
|
||||||
| "plugin_unload";
|
| "plugin_unload"
|
||||||
|
| "agent_update";
|
||||||
|
|
||||||
export type Command = {
|
export type Command = {
|
||||||
type: "command";
|
type: "command";
|
||||||
|
|||||||
@@ -1,3 +1,4 @@
|
|||||||
|
import { createHash } from "crypto";
|
||||||
import fs from "fs/promises";
|
import fs from "fs/promises";
|
||||||
import path from "path";
|
import path from "path";
|
||||||
import { v4 as uuidv4 } from "uuid";
|
import { v4 as uuidv4 } from "uuid";
|
||||||
@@ -298,12 +299,13 @@ export async function handleDeployRoutes(
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const hash = createHash("sha256").update(bytes).digest("hex");
|
||||||
target.ws.send(
|
target.ws.send(
|
||||||
encodeMessage({
|
encodeMessage({
|
||||||
type: "command",
|
type: "command",
|
||||||
commandType: "agent_update",
|
commandType: "agent_update",
|
||||||
id: uuidv4(),
|
id: uuidv4(),
|
||||||
payload: { path: destPath },
|
payload: { path: destPath, hash },
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user