websocket: fix protocol message handling

lesismal · Jul 16, 2024 · a7eb6fb · a7eb6fb
1 parent ebb0bdb
commit a7eb6fb
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/nbhttp/websocket/conn.go b/nbhttp/websocket/conn.go
@@ -361,6 +361,7 @@ func (c *Conn) Parse(data []byte) error {
 	var message []byte
 	var msgType MessageType
 	var protocolMessage []byte
+	var isProtocolMessage bool
 	var opcode MessageType
 	var ok, fin, compress bool
 	var totalFrameSize int
@@ -448,6 +449,7 @@ func (c *Conn) Parse(data []byte) error {
 					}
 				}
 			case PingMessage, PongMessage, CloseMessage:
+				isProtocolMessage = true
 				if bl > 0 {
 					protocolMessage = allocator.Malloc(len(body))
 					copy(protocolMessage, body)
@@ -482,9 +484,10 @@ func (c *Conn) Parse(data []byte) error {
 			c.handleDataFrame(msgType, fin, frame)
 			frame = nil
 		}
-		if fin && opcode == CloseMessage { //protocolMessage != nil {
+		if isProtocolMessage {
 			c.handleProtocolMessage(opcode, protocolMessage)
 			protocolMessage = nil
+			isProtocolMessage = false
 		}
 
 		// need more data