websocket: pre-calculated length (#3284)

tsctx · web-flow · commit 13523fd7fe66 · 2024-05-21T22:35:50.000+09:00
diff --git a/lib/web/websocket/receiver.js b/lib/web/websocket/receiver.js
@@ -11,6 +11,7 @@ const {
   failWebsocketConnection,
   websocketMessageReceived,
   utf8Decode,
+  bufferConcat,
   isControlFrame,
   isTextBinaryFrame,
   isContinuationFrame
@@ -33,6 +34,7 @@ class ByteParser extends Writable {
 
   #info = {}
   #fragments = []
+  #fragmentsBytes = 0
 
   /** @type {Map<string, PerMessageDeflate>} */
   #extensions
@@ -208,15 +210,17 @@ class ByteParser extends Writable {
         } else {
           if (!this.#info.compressed) {
             this.#fragments.push(body)
+            this.#fragmentsBytes += body.length
 
             // If the frame is not fragmented, a message has been received.
             // If the frame is fragmented, it will terminate with a fin bit set
             // and an opcode of 0 (continuation), therefore we handle that when
             // parsing continuation frames, not here.
             if (!this.#info.fragmented && this.#info.fin) {
-              const fullMessage = Buffer.concat(this.#fragments)
+              const fullMessage = bufferConcat(this.#fragments, this.#fragmentsBytes)
               websocketMessageReceived(this.ws, this.#info.binaryType, fullMessage)
               this.#fragments.length = 0
+              this.#fragmentsBytes = 0
             }
 
             this.#state = parserStates.INFO
@@ -236,12 +240,13 @@ class ByteParser extends Writable {
                 return
               }
 
-              websocketMessageReceived(this.ws, this.#info.binaryType, Buffer.concat(this.#fragments))
+              websocketMessageReceived(this.ws, this.#info.binaryType, bufferConcat(this.#fragments, this.#fragmentsBytes))
 
               this.#loop = true
               this.#state = parserStates.INFO
               this.run(callback)
               this.#fragments.length = 0
+              this.#fragmentsBytes = 0
             })
 
             this.#loop = false
diff --git a/lib/web/websocket/util.js b/lib/web/websocket/util.js
@@ -294,6 +294,29 @@ const utf8Decode = hasIntl
     throw new TypeError('Invalid utf-8 received.')
   }
 
+function bufferConcat (buffers, length) {
+  if (buffers.length === 0 || length === 0) {
+    return Buffer.allocUnsafe(0)
+  }
+  let offset = 0
+  if (length === undefined) {
+    length = 0
+    for (let i = 0; i < buffers.length; ++i) {
+      length += buffers[i].length
+    }
+  }
+  const output = Buffer.allocUnsafe(length)
+  for (let i = 0; i < buffers.length; ++i) {
+    const buffer = buffers[i]
+    output.set(buffer, offset)
+    offset += buffer.length
+  }
+  if (offset < length) {
+    output.fill(0, offset, length)
+  }
+  return output
+}
+
 module.exports = {
   isConnecting,
   isEstablished,
@@ -305,6 +328,7 @@ module.exports = {
   failWebsocketConnection,
   websocketMessageReceived,
   utf8Decode,
+  bufferConcat,
   isControlFrame,
   isContinuationFrame,
   isTextBinaryFrame,
