From a7eb6fbc5690493f34ec7ccf948570ae0ae28bbf Mon Sep 17 00:00:00 2001
From: lesismal <lesismal@gmail.com>
Date: Tue, 16 Jul 2024 16:44:16 +0800
Subject: [PATCH] websocket: fix protocol message handling

---
 nbhttp/websocket/conn.go | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/nbhttp/websocket/conn.go b/nbhttp/websocket/conn.go
index 79cde0bc..88260dd8 100644
--- 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