diff --git a/nbhttp/websocket/conn.go b/nbhttp/websocket/conn.go index 9a7b8232..46e540b0 100644 --- a/nbhttp/websocket/conn.go +++ b/nbhttp/websocket/conn.go @@ -530,6 +530,13 @@ func (c *Conn) WriteClose(code int, reason string) error { // WriteMessage . func (c *Conn) WriteMessage(messageType MessageType, data []byte) error { + c.mux.Lock() + defer c.mux.Unlock() + + if c.closed { + return net.ErrClosed + } + switch messageType { case TextMessage: case BinaryMessage: @@ -690,6 +697,13 @@ func (c *Conn) CloseAndClean(err error) { // WriteFrame . func (c *Conn) WriteFrame(messageType MessageType, sendOpcode, fin bool, data []byte) error { + c.mux.Lock() + defer c.mux.Unlock() + + if c.closed { + return net.ErrClosed + } + return c.writeFrame(messageType, sendOpcode, fin, data, false) } @@ -751,16 +765,8 @@ func (c *Conn) writeFrame(messageType MessageType, sendOpcode, fin bool, data [] buf[0] |= byte(0x80) } - c.mux.Lock() - if c.closed { - c.mux.Unlock() - c.Engine.BodyAllocator.Free(buf) - return net.ErrClosed - } - if c.sendQueue != nil { if c.sendQueueSize > 0 && len(c.sendQueue) >= int(c.sendQueueSize) { - c.mux.Unlock() c.Engine.BodyAllocator.Free(buf) return ErrMessageSendQuqueIsFull } @@ -769,7 +775,6 @@ func (c *Conn) writeFrame(messageType MessageType, sendOpcode, fin bool, data [] if isHead { c.sendQueue[0] = nil } - c.mux.Unlock() if isHead { go func() { @@ -807,7 +812,6 @@ func (c *Conn) writeFrame(messageType MessageType, sendOpcode, fin bool, data [] } return nil } - c.mux.Unlock() _, err := c.Conn.Write(buf) c.Engine.BodyAllocator.Free(buf)