fix: agent ws connect panic

This commit is contained in:
engigu
2026-01-01 12:28:34 +08:00
parent a5f3ec7d9b
commit 1c3b153742
2 changed files with 81 additions and 14 deletions
+54 -9
View File
@@ -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()