websocket: fix writeFrame order in concurrent WriteMessage opration

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)