package controlplane import ( "errors" "fmt" "strings" "time" ) const restartCooldown = 60 * time.Second // AgentRestart is the safe lifecycle state rendered by the controller. It // contains no remote process details or credentials. type AgentRestart struct { RequestID string `json:"requestID"` AgentID string `json:"agentID"` Status string `json:"status"` RequestedAt string `json:"requestedAt"` ConfirmedAt string `json:"confirmedAt"` Failure string `json:"failure"` } // RestartAgent sends an authenticated command over the existing paired WSS // session. The controller deliberately never starts or terminates a remote // machine process itself. func (c *ControlPlane) RestartAgent(agentID string) (AgentRestart, error) { agentID = strings.TrimSpace(agentID) if agentID == "" { return AgentRestart{}, errors.New("Agent 标识无效") } c.expireStaleAgents() var status string if err := c.db.QueryRow("SELECT status FROM managed_agents WHERE agent_id = ?", agentID).Scan(&status); err != nil { return AgentRestart{}, errors.New("未找到 Agent") } client := c.clientFor(agentID) if status != "online" || client == nil { c.recordAudit("remote.restart_rejected", "controller", "agent:"+agentID, "failed", "目标 Agent 当前离线,无法通过已认证连接发送重启指令。") return AgentRestart{}, errors.New("目标 Agent 当前离线,无法重启;请等待其重新连接") } var existing AgentRestart err := c.db.QueryRow("SELECT request_id, agent_id, status, requested_at, COALESCE(confirmed_at, ''), COALESCE(failure, '') FROM agent_restart_operations WHERE agent_id = ?", agentID).Scan(&existing.RequestID, &existing.AgentID, &existing.Status, &existing.RequestedAt, &existing.ConfirmedAt, &existing.Failure) if err == nil { if existing.Status == "restarting" || existing.Status == "sent" { return existing, nil // idempotent repeated click; do not send another command. } if existing.Status == "reconnected" { if requested, parseErr := time.Parse(time.RFC3339Nano, existing.RequestedAt); parseErr == nil && time.Since(requested) < restartCooldown { return existing, errors.New("该 Agent 刚完成重启,请稍后再试") } } } requestID, err := randomHex(16) if err != nil { return AgentRestart{}, err } now := time.Now().UTC().Format(time.RFC3339Nano) op := AgentRestart{RequestID: requestID, AgentID: agentID, Status: "restarting", RequestedAt: now} if _, err := c.db.Exec(`INSERT INTO agent_restart_operations(request_id, agent_id, status, requested_at, confirmed_at, failure) VALUES (?, ?, ?, ?, '', '') ON CONFLICT(agent_id) DO UPDATE SET request_id=excluded.request_id, status=excluded.status, requested_at=excluded.requested_at, confirmed_at='', failure=''`, requestID, agentID, op.Status, now); err != nil { return AgentRestart{}, fmt.Errorf("无法记录重启操作: %w", err) } c.recordAudit("remote.restart_requested", "controller", "agent:"+agentID, "restarting", "已通过认证控制连接请求 Agent 重启。") if err := client.send(map[string]any{"type": "agent_restart", "schema_version": protocolVersion, "request_id": requestID, "sent_at": now}); err != nil { _, _ = c.db.Exec("UPDATE agent_restart_operations SET status='failed', failure=? WHERE agent_id=? AND request_id=?", "重启指令未发送成功。", agentID, requestID) c.recordAudit("remote.restart_delivery", "controller", "agent:"+agentID, "failed", "重启指令未发送成功。") return AgentRestart{}, errors.New("重启指令未发送成功,请刷新后重试") } c.recordAudit("remote.restart_delivery", "controller", "agent:"+agentID, "sent", "重启指令已下发,等待 Agent 确认。") return op, nil } func (c *ControlPlane) applyRestartAck(agentID string, message map[string]any) { requestID, _ := message["request_id"].(string) status, _ := message["status"].(string) errorText, _ := message["error"].(string) if requestID == "" || (status != "accepted" && status != "restarting" && status != "failed") { return } var expected string if c.db.QueryRow("SELECT request_id FROM agent_restart_operations WHERE agent_id=?", agentID).Scan(&expected) != nil || expected != requestID { return } if status == "failed" { detail := limit(sanitizeDisplayText(errorText), 1000) if detail == "" { detail = "Agent 未能启动替代进程。" } _, _ = c.db.Exec("UPDATE agent_restart_operations SET status='failed', confirmed_at=?, failure=? WHERE agent_id=?", time.Now().UTC().Format(time.RFC3339Nano), detail, agentID) c.recordAudit("remote.restart_confirmed", "agent:"+agentID, "agent:"+agentID, "failed", detail) return } _, _ = c.db.Exec("UPDATE agent_restart_operations SET status='restarting', confirmed_at=?, failure='' WHERE agent_id=?", time.Now().UTC().Format(time.RFC3339Nano), agentID) c.recordAudit("remote.restart_confirmed", "agent:"+agentID, "agent:"+agentID, "restarting", "Agent 已确认正在安全结束任务并重启。") } func (c *ControlPlane) markRestartReconnected(agentID string) { result, err := c.db.Exec("UPDATE agent_restart_operations SET status='reconnected', confirmed_at=?, failure='' WHERE agent_id=? AND status IN ('sent', 'restarting')", time.Now().UTC().Format(time.RFC3339Nano), agentID) if err == nil { if changed, _ := result.RowsAffected(); changed == 1 { c.recordAudit("remote.restart_reconnected", "agent:"+agentID, "agent:"+agentID, "reconnected", "替代 Agent 已重新通过认证并恢复连接。") } } }