websocket: fix transfer conn concurrent op

lesismal · Nov 20, 2023 · 8512d68 · 8512d68
1 parent f02683a
commit 8512d68
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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 {