[feature] Add "fragments" as possible value for binaryType (#1018)

Author: guymguym

.gitignore

@@ -1,3 +1,4 @@
 node_modules/
 coverage/
 npm-debug.log
+.vscode/

doc/ws.md

@@ -277,7 +277,10 @@ Register an event listener emulating the `EventTarget` interface.
 - {String}
 
 A string indicating the type of binary data being transmitted by the connection.
-This should be either "nodebuffer" or "arraybuffer". Defaults to "nodebuffer".
+This should be either "nodebuffer" or "arraybuffer" or "fragments". Defaults to "nodebuffer".
+Type "fragments" will emit the array of fragments as received from the sender,
+without copyfull concatenation, which is useful for the performance of binary protocols 
+transfering large messages with multiple fragments.
 
 ### websocket.bufferedAmount
 

lib/Constants.js

@@ -3,3 +3,5 @@
 exports.GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11';
 exports.EMPTY_BUFFER = Buffer.alloc(0);
 exports.NOOP = () => {};
+
+exports.BINARY_TYPES = ['nodebuffer', 'arraybuffer', 'fragments'];

lib/EventTarget.js

@@ -28,7 +28,7 @@ class MessageEvent extends Event {
   /**
    * Create a new `MessageEvent`.
    *
-   * @param {(String|Buffer|ArrayBuffer)} data The received data
+   * @param {(String|Buffer|ArrayBuffer|Buffer[])} data The received data
    * @param {Boolean} isBinary Specifies if `data` is binary
    * @param {WebSocket} target A reference to the target to which the event was dispatched
    */
@@ -100,9 +100,6 @@ const EventTarget = {
     if (typeof listener !== 'function') return;
 
     function onMessage (data, flags) {
-      if (flags.binary && this.binaryType === 'arraybuffer') {
-        data = new Uint8Array(data).buffer;
-      }
       listener.call(this, new MessageEvent(data, !!flags.binary, this));
     }
 

lib/Receiver.js

@@ -28,10 +28,12 @@ class Receiver {
    *
    * @param {Object} extensions An object containing the negotiated extensions
    * @param {Number} maxPayload The maximum allowed message length
+   * @param {String} binaryType The type for binary data, see constants.BINARY_TYPES
    */
-  constructor (extensions, maxPayload) {
+  constructor (extensions, maxPayload, binaryType) {
     this.extensions = extensions || {};
     this.maxPayload = maxPayload | 0;
+    this.binaryType = binaryType || constants.BINARY_TYPES[0];
 
     this.bufferedBytes = 0;
     this.buffers = [];
@@ -352,20 +354,28 @@ class Receiver {
    */
   dataMessage () {
     if (this.fin) {
-      const buf = this.fragments.length > 1
-        ? Buffer.concat(this.fragments, this.messageLength)
-        : this.fragments.length === 1
-          ? this.fragments[0]
-          : constants.EMPTY_BUFFER;
-
+      const fragments = this.fragments;
+      const messageLength = this.messageLength;
       this.totalPayloadLength = 0;
-      this.fragments.length = 0;
+      this.fragments = [];
       this.messageLength = 0;
       this.fragmented = 0;
 
       if (this.opcode === 2) {
-        this.onmessage(buf, { masked: this.masked, binary: true });
+        var data;
+
+        if (this.binaryType === 'nodebuffer') {
+          data = bufferFromFragments(fragments, messageLength);
+        } else if (this.binaryType === 'arraybuffer') {
+          data = new Uint8Array(bufferFromFragments(fragments, messageLength)).buffer;
+        } else {
+          data = fragments;
+        }
+
+        this.onmessage(data, { masked: this.masked, binary: true });
       } else {
+        const buf = bufferFromFragments(fragments, messageLength);
+
         if (!isValidUTF8(buf)) {
           this.error(new Error('invalid utf8 sequence'), 1007);
           return;
@@ -509,4 +519,20 @@ class Receiver {
   }
 }
 
+/**
+ * Make a buffer from a list of fragments.
+ * Optimized for the common case of a single fragment in order to
+ * avoid copyfull concat and simply return the fragment buffer.
+ *
+ * @param {Buffer[]} fragments
+ * @param {Number} messageLength
+ * @return {Buffer}
+ * @private
+ */
+function bufferFromFragments (fragments, messageLength) {
+  if (fragments.length === 1) return fragments[0];
+  if (fragments.length > 1) return Buffer.concat(fragments, messageLength);
+  return constants.EMPTY_BUFFER;
+}
+
 module.exports = Receiver;

lib/WebSocket.js

@@ -53,7 +53,7 @@ class WebSocket extends EventEmitter {
     this.protocol = '';
 
     this._finalize = this.finalize.bind(this);
-    this._binaryType = 'nodebuffer';
+    this._binaryType = constants.BINARY_TYPES[0];
     this._finalizeCalled = false;
     this._closeMessage = null;
     this._closeTimer = null;
@@ -98,10 +98,17 @@ class WebSocket extends EventEmitter {
   }
 
   set binaryType (type) {
-    if (type === 'arraybuffer' || type === 'nodebuffer') {
-      this._binaryType = type;
-    } else {
-      throw new SyntaxError('unsupported binaryType: must be either "nodebuffer" or "arraybuffer"');
+    // silently ignore unsupported types
+    if (constants.BINARY_TYPES.indexOf(type) < 0) {
+      return;
+    }
+
+    this._binaryType = type;
+
+    // update the receiver if already created,
+    // if not then it will take the value once created
+    if (this._receiver) {
+      this._receiver.binaryType = type;
     }
   }
 
@@ -116,7 +123,7 @@ class WebSocket extends EventEmitter {
     socket.setTimeout(0);
     socket.setNoDelay();
 
-    this._receiver = new Receiver(this.extensions, this.maxPayload);
+    this._receiver = new Receiver(this.extensions, this.maxPayload, this.binaryType);
     this._sender = new Sender(socket, this.extensions);
     this._ultron = new Ultron(socket);
     this._socket = socket;

test/Receiver.test.js

@@ -5,6 +5,7 @@ const crypto = require('crypto');
 
 const PerMessageDeflate = require('../lib/PerMessageDeflate');
 const Receiver = require('../lib/Receiver');
+const Sender = require('../lib/Sender');
 const util = require('./hybi-util');
 
 describe('Receiver', function () {
@@ -827,4 +828,85 @@ describe('Receiver', function () {
       });
     });
   });
+
+  it('can emit nodebuffer of fragmented binary message', function (done) {
+    const p = new Receiver();
+    const frags = [
+      crypto.randomBytes(7321),
+      crypto.randomBytes(137),
+      crypto.randomBytes(285787),
+      crypto.randomBytes(3)
+    ];
+
+    p.binaryType = 'nodebuffer';
+    p.onmessage = function (data) {
+      assert.ok(Buffer.isBuffer(data));
+      assert.ok(data.equals(Buffer.concat(frags)));
+      done();
+    };
+
+    addBinaryFragments(p, frags);
+  });
+
+  it('can emit arraybuffer of fragmented binary message', function (done) {
+    const p = new Receiver();
+    const frags = [
+      crypto.randomBytes(19221),
+      crypto.randomBytes(954),
+      crypto.randomBytes(623987)
+    ];
+
+    p.binaryType = 'arraybuffer';
+    p.onmessage = function (data) {
+      assert.ok(data instanceof ArrayBuffer);
+      assert.ok(Buffer.from(data).equals(Buffer.concat(frags)));
+      done();
+    };
+
+    addBinaryFragments(p, frags);
+  });
+
+  it('can emit fragments of fragmented binary message', function (done) {
+    const p = new Receiver();
+    const frags = [
+      crypto.randomBytes(17),
+      crypto.randomBytes(419872),
+      crypto.randomBytes(83),
+      crypto.randomBytes(9928),
+      crypto.randomBytes(1)
+    ];
+
+    p.binaryType = 'fragments';
+    p.onmessage = function (data) {
+      assert.ok(Array.isArray(data));
+      assert.ok(data.length === frags.length);
+      for (let i = 0; i < frags.length; ++i) {
+        assert.ok(Buffer.isBuffer(data[i]));
+        assert.ok(frags[i].equals(data[i]));
+      }
+      done();
+    };
+
+    addBinaryFragments(p, frags);
+  });
 });
+
+/**
+ * Adds a list of binary fragments to the receiver prefixing each one
+ * with info byte and length.
+ *
+ * @param {Receiver} receiver
+ * @param {Buffer[]} frags
+ * @private
+ */
+function addBinaryFragments (receiver, frags) {
+  for (let i = 0; i < frags.length; ++i) {
+    Sender.frame(frags[i], {
+      opcode: i === 0 ? 2 : 0,
+      fin: i + 1 === frags.length,
+      mask: false,
+      rsv1: false,
+      readOnly: true
+    }).forEach(buf => receiver.add(buf));
+  }
+}

test/WebSocket.test.js

@@ -1168,12 +1168,23 @@ describe('WebSocket', function () {
       assert.strictEqual(ws.onopen, listener);
     });
 
-    it('should throw an error when setting an invalid binary type', function () {
+    it('should ignore when setting an invalid binary type', function () {
       const ws = new WebSocket('ws://localhost', { agent: new CustomAgent() });
 
-      assert.throws(() => {
-        ws.binaryType = 'foo';
-      }, /^SyntaxError: unsupported binaryType: must be either "nodebuffer" or "arraybuffer"$/);
+      ws.binaryType = 'nodebuffer';
+      assert.ok(ws.binaryType === 'nodebuffer');
+      ws.binaryType = 'foo';
+      assert.ok(ws.binaryType === 'nodebuffer');
+      ws.binaryType = 'arraybuffer';
+      assert.ok(ws.binaryType === 'arraybuffer');
+      ws.binaryType = '';
+      assert.ok(ws.binaryType === 'arraybuffer');
+      ws.binaryType = 'fragments';
+      assert.ok(ws.binaryType === 'fragments');
+      ws.binaryType = 'buffer';
+      assert.ok(ws.binaryType === 'fragments');
+      ws.binaryType = 'nodebuffer';
+      assert.ok(ws.binaryType === 'nodebuffer');
     });
 
     it('should work the same as the EventEmitter api', function (done) {
@@ -1420,6 +1431,44 @@ describe('WebSocket', function () {
         };
       });
     });
+
+    it('should allow to update binaryType after receiver created', function (done) {
+      server.createServer(++port, (srv) => {
+        const ws = new WebSocket(`ws://localhost:${port}`);
+
+        function testType (binaryType, callback) {
+          const buf = Buffer.from(binaryType);
+          ws.binaryType = binaryType;
+          ws.onmessage = (messageEvent) => {
+            if (binaryType === 'nodebuffer') {
+              assert.ok(Buffer.isBuffer(messageEvent.data));
+              assert.ok(messageEvent.data.equals(buf));
+            } else if (binaryType === 'arraybuffer') {
+              assert.ok(messageEvent.data instanceof ArrayBuffer);
+              assert.ok(Buffer.from(messageEvent.data).equals(buf));
+            } else if (binaryType === 'fragments') {
+              assert.ok(Array.isArray(messageEvent.data));
+              assert.ok(messageEvent.data.length === 1);
+              assert.ok(Buffer.from(messageEvent.data[0]).equals(buf));
+            }
+            callback();
+          };
+          ws.send(buf);
+        }
+
+        ws.onopen =
+          () => testType('nodebuffer',
+            () => testType('arraybuffer',
+              () => testType('fragments',
+                () => {
+                  srv.close(done);
+                  ws.terminate();
+                }
+              )
+            )
+          );
+      });
+    });
   });
 
   describe('ssl', function () {