From f7b89bbd9887d99fd7b5e198a471c4d66f003bec Mon Sep 17 00:00:00 2001 From: likeaturtle Date: Mon, 23 Feb 2026 09:43:47 +0800 Subject: [PATCH] Refer to the implementation of onebot.go to protect concurrent writes. --- pkg/channels/websocket/channel.go | 58 ++++++++++++++++++++++--------- 1 file changed, 42 insertions(+), 16 deletions(-) diff --git a/pkg/channels/websocket/channel.go b/pkg/channels/websocket/channel.go index e06bea39f..d44de4cd3 100644 --- a/pkg/channels/websocket/channel.go +++ b/pkg/channels/websocket/channel.go @@ -26,6 +26,31 @@ var chatHTML []byte //go:embed logo.jpg var logoImage []byte +// clientConn wraps a WebSocket connection with a write mutex for safe concurrent writes +type clientConn struct { + conn *websocket.Conn + writeMu sync.Mutex +} + +// writeMessage safely writes a message to the WebSocket connection +func (c *clientConn) writeMessage(messageType int, data []byte) error { + c.writeMu.Lock() + defer c.writeMu.Unlock() + return c.conn.WriteMessage(messageType, data) +} + +// writeControl safely writes a control message to the WebSocket connection +func (c *clientConn) writeControl(messageType int, data []byte, deadline time.Time) error { + c.writeMu.Lock() + defer c.writeMu.Unlock() + return c.conn.WriteControl(messageType, data, deadline) +} + +// close closes the underlying WebSocket connection +func (c *clientConn) close() error { + return c.conn.Close() +} + // Channel implements the Channel interface for WebSocket connections type Channel struct { config config.WebSocketConfig @@ -34,7 +59,7 @@ type Channel struct { allowList []string server *http.Server upgrader websocket.Upgrader - clients sync.Map // map[string]*websocket.Conn + clients sync.Map // map[string]*clientConn clientsMu sync.RWMutex ctx context.Context cancel context.CancelFunc @@ -182,8 +207,8 @@ func (c *Channel) Stop(ctx context.Context) error { // Close all client connections c.clients.Range(func(key, value any) bool { - if conn, ok := value.(*websocket.Conn); ok { - conn.Close() + if client, ok := value.(*clientConn); ok { + client.close() } c.clients.Delete(key) return true @@ -235,8 +260,8 @@ func (c *Channel) Send(ctx context.Context, msg bus.OutboundMessage) error { if msg.ChatID != "" && msg.ChatID != "broadcast" { // Send to specific client if conn, ok := c.clients.Load(msg.ChatID); ok { - if wsConn, ok := conn.(*websocket.Conn); ok { - err := wsConn.WriteMessage(websocket.TextMessage, data) + if client, ok := conn.(*clientConn); ok { + err := client.writeMessage(websocket.TextMessage, data) if err != nil { // Connection may be dead, clean it up logger.WarnCF("websocket", "Failed to send to client, removing connection", map[string]any{ @@ -256,8 +281,8 @@ func (c *Channel) Send(ctx context.Context, msg bus.OutboundMessage) error { var lastErr error deadClients := make([]any, 0) c.clients.Range(func(key, value any) bool { - if conn, ok := value.(*websocket.Conn); ok { - if err := conn.WriteMessage(websocket.TextMessage, data); err != nil { + if client, ok := value.(*clientConn); ok { + if err := client.writeMessage(websocket.TextMessage, data); err != nil { logger.WarnCF("websocket", "Failed to send to client", map[string]any{ "client": key, "error": err.Error(), @@ -300,20 +325,21 @@ func (c *Channel) handleWebSocket(w http.ResponseWriter, r *http.Request) { // Generate client ID from remote address clientID := r.RemoteAddr - c.clients.Store(clientID, conn) + client := &clientConn{conn: conn} + c.clients.Store(clientID, client) logger.InfoCF("websocket", "New client connected", map[string]any{ "client_id": clientID, }) // Handle client messages - go c.handleClient(clientID, conn) + go c.handleClient(clientID, client) } // handleClient processes messages from a WebSocket client -func (c *Channel) handleClient(clientID string, conn *websocket.Conn) { +func (c *Channel) handleClient(clientID string, client *clientConn) { defer func() { - conn.Close() + client.close() c.clients.Delete(clientID) logger.InfoCF("websocket", "Client disconnected", map[string]any{ "client_id": clientID, @@ -321,9 +347,9 @@ func (c *Channel) handleClient(clientID string, conn *websocket.Conn) { }() // Set up ping/pong handlers for connection health check - conn.SetReadDeadline(time.Now().Add(60 * time.Second)) - conn.SetPongHandler(func(string) error { - conn.SetReadDeadline(time.Now().Add(60 * time.Second)) + client.conn.SetReadDeadline(time.Now().Add(60 * time.Second)) + client.conn.SetPongHandler(func(string) error { + client.conn.SetReadDeadline(time.Now().Add(60 * time.Second)) return nil }) @@ -338,7 +364,7 @@ func (c *Channel) handleClient(clientID string, conn *websocket.Conn) { // Read messages in a goroutine go func() { for { - _, message, err := conn.ReadMessage() + _, message, err := client.conn.ReadMessage() if err != nil { if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) { logger.ErrorCF("websocket", "WebSocket error", map[string]any{ @@ -358,7 +384,7 @@ func (c *Channel) handleClient(clientID string, conn *websocket.Conn) { return case <-pingTicker.C: // Send ping to check connection health - if err := conn.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(10*time.Second)); err != nil { + if err := client.writeControl(websocket.PingMessage, []byte{}, time.Now().Add(10*time.Second)); err != nil { logger.WarnCF("websocket", "Failed to send ping", map[string]any{ "client_id": clientID, "error": err.Error(),