- fix websocket tls dirty memory, solve: #22

lesismal · Jun 17, 2021 · 9ebb716 · 9ebb716
1 parent 1de2f1c
commit 9ebb716
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 11 deletions.
diff --git a/nbhttp/server.go b/nbhttp/server.go
@@ -505,8 +505,8 @@ func NewServerTLS(conf Config, handler http.Handler, messageHandlerExecutor func
 		if tlsConn, ok := parser.Processor.Conn().(*tls.Conn); ok {
 			svr.ParserExecutor(c.Hash(), func() {
 				tlsConn.Append(data)
+				buffer := parser.TLSBuffer
 				for {
-					buffer := parser.TLSBuffer
 					n, err := tlsConn.Read(buffer)
 					if err != nil {
 						c.CloseWithError(err)

diff --git a/nbhttp/websocket/upgrader.go b/nbhttp/websocket/upgrader.go
@@ -180,19 +180,20 @@ func (u *Upgrader) Read(p *nbhttp.Parser, data []byte) error {
 		}
 	}
 
-	if consumed && bufLen == 0 {
-		bufLen = len(u.buffer)
-		if bufLen > 0 {
-			var newBuf []byte
-			if bufLen < 1024 {
-				newBuf = u.Server.Malloc(1024)[:bufLen]
-			} else {
-				newBuf = u.Server.Malloc(bufLen)
+	if consumed {
+		if bufLen == 0 {
+			bufLen = len(u.buffer)
+			if bufLen > 0 {
+				newBuf := u.Server.Malloc(bufLen)
+				copy(newBuf, u.buffer)
+				u.buffer = newBuf
 			}
-			copy(newBuf, u.buffer)
-			u.buffer = newBuf
 		}
 		u.Server.Free(buffer)
+	} else if p.TLSBuffer != nil {
+		u.buffer = u.Server.Malloc(len(data))
+		copy(u.buffer, data)
+		u.Server.Free(buffer)
 	}
 
 	return nil