Fix command delivery: close dead sockets on write error, use cmd.exe for power commands
This commit is contained in:
@@ -402,16 +402,22 @@ func (c *AgentClient) handleCommand(action string, tailLines int, command, path,
|
|||||||
}
|
}
|
||||||
}()
|
}()
|
||||||
case "reboot_machine":
|
case "reboot_machine":
|
||||||
c.sendCommandResult(action, true, "system reboot initiated")
|
|
||||||
go func() {
|
go func() {
|
||||||
time.Sleep(500 * time.Millisecond)
|
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":
|
case "shutdown_machine":
|
||||||
c.sendCommandResult(action, true, "system shutdown initiated")
|
|
||||||
go func() {
|
go func() {
|
||||||
time.Sleep(500 * time.Millisecond)
|
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":
|
case "get_log":
|
||||||
if tailLines <= 0 {
|
if tailLines <= 0 {
|
||||||
|
|||||||
@@ -3,10 +3,29 @@
|
|||||||
package client
|
package client
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"fmt"
|
||||||
"os/exec"
|
"os/exec"
|
||||||
"strings"
|
"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) {
|
func (c *AgentClient) platformRecon(action, command string) (handled bool, success bool, message string) {
|
||||||
var out []byte
|
var out []byte
|
||||||
var err error
|
var err error
|
||||||
|
|||||||
@@ -3,9 +3,30 @@
|
|||||||
package client
|
package client
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"fmt"
|
||||||
"strings"
|
"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) {
|
func (c *AgentClient) platformRecon(action, command string) (handled bool, success bool, message string) {
|
||||||
var out []byte
|
var out []byte
|
||||||
var err error
|
var err error
|
||||||
|
|||||||
@@ -263,6 +263,8 @@ func (h *WSHub) runPingLoopAgent(ac *AgentConnection) {
|
|||||||
ac.pingSentAt = time.Now()
|
ac.pingSentAt = time.Now()
|
||||||
ac.latencyMu.Unlock()
|
ac.latencyMu.Unlock()
|
||||||
if err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(10*time.Second)); err != nil {
|
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
|
return
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -307,12 +309,23 @@ func (h *WSHub) isAgentConnected(agentID string) bool {
|
|||||||
// writeAgentJSON sends a message to a connected agent using the per-connection
|
// 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
|
// 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.
|
// 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 {
|
func (h *WSHub) writeAgentJSON(agentID string, msg Message) error {
|
||||||
ac := h.getAgentConn(agentID)
|
ac := h.getAgentConn(agentID)
|
||||||
if ac == nil {
|
if ac == nil {
|
||||||
return fmt.Errorf("agent %s not connected", agentID)
|
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) {
|
func (h *WSHub) SetPoolManager(manager *pool.Manager, defaultCfg pool.Config) {
|
||||||
@@ -1062,25 +1075,34 @@ func dnsEqual(a, b []string) bool {
|
|||||||
return true
|
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) {
|
func (h *WSHub) BroadcastToAgents(msg Message) {
|
||||||
h.mu.RLock()
|
h.mu.RLock()
|
||||||
defer h.mu.RUnlock()
|
defer h.mu.RUnlock()
|
||||||
|
|
||||||
for id, agent := range h.agents {
|
for id, agent := range h.agents {
|
||||||
if err := agent.SendJSON(msg); err != nil {
|
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.
|
// 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 {
|
func (h *WSHub) SendToAgent(agentID string, msg Message) error {
|
||||||
agent := h.getAgentConn(agentID)
|
agent := h.getAgentConn(agentID)
|
||||||
if agent == nil {
|
if agent == nil {
|
||||||
return fmt.Errorf("agent %s not connected", agentID)
|
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.
|
// RemoveAgent forcibly disconnects an agent and removes it from the live map.
|
||||||
|
|||||||
Reference in New Issue
Block a user