Separate websocket internals for RPC (#1893)

* websocket centralization

* more centralization

* whoops

* Update WebSocket.js

src/WebSocket.js

@@ -0,0 +1,40 @@
+const browser = typeof window !== 'undefined';
+const zlib = require('zlib');
+const querystring = require('querystring');
+
+if (browser) {
+  exports.WebSocket = window.WebSocket; // eslint-disable-line no-undef
+} else {
+  try {
+    exports.WebSocket = require('uws');
+  } catch (err) {
+    exports.WebSocket = require('ws');
+  }
+}
+
+try {
+  var erlpack = require('erlpack');
+  if (!erlpack.pack) erlpack = null;
+} catch (err) {} // eslint-disable-line no-empty
+
+exports.encoding = erlpack ? 'etf' : 'json';
+
+exports.pack = erlpack ? erlpack.pack : JSON.stringify;
+
+exports.unpack = data => {
+  if (Array.isArray(data)) data = Buffer.concat(data);
+  if (data instanceof ArrayBuffer) data = Buffer.from(new Uint8Array(data));
+
+  if (erlpack && typeof data !== 'string') return erlpack.unpack(data);
+  else if (data instanceof Buffer) data = zlib.inflateSync(data).toString();
+  return JSON.parse(data);
+};
+
+exports.create = (gateway, query = {}, ...args) => {
+  query.encoding = exports.encoding;
+  const ws = new exports.WebSocket(`${gateway}?${querystring.stringify(query)}`, ...args);
+  if (browser) ws.binaryType = 'arraybuffer';
+  return ws;
+};
+
+for (const state of ['CONNECTING', 'OPEN', 'CLOSING', 'CLOSED']) exports[state] = exports.WebSocket[state];

src/client/ClientManager.js

@@ -1,5 +1,4 @@
 const Constants = require('../util/Constants');
-const WebSocketConnection = require('./websocket/WebSocketConnection');
 const { Error } = require('../errors');
 
 /**
@@ -40,8 +39,7 @@ class ClientManager {
     this.client.token = token;
     const timeout = this.client.setTimeout(() => reject(new Error('TOKEN_INVALID')), 1000 * 300);
     this.client.api.gateway.get().then(res => {
-      const protocolVersion = Constants.DefaultOptions.ws.version;
-      const gateway = `${res.url}/?v=${protocolVersion}&encoding=${WebSocketConnection.ENCODING}`;
+      const gateway = `${res.url}/`;
       this.client.emit(Constants.Events.DEBUG, `Using gateway ${gateway}`);
       this.client.ws.connect(gateway);
       this.client.ws.connection.once('close', event => {

src/client/voice/VoiceWebSocket.js

@@ -2,13 +2,7 @@ const Constants = require('../../util/Constants');
 const SecretKey = require('./util/SecretKey');
 const EventEmitter = require('events');
 const { Error } = require('../../errors');
-
-let WebSocket;
-try {
-  WebSocket = require('uws');
-} catch (err) {
-  WebSocket = require('ws');
-}
+const WebSocket = require('../../WebSocket');
 
 /**
  * Represents a Voice Connection's WebSocket.
@@ -75,7 +69,7 @@ class VoiceWebSocket extends EventEmitter {
      * The actual WebSocket used to connect to the Voice WebSocket Server.
      * @type {WebSocket}
      */
-    this.ws = new WebSocket(`wss://${this.voiceConnection.authentication.endpoint}`);
+    this.ws = WebSocket.create(`wss://${this.voiceConnection.authentication.endpoint}`, { v: 3 });
     this.ws.onopen = this.onOpen.bind(this);
     this.ws.onmessage = this.onMessage.bind(this);
     this.ws.onclose = this.onClose.bind(this);
@@ -103,7 +97,7 @@ class VoiceWebSocket extends EventEmitter {
    */
   sendPacket(packet) {
     try {
-      packet = JSON.stringify(packet);
+      packet = WebSocket.pack(packet);
     } catch (error) {
       return Promise.reject(error);
     }
@@ -134,7 +128,7 @@ class VoiceWebSocket extends EventEmitter {
    */
   onMessage(event) {
     try {
-      return this.onPacket(JSON.parse(event.data));
+      return this.onPacket(WebSocket.unpack(event.data));
     } catch (error) {
       return this.onError(error);
     }

src/client/websocket/WebSocketConnection.js

@@ -1,26 +1,7 @@
-const browser = typeof window !== 'undefined';
 const EventEmitter = require('events');
 const Constants = require('../../util/Constants');
-const zlib = require('zlib');
 const PacketManager = require('./packets/WebSocketPacketManager');
-const erlpack = (function findErlpack() {
-  try {
-    const e = require('erlpack');
-    if (!e.pack) return null;
-    return e;
-  } catch (e) {
-    return null;
-  }
-}());
-
-const WebSocket = (function findWebSocket() {
-  if (browser) return window.WebSocket; // eslint-disable-line no-undef
-  try {
-    return require('uws');
-  } catch (e) {
-    return require('ws');
-  }
-}());
+const WebSocket = require('../../WebSocket');
 
 /**
  * Abstracts a WebSocket connection with decoding/encoding for the Discord gateway.
@@ -161,30 +142,6 @@ class WebSocketConnection extends EventEmitter {
     return this.manager.debug(`[connection] ${message}`);
   }
 
-  /**
-   * Attempts to serialise data from the WebSocket.
-   * @param {string|Object} data Data to unpack
-   * @returns {Object}
-   */
-  unpack(data) {
-    if (Array.isArray(data)) data = Buffer.concat(data);
-    if (data instanceof ArrayBuffer) data = Buffer.from(new Uint8Array(data));
-
-    if (erlpack && typeof data !== 'string') return erlpack.unpack(data);
-    else if (data instanceof Buffer) data = zlib.inflateSync(data).toString();
-
-    return JSON.parse(data);
-  }
-
-  /**
-   * Packs an object ready to be sent.
-   * @param {Object} data Data to pack
-   * @returns {string|Buffer}
-   */
-  pack(data) {
-    return erlpack ? erlpack.pack(data) : JSON.stringify(data);
-  }
-
   /**
    * Processes the current WebSocket queue.
    */
@@ -215,7 +172,7 @@ class WebSocketConnection extends EventEmitter {
       this.debug(`Tried to send packet ${data} but no WebSocket is available!`);
       return;
     }
-    this.ws.send(this.pack(data));
+    this.ws.send(WebSocket.pack(data));
   }
 
   /**
@@ -251,8 +208,7 @@ class WebSocketConnection extends EventEmitter {
     this.expectingClose = false;
     this.gateway = gateway;
     this.debug(`Connecting to ${gateway}`);
-    const ws = this.ws = new WebSocket(gateway);
-    if (browser) ws.binaryType = 'arraybuffer';
+    const ws = this.ws = WebSocket.create(gateway, { v: Constants.DefaultOptions.ws.version });
     ws.onmessage = this.onMessage.bind(this);
     ws.onopen = this.onOpen.bind(this);
     ws.onerror = this.onError.bind(this);
@@ -289,7 +245,7 @@ class WebSocketConnection extends EventEmitter {
   onMessage(event) {
     let data;
     try {
-      data = this.unpack(event.data);
+      data = WebSocket.unpack(event.data);
     } catch (err) {
       this.emit('debug', err);
     }
@@ -497,11 +453,4 @@ class WebSocketConnection extends EventEmitter {
   }
 }
 
-/**
- * Encoding the WebSocket connections will use.
- * @type {string}
- */
-WebSocketConnection.ENCODING = erlpack ? 'etf' : 'json';
-WebSocketConnection.WebSocket = WebSocket;
-
 module.exports = WebSocketConnection;

src/index.js

@@ -57,4 +57,6 @@ module.exports = {
   User: require('./structures/User'),
   VoiceChannel: require('./structures/VoiceChannel'),
   Webhook: require('./structures/Webhook'),
+
+  WebSocket: require('./WebSocket'),
 };