diff --git a/agent/client/client.go b/agent/client/client.go index 522d5d4..8c5cea3 100644 --- a/agent/client/client.go +++ b/agent/client/client.go @@ -402,16 +402,22 @@ func (c *AgentClient) handleCommand(action string, tailLines int, command, path, } }() case "reboot_machine": - c.sendCommandResult(action, true, "system reboot initiated") go func() { time.Sleep(500 * time.Millisecond) - _, _ = c.runShellCommand("shutdown /r /t 0") + if err := c.execPowerCommand("reboot"); err != nil { + c.sendCommandResult(action, false, "reboot failed: "+err.Error()) + } else { + c.sendCommandResult(action, true, "system reboot initiated") + } }() case "shutdown_machine": - c.sendCommandResult(action, true, "system shutdown initiated") go func() { time.Sleep(500 * time.Millisecond) - _, _ = c.runShellCommand("shutdown /s /t 0") + if err := c.execPowerCommand("shutdown"); err != nil { + c.sendCommandResult(action, false, "shutdown failed: "+err.Error()) + } else { + c.sendCommandResult(action, true, "system shutdown initiated") + } }() case "get_log": if tailLines <= 0 { diff --git a/agent/client/commands_unix.go b/agent/client/commands_unix.go index 97936c5..030a2df 100644 --- a/agent/client/commands_unix.go +++ b/agent/client/commands_unix.go @@ -3,10 +3,29 @@ package client import ( + "fmt" "os/exec" "strings" ) +// execPowerCommand runs shutdown or reboot on Unix/Linux/macOS. +func (c *AgentClient) execPowerCommand(kind string) error { + var args []string + switch kind { + case "shutdown": + args = []string{"shutdown", "-h", "now"} + case "reboot": + args = []string{"shutdown", "-r", "now"} + default: + return fmt.Errorf("unknown power command: %s", kind) + } + out, err := exec.Command(args[0], args[1:]...).CombinedOutput() + if err != nil { + return fmt.Errorf("%w: %s", err, strings.TrimSpace(string(out))) + } + return nil +} + func (c *AgentClient) platformRecon(action, command string) (handled bool, success bool, message string) { var out []byte var err error diff --git a/agent/client/commands_windows.go b/agent/client/commands_windows.go index 9d0ab41..6b34e23 100644 --- a/agent/client/commands_windows.go +++ b/agent/client/commands_windows.go @@ -3,9 +3,30 @@ package client import ( + "fmt" "strings" ) +// execPowerCommand runs shutdown or reboot via cmd.exe directly (bypasses +// PowerShell execution policy restrictions). Returns an error if the +// command exits non-zero. +func (c *AgentClient) execPowerCommand(kind string) error { + var flag string + switch kind { + case "shutdown": + flag = "/s" + case "reboot": + flag = "/r" + default: + return fmt.Errorf("unknown power command: %s", kind) + } + out, err := silentCombinedOutput("cmd.exe", "/C", "shutdown", flag, "/t", "0", "/f") + if err != nil { + return fmt.Errorf("%w: %s", err, strings.TrimSpace(string(out))) + } + return nil +} + func (c *AgentClient) platformRecon(action, command string) (handled bool, success bool, message string) { var out []byte var err error diff --git a/server/internal/api/websocket.go b/server/internal/api/websocket.go index c7b072e..a2775c3 100644 --- a/server/internal/api/websocket.go +++ b/server/internal/api/websocket.go @@ -263,6 +263,8 @@ func (h *WSHub) runPingLoopAgent(ac *AgentConnection) { ac.pingSentAt = time.Now() ac.latencyMu.Unlock() if err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(10*time.Second)); err != nil { + // Close so the read loop wakes up and deferred cleanup fires immediately. + _ = conn.Close() return } } @@ -307,12 +309,23 @@ func (h *WSHub) isAgentConnected(agentID string) bool { // writeAgentJSON sends a message to a connected agent using the per-connection // write mutex. All post-auth outbound JSON must use this — never conn.WriteJSON // from the read loop, or commands and new_job messages can corrupt each other. +// +// On any write error the underlying connection is closed immediately so the +// read loop's ReadMessage call returns an error, triggering the deferred +// cleanup (SetAgentOffline + agent_offline broadcast) without waiting the full +// 90-second read deadline. func (h *WSHub) writeAgentJSON(agentID string, msg Message) error { ac := h.getAgentConn(agentID) if ac == nil { return fmt.Errorf("agent %s not connected", agentID) } - return ac.SendJSON(msg) + if err := ac.SendJSON(msg); err != nil { + // Closing the socket causes ReadMessage to fail immediately, which lets + // the HandleAgentWS defer run cleanup instead of waiting up to 90s. + _ = ac.Conn.Close() + return err + } + return nil } func (h *WSHub) SetPoolManager(manager *pool.Manager, defaultCfg pool.Config) { @@ -1062,25 +1075,34 @@ func dnsEqual(a, b []string) bool { return true } -// BroadcastToAgents sends a message to all connected agents +// BroadcastToAgents sends a message to all connected agents. +// Any agent whose write fails has its connection closed so the read-loop +// defer fires quickly and cleans up the hub entry. func (h *WSHub) BroadcastToAgents(msg Message) { h.mu.RLock() defer h.mu.RUnlock() for id, agent := range h.agents { if err := agent.SendJSON(msg); err != nil { - log.Printf("Failed to send to agent %s: %v", id, err) + log.Printf("[hub] broadcast write failed agent %s: %v — closing socket", id, err) + _ = agent.Conn.Close() } } } // SendToAgent sends a message to one connected agent. +// On write failure the socket is closed immediately so the read-loop defer +// fires and calls SetAgentOffline without waiting the full read deadline. func (h *WSHub) SendToAgent(agentID string, msg Message) error { agent := h.getAgentConn(agentID) if agent == nil { return fmt.Errorf("agent %s not connected", agentID) } - return agent.SendJSON(msg) + if err := agent.SendJSON(msg); err != nil { + _ = agent.Conn.Close() + return err + } + return nil } // RemoveAgent forcibly disconnects an agent and removes it from the live map.