From 8512d684eb45119a06dbd78ae36bae3845cafc8b Mon Sep 17 00:00:00 2001
From: lesismal <lesismal@gmail.com>
Date: Tue, 21 Nov 2023 01:30:55 +0800
Subject: [PATCH] websocket: fix transfer conn concurrent op

---
 nbhttp/parser.go             | 2 +-
 nbhttp/websocket/conn.go     | 2 +-
 nbhttp/websocket/upgrader.go | 8 ++++++--
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/nbhttp/parser.go b/nbhttp/parser.go
index 92c45722..de5567a6 100644
--- a/nbhttp/parser.go
+++ b/nbhttp/parser.go
@@ -136,7 +136,6 @@ func (p *Parser) Read(data []byte) error {
 		return nil
 	}
 
-	var c byte
 	var start = 0
 	var offset = len(p.cache)
 	if offset > 0 {
@@ -161,6 +160,7 @@ UPGRADER:
 		return err
 	}
 
+	var c byte
 	for i := offset; i < len(data); i++ {
 		if p.Reader != nil {
 			goto UPGRADER
diff --git a/nbhttp/websocket/conn.go b/nbhttp/websocket/conn.go
index d9c03fb3..8d9afa5f 100644
--- a/nbhttp/websocket/conn.go
+++ b/nbhttp/websocket/conn.go
@@ -306,7 +306,7 @@ func (c *Conn) nextFrame() (opcode MessageType, body []byte, ok, fin, res1, res2
 func (c *Conn) Read(p *nbhttp.Parser, data []byte) error {
 	oldLen := len(c.buffer)
 	readLimit := c.Engine.ReadLimit
-	if readLimit > 0 && ((oldLen+len(data) > readLimit) || ((oldLen + len(c.message) + len(data)) > readLimit)) {
+	if readLimit > 0 && (oldLen+len(data) > readLimit) {
 		return nbhttp.ErrTooLong
 	}
 
diff --git a/nbhttp/websocket/upgrader.go b/nbhttp/websocket/upgrader.go
index f30c7b6e..7a02c42f 100644
--- a/nbhttp/websocket/upgrader.go
+++ b/nbhttp/websocket/upgrader.go
@@ -275,7 +275,9 @@ func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeade
 					parser.Execute = nbhttp.SyncExecutor
 				}
 				wsc = NewConn(u, vt, subprotocol, compress, false)
-				nbc.SetSession(wsc)
+				parser.Reader = wsc
+				parser.Engine = engine
+				nbc.SetSession(parser)
 				nbc.OnData(func(c *nbio.Conn, data []byte) {
 					defer func() {
 						if err := recover(); err != nil {
@@ -349,7 +351,9 @@ func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeade
 				parser.Execute = nbhttp.SyncExecutor
 			}
 			wsc = NewConn(u, nbc, subprotocol, compress, false)
-			nbc.SetSession(wsc)
+			parser.Reader = wsc
+			parser.Engine = engine
+			nbc.SetSession(parser)
 			nbc.OnData(func(c *nbio.Conn, data []byte) {
 				defer func() {
 					if err := recover(); err != nil {