Refer to the implementation of onebot.go to protect concurrent writes.
This commit is contained in:
parent
fbd3ea34ee
commit
f7b89bbd98
1 changed files with 42 additions and 16 deletions
|
|
@ -26,6 +26,31 @@ var chatHTML []byte
|
||||||
//go:embed logo.jpg
|
//go:embed logo.jpg
|
||||||
var logoImage []byte
|
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
|
// Channel implements the Channel interface for WebSocket connections
|
||||||
type Channel struct {
|
type Channel struct {
|
||||||
config config.WebSocketConfig
|
config config.WebSocketConfig
|
||||||
|
|
@ -34,7 +59,7 @@ type Channel struct {
|
||||||
allowList []string
|
allowList []string
|
||||||
server *http.Server
|
server *http.Server
|
||||||
upgrader websocket.Upgrader
|
upgrader websocket.Upgrader
|
||||||
clients sync.Map // map[string]*websocket.Conn
|
clients sync.Map // map[string]*clientConn
|
||||||
clientsMu sync.RWMutex
|
clientsMu sync.RWMutex
|
||||||
ctx context.Context
|
ctx context.Context
|
||||||
cancel context.CancelFunc
|
cancel context.CancelFunc
|
||||||
|
|
@ -182,8 +207,8 @@ func (c *Channel) Stop(ctx context.Context) error {
|
||||||
|
|
||||||
// Close all client connections
|
// Close all client connections
|
||||||
c.clients.Range(func(key, value any) bool {
|
c.clients.Range(func(key, value any) bool {
|
||||||
if conn, ok := value.(*websocket.Conn); ok {
|
if client, ok := value.(*clientConn); ok {
|
||||||
conn.Close()
|
client.close()
|
||||||
}
|
}
|
||||||
c.clients.Delete(key)
|
c.clients.Delete(key)
|
||||||
return true
|
return true
|
||||||
|
|
@ -235,8 +260,8 @@ func (c *Channel) Send(ctx context.Context, msg bus.OutboundMessage) error {
|
||||||
if msg.ChatID != "" && msg.ChatID != "broadcast" {
|
if msg.ChatID != "" && msg.ChatID != "broadcast" {
|
||||||
// Send to specific client
|
// Send to specific client
|
||||||
if conn, ok := c.clients.Load(msg.ChatID); ok {
|
if conn, ok := c.clients.Load(msg.ChatID); ok {
|
||||||
if wsConn, ok := conn.(*websocket.Conn); ok {
|
if client, ok := conn.(*clientConn); ok {
|
||||||
err := wsConn.WriteMessage(websocket.TextMessage, data)
|
err := client.writeMessage(websocket.TextMessage, data)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
// Connection may be dead, clean it up
|
// Connection may be dead, clean it up
|
||||||
logger.WarnCF("websocket", "Failed to send to client, removing connection", map[string]any{
|
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
|
var lastErr error
|
||||||
deadClients := make([]any, 0)
|
deadClients := make([]any, 0)
|
||||||
c.clients.Range(func(key, value any) bool {
|
c.clients.Range(func(key, value any) bool {
|
||||||
if conn, ok := value.(*websocket.Conn); ok {
|
if client, ok := value.(*clientConn); ok {
|
||||||
if err := conn.WriteMessage(websocket.TextMessage, data); err != nil {
|
if err := client.writeMessage(websocket.TextMessage, data); err != nil {
|
||||||
logger.WarnCF("websocket", "Failed to send to client", map[string]any{
|
logger.WarnCF("websocket", "Failed to send to client", map[string]any{
|
||||||
"client": key,
|
"client": key,
|
||||||
"error": err.Error(),
|
"error": err.Error(),
|
||||||
|
|
@ -300,20 +325,21 @@ func (c *Channel) handleWebSocket(w http.ResponseWriter, r *http.Request) {
|
||||||
|
|
||||||
// Generate client ID from remote address
|
// Generate client ID from remote address
|
||||||
clientID := r.RemoteAddr
|
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{
|
logger.InfoCF("websocket", "New client connected", map[string]any{
|
||||||
"client_id": clientID,
|
"client_id": clientID,
|
||||||
})
|
})
|
||||||
|
|
||||||
// Handle client messages
|
// Handle client messages
|
||||||
go c.handleClient(clientID, conn)
|
go c.handleClient(clientID, client)
|
||||||
}
|
}
|
||||||
|
|
||||||
// handleClient processes messages from a WebSocket 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() {
|
defer func() {
|
||||||
conn.Close()
|
client.close()
|
||||||
c.clients.Delete(clientID)
|
c.clients.Delete(clientID)
|
||||||
logger.InfoCF("websocket", "Client disconnected", map[string]any{
|
logger.InfoCF("websocket", "Client disconnected", map[string]any{
|
||||||
"client_id": clientID,
|
"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
|
// Set up ping/pong handlers for connection health check
|
||||||
conn.SetReadDeadline(time.Now().Add(60 * time.Second))
|
client.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
|
||||||
conn.SetPongHandler(func(string) error {
|
client.conn.SetPongHandler(func(string) error {
|
||||||
conn.SetReadDeadline(time.Now().Add(60 * time.Second))
|
client.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
|
||||||
return nil
|
return nil
|
||||||
})
|
})
|
||||||
|
|
||||||
|
|
@ -338,7 +364,7 @@ func (c *Channel) handleClient(clientID string, conn *websocket.Conn) {
|
||||||
// Read messages in a goroutine
|
// Read messages in a goroutine
|
||||||
go func() {
|
go func() {
|
||||||
for {
|
for {
|
||||||
_, message, err := conn.ReadMessage()
|
_, message, err := client.conn.ReadMessage()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
|
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
|
||||||
logger.ErrorCF("websocket", "WebSocket error", map[string]any{
|
logger.ErrorCF("websocket", "WebSocket error", map[string]any{
|
||||||
|
|
@ -358,7 +384,7 @@ func (c *Channel) handleClient(clientID string, conn *websocket.Conn) {
|
||||||
return
|
return
|
||||||
case <-pingTicker.C:
|
case <-pingTicker.C:
|
||||||
// Send ping to check connection health
|
// 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{
|
logger.WarnCF("websocket", "Failed to send ping", map[string]any{
|
||||||
"client_id": clientID,
|
"client_id": clientID,
|
||||||
"error": err.Error(),
|
"error": err.Error(),
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue