perf(websocket): mask frames without a mask array, four bytes per step

lib/web/websocket/frame.js

@@ -12,12 +12,58 @@ const randomFillSync = runtimeFeatures.has('crypto')
   ? require('node:crypto').randomFillSync
   : null
 
-function generateMask () {
+/**
+ * Writes the next mask key into target at offset.
+ * @param {Uint8Array|number[]} target
+ * @param {number} offset
+ */
+function writeMask (target, offset) {
   if (bufIdx === BUFFER_SIZE) {
     bufIdx = 0
     randomFillSync((buffer ??= Buffer.allocUnsafeSlow(BUFFER_SIZE)), 0, BUFFER_SIZE)
   }
-  return [buffer[bufIdx++], buffer[bufIdx++], buffer[bufIdx++], buffer[bufIdx++]]
+  target[offset] = buffer[bufIdx++]
+  target[offset + 1] = buffer[bufIdx++]
+  target[offset + 2] = buffer[bufIdx++]
+  target[offset + 3] = buffer[bufIdx++]
+}
+
+function generateMask () {
+  const mask = [0, 0, 0, 0]
+  writeMask(mask, 0)
+  return mask
+}
+
+/**
+ * Writes source[0, length) XOR the mask key at mask[maskOffset, maskOffset + 4)
+ * into target, starting at targetOffset. source and target may be the same
+ * buffer when targetOffset is 0.
+ * @param {Uint8Array} source
+ * @param {Uint8Array} target
+ * @param {number} targetOffset
+ * @param {Uint8Array} mask
+ * @param {number} maskOffset
+ * @param {number} length
+ */
+function maskPayload (source, target, targetOffset, mask, maskOffset, length) {
+  const mask0 = mask[maskOffset]
+  const mask1 = mask[maskOffset + 1]
+  const mask2 = mask[maskOffset + 2]
+  const mask3 = mask[maskOffset + 3]
+  const end = length - (length & 3)
+
+  // Four bytes per step, then the remaining zero to three bytes.
+  for (let i = 0; i < end; i += 4) {
+    target[targetOffset + i] = source[i] ^ mask0
+    target[targetOffset + i + 1] = source[i + 1] ^ mask1
+    target[targetOffset + i + 2] = source[i + 2] ^ mask2
+    target[targetOffset + i + 3] = source[i + 3] ^ mask3
+  }
+
+  const rest = length - end
+  if (rest > 0) target[targetOffset + end] = source[end] ^ mask0
+  if (rest > 1) target[targetOffset + end + 1] = source[end + 1] ^ mask1
+  if (rest > 2) target[targetOffset + end + 2] = source[end + 2] ^ mask2
 }
 
 class WebsocketFrameSend {
@@ -30,7 +76,6 @@ class WebsocketFrameSend {
 
   createFrame (opcode) {
     const frameData = this.frameData
-    const maskKey = generateMask()
     const bodyLength = frameData?.byteLength ?? 0
 
     /** @type {number} */
@@ -53,10 +98,7 @@ class WebsocketFrameSend {
     buffer[0] = (buffer[0] & 0xF0) + opcode // opcode
 
     /*! ws. MIT License. Einar Otto Stangvik <einaros@gmail.com> */
-    buffer[offset - 4] = maskKey[0]
-    buffer[offset - 3] = maskKey[1]
-    buffer[offset - 2] = maskKey[2]
-    buffer[offset - 1] = maskKey[3]
+    writeMask(buffer, offset - 4)
 
     buffer[1] = payloadLength
 
@@ -71,9 +113,7 @@ class WebsocketFrameSend {
     buffer[1] |= 0x80 // MASK
 
     // mask body
-    for (let i = 0; i < bodyLength; ++i) {
-      buffer[offset + i] = frameData[i] ^ maskKey[i & 3]
-    }
+    maskPayload(frameData, buffer, offset, buffer, offset - 4, bodyLength)
 
     return buffer
   }
@@ -82,15 +122,8 @@ class WebsocketFrameSend {
    * @param {Uint8Array} buffer
    */
   static createFastTextFrame (buffer) {
-    const maskKey = generateMask()
-
     const bodyLength = buffer.length
 
-    // mask body
-    for (let i = 0; i < bodyLength; ++i) {
-      buffer[i] ^= maskKey[i & 3]
-    }
-
     let payloadLength = bodyLength
     let offset = 6
 
@@ -102,13 +135,15 @@ class WebsocketFrameSend {
       payloadLength = 126
     }
     const head = Buffer.allocUnsafeSlow(offset)
+    const maskOffset = offset - 4
+
+    writeMask(head, maskOffset)
+
+    // mask body
+    maskPayload(buffer, buffer, 0, head, maskOffset, bodyLength)
 
     head[0] = 0x80 /* FIN */ | opcodes.TEXT /* opcode TEXT */
     head[1] = payloadLength | 0x80 /* MASK */
-    head[offset - 4] = maskKey[0]
-    head[offset - 3] = maskKey[1]
-    head[offset - 2] = maskKey[2]
-    head[offset - 1] = maskKey[3]
 
     if (payloadLength === 126) {
       head.writeUInt16BE(bodyLength, 2)

test/websocket/frame.js

@@ -37,3 +37,30 @@ test('Writing 16-bit frame length value at correct offset when buffer has a non-
   t.assert.strictEqual(frame[3], payloadLength & 0xff)
   t.assert.strictEqual(smallBuffer.length, 1) // ensure smallBuffer can't be garbage-collected too soon
 })
+
+test('Masked payloads unmask to the original bytes for every tail length', (t) => {
+  const lengths = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 125, 126, 65535, 65536]
+
+  const payloadOffset = (head) => {
+    const length = head[1] & 0x7f
+    return length === 127 ? 14 : length === 126 ? 8 : 6
+  }
+
+  const unmask = (body, mask) => Buffer.from(body.map((byte, i) => byte ^ mask[i & 3]))
+
+  for (const length of lengths) {
+    const payload = Buffer.alloc(length)
+    for (let i = 0; i < length; i++) payload[i] = (i * 31 + 7) & 0xff
+
+    const frame = new WebsocketFrameSend(payload).createFrame(opcodes.BINARY)
+    const offset = payloadOffset(frame)
+    t.assert.strictEqual(frame.length, offset + length)
+    t.assert.deepStrictEqual(unmask(frame.subarray(offset), frame.subarray(offset - 4, offset)), payload)
+
+    const [head, body] = WebsocketFrameSend.createFastTextFrame(Buffer.from(payload))
+    const headOffset = payloadOffset(head)
+    t.assert.strictEqual(head.length, headOffset)
+    t.assert.strictEqual(body.length, length)
+    t.assert.deepStrictEqual(unmask(body, head.subarray(headOffset - 4, headOffset)), payload)
+  }
+})