From 1c3b153742483f25247d8413e192ab9ce0b12187 Mon Sep 17 00:00:00 2001
From: engigu <sayheya@qq.com>
Date: Thu, 1 Jan 2026 12:28:34 +0800
Subject: [PATCH] fix: agent ws connect panic

---
 internal/controllers/agent_controller.go | 32 ++++++++++--
 internal/services/agent_ws_service.go    | 63 ++++++++++++++++++++----
 2 files changed, 81 insertions(+), 14 deletions(-)

diff --git a/internal/controllers/agent_controller.go b/internal/controllers/agent_controller.go
index 73ef0ee..662aed4 100644
--- a/internal/controllers/agent_controller.go
+++ b/internal/controllers/agent_controller.go
@@ -407,17 +407,28 @@ func (c *AgentController) WSConnect(ctx *gin.Context) {
 // wsReadPump 读取消息
 func (c *AgentController) wsReadPump(ac *services.AgentConnection, agent *models.Agent) {
 	defer func() {
+		if r := recover(); r != nil {
+			logger.Errorf("[AgentWS] Agent #%d wsReadPump panic: %v", agent.ID, r)
+		}
 		c.wsManager.Unregister(agent.ID)
 	}()
 
-	ac.Conn.SetReadDeadline(time.Now().Add(90 * time.Second))
+	// 检查连接是否有效
+	if ac == nil || ac.IsClosed() {
+		logger.Warnf("[AgentWS] Agent #%d 连接无效", agent.ID)
+		return
+	}
+
+	ac.SetReadDeadline(time.Now().Add(90 * time.Second))
+	// 注意：SetPongHandler 需要直接访问 Conn，但这里我们在连接建立后立即设置
+	// 所以是安全的，因为此时连接还没有被其他 goroutine 关闭
 	ac.Conn.SetPongHandler(func(string) error {
-		ac.Conn.SetReadDeadline(time.Now().Add(90 * time.Second))
+		ac.SetReadDeadline(time.Now().Add(90 * time.Second))
 		return nil
 	})
 
 	for {
-		_, message, err := ac.Conn.ReadMessage()
+		_, message, err := ac.ReadMessage()
 		if err != nil {
 			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
 				logger.Warnf("[AgentWS] Agent #%d 读取错误: %v", agent.ID, err)
@@ -436,6 +447,12 @@ func (c *AgentController) wsReadPump(ac *services.AgentConnection, agent *models
 
 // wsWritePump 写入消息
 func (c *AgentController) wsWritePump(ac *services.AgentConnection) {
+	defer func() {
+		if r := recover(); r != nil {
+			logger.Errorf("[AgentWS] Agent #%d wsWritePump panic: %v", ac.AgentID, r)
+		}
+	}()
+
 	ticker := time.NewTicker(30 * time.Second)
 	defer ticker.Stop()
 
@@ -445,12 +462,17 @@ func (c *AgentController) wsWritePump(ac *services.AgentConnection) {
 			if !ok {
 				return
 			}
+			if ac.IsClosed() {
+				return
+			}
 			if err := ac.WriteMessage(message); err != nil {
 				return
 			}
 		case <-ticker.C:
-			ac.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
-			if err := ac.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
+			if ac.IsClosed() {
+				return
+			}
+			if err := ac.WritePing(); err != nil {
 				return
 			}
 		}
diff --git a/internal/services/agent_ws_service.go b/internal/services/agent_ws_service.go
index 9faca90..98dccf0 100644
--- a/internal/services/agent_ws_service.go
+++ b/internal/services/agent_ws_service.go
@@ -30,12 +30,13 @@ const (
 
 // AgentConnection Agent WebSocket 连接
 type AgentConnection struct {
-	AgentID   uint
-	IP        string
-	Conn      *websocket.Conn
-	Send      chan []byte
-	LastPing  time.Time
-	mu        sync.Mutex
+	AgentID  uint
+	IP       string
+	Conn     *websocket.Conn
+	Send     chan []byte
+	LastPing time.Time
+	closed   bool
+	mu       sync.Mutex
 }
 
 // WSMessage WebSocket 消息结构
@@ -266,27 +267,71 @@ func (m *AgentWSManager) cleanupLoop() {
 func (c *AgentConnection) Close() {
 	c.mu.Lock()
 	defer c.mu.Unlock()
+	if c.closed {
+		return
+	}
+	c.closed = true
 	if c.Conn != nil {
 		c.Conn.Close()
-		c.Conn = nil
 	}
 	if c.Send != nil {
 		close(c.Send)
-		c.Send = nil
 	}
 }
 
+// IsClosed 检查连接是否已关闭
+func (c *AgentConnection) IsClosed() bool {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.closed
+}
+
 // WriteMessage 写入消息
 func (c *AgentConnection) WriteMessage(data []byte) error {
 	c.mu.Lock()
 	defer c.mu.Unlock()
-	if c.Conn == nil {
+	if c.closed || c.Conn == nil {
 		return nil
 	}
 	c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
 	return c.Conn.WriteMessage(websocket.TextMessage, data)
 }
 
+// SetReadDeadline 设置读取超时
+func (c *AgentConnection) SetReadDeadline(t time.Time) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.closed || c.Conn == nil {
+		return nil
+	}
+	return c.Conn.SetReadDeadline(t)
+}
+
+// ReadMessage 读取消息
+func (c *AgentConnection) ReadMessage() (int, []byte, error) {
+	// 不加锁，因为 ReadMessage 是阻塞的
+	// 但需要先检查连接状态
+	c.mu.Lock()
+	if c.closed || c.Conn == nil {
+		c.mu.Unlock()
+		return 0, nil, websocket.ErrCloseSent
+	}
+	conn := c.Conn
+	c.mu.Unlock()
+	return conn.ReadMessage()
+}
+
+// WritePing 发送 ping 消息
+func (c *AgentConnection) WritePing() error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.closed || c.Conn == nil {
+		return nil
+	}
+	c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
+	return c.Conn.WriteMessage(websocket.PingMessage, nil)
+}
+
 // UpdatePing 更新心跳时间
 func (c *AgentConnection) UpdatePing() {
 	c.LastPing = time.Now()