diff --git a/src/WebSocket.js b/src/WebSocket.js new file mode 100644 index 000000000..009a0f798 --- /dev/null +++ b/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]; diff --git a/src/client/ClientManager.js b/src/client/ClientManager.js index 3ffc82e7c..0a6ec95af 100644 --- a/src/client/ClientManager.js +++ b/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 => { diff --git a/src/client/voice/VoiceWebSocket.js b/src/client/voice/VoiceWebSocket.js index c5927d57a..9db347592 100644 --- a/src/client/voice/VoiceWebSocket.js +++ b/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); } diff --git a/src/client/websocket/WebSocketConnection.js b/src/client/websocket/WebSocketConnection.js index 42fda40e8..423975ffd 100644 --- a/src/client/websocket/WebSocketConnection.js +++ b/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; diff --git a/src/index.js b/src/index.js index ed966e000..d3928d6bb 100644 --- a/src/index.js +++ b/src/index.js @@ -57,4 +57,6 @@ module.exports = { User: require('./structures/User'), VoiceChannel: require('./structures/VoiceChannel'), Webhook: require('./structures/Webhook'), + + WebSocket: require('./WebSocket'), };