Update WebSocketShard.js

This commit is contained in:
Elysia 2024-01-13 12:15:42 +07:00
parent 0c42a5b409
commit 73734f310e

View File

@ -1,11 +1,11 @@
'use strict'; 'use strict';
const EventEmitter = require('node:events'); const EventEmitter = require('node:events');
const http = require('node:http');
const { setTimeout, setInterval, clearTimeout } = require('node:timers'); const { setTimeout, setInterval, clearTimeout } = require('node:timers');
const WebSocket = require('../../WebSocket'); const WebSocket = require('../../WebSocket');
const { Status, Events, ShardEvents, Opcodes, WSEvents, WSCodes } = require('../../util/Constants'); const { Status, Events, ShardEvents, Opcodes, WSEvents, WSCodes } = require('../../util/Constants');
const Intents = require('../../util/Intents'); const Intents = require('../../util/Intents');
const { getProxyObject } = require('../../util/Util');
const STATUS_KEYS = Object.keys(Status); const STATUS_KEYS = Object.keys(Status);
const CONNECTION_STATE = Object.keys(WebSocket.WebSocket); const CONNECTION_STATE = Object.keys(WebSocket.WebSocket);
@ -18,6 +18,7 @@ try {
/** /**
* Represents a Shard's WebSocket connection * Represents a Shard's WebSocket connection
* @extends {EventEmitter}
*/ */
class WebSocketShard extends EventEmitter { class WebSocketShard extends EventEmitter {
constructor(manager, id) { constructor(manager, id) {
@ -35,6 +36,13 @@ class WebSocketShard extends EventEmitter {
*/ */
this.id = id; this.id = id;
/**
* The resume URL for this shard
* @type {?string}
* @private
*/
this.resumeURL = null;
/** /**
* The current status of the shard * The current status of the shard
* @type {Status} * @type {Status}
@ -62,13 +70,6 @@ class WebSocketShard extends EventEmitter {
*/ */
this.sessionId = null; this.sessionId = null;
/**
* URL to use when resuming
* @type {?string}
* @private
*/
this.resumeURL = null;
/** /**
* The previous heartbeat ping of the shard * The previous heartbeat ping of the shard
* @type {number} * @type {number}
@ -200,12 +201,12 @@ class WebSocketShard extends EventEmitter {
connect() { connect() {
const { client } = this.manager; const { client } = this.manager;
const gateway = this.resumeURL ?? this.manager.gateway;
if (this.connection?.readyState === WebSocket.OPEN && this.status === Status.READY) { if (this.connection?.readyState === WebSocket.OPEN && this.status === Status.READY) {
return Promise.resolve(); return Promise.resolve();
} }
const gateway = this.resumeURL ?? this.manager.gateway;
return new Promise((resolve, reject) => { return new Promise((resolve, reject) => {
const cleanup = () => { const cleanup = () => {
this.removeListener(ShardEvents.CLOSE, onClose); this.removeListener(ShardEvents.CLOSE, onClose);
@ -271,7 +272,7 @@ class WebSocketShard extends EventEmitter {
Version : ${client.options.ws.version} Version : ${client.options.ws.version}
Encoding : ${WebSocket.encoding} Encoding : ${WebSocket.encoding}
Compression: ${zlib ? 'zlib-stream' : 'none'} Compression: ${zlib ? 'zlib-stream' : 'none'}
Proxy : ${client.options.proxy || 'none'}`, Agent : ${client.options.ws.agent instanceof http.Agent}`,
); );
this.status = this.status === Status.DISCONNECTED ? Status.RECONNECTING : Status.CONNECTING; this.status = this.status === Status.DISCONNECTED ? Status.RECONNECTING : Status.CONNECTING;
@ -279,13 +280,11 @@ class WebSocketShard extends EventEmitter {
this.setWsCloseTimeout(-1); this.setWsCloseTimeout(-1);
this.connectedAt = Date.now(); this.connectedAt = Date.now();
let args = { handshakeTimeout: 30_000 };
if (client.options.proxy.length > 0) {
args.agent = getProxyObject(client.options.proxy);
this.debug(`Using proxy ${client.options.proxy}`, args);
}
// Adding a handshake timeout to just make sure no zombie connection appears. // Adding a handshake timeout to just make sure no zombie connection appears.
const ws = (this.connection = WebSocket.create(gateway, wsQuery, args)); const ws = (this.connection = WebSocket.create(gateway, wsQuery, {
handshakeTimeout: 30_000,
agent: client.options.ws.agent instanceof http.Agent ? client.options.ws.agent : undefined,
}));
ws.onopen = this.onOpen.bind(this); ws.onopen = this.onOpen.bind(this);
ws.onmessage = this.onMessage.bind(this); ws.onmessage = this.onMessage.bind(this);
ws.onerror = this.onError.bind(this); ws.onerror = this.onError.bind(this);
@ -433,11 +432,11 @@ class WebSocketShard extends EventEmitter {
*/ */
this.emit(ShardEvents.READY); this.emit(ShardEvents.READY);
this.sessionId = packet.d.session_id;
this.resumeURL = packet.d.resume_gateway_url; this.resumeURL = packet.d.resume_gateway_url;
this.sessionId = packet.d.session_id;
this.expectedGuilds = new Set(packet.d.guilds.filter(d => d?.unavailable == true).map(d => d.id)); this.expectedGuilds = new Set(packet.d.guilds.filter(d => d?.unavailable == true).map(d => d.id));
this.status = Status.WAITING_FOR_GUILDS; this.status = Status.WAITING_FOR_GUILDS;
this.debug(`[READY] Session ${this.sessionId} | ResumeURL ${this.resumeURL}`); this.debug(`[READY] Session ${this.sessionId} | Resume url ${this.resumeURL}.`);
this.lastHeartbeatAcked = true; this.lastHeartbeatAcked = true;
this.sendHeartbeat('ReadyHeartbeat'); this.sendHeartbeat('ReadyHeartbeat');
break; break;
@ -538,21 +537,23 @@ class WebSocketShard extends EventEmitter {
const { waitGuildTimeout } = this.manager.client.options; const { waitGuildTimeout } = this.manager.client.options;
this.readyTimeout = setTimeout(() => { this.readyTimeout = setTimeout(
this.debug( () => {
`Shard ${hasGuildsIntent ? 'did' : 'will'} not receive any more guild packets` + this.debug(
`${hasGuildsIntent ? ` in ${waitGuildTimeout} ms` : ''}.\nUnavailable guild count: ${ `Shard ${hasGuildsIntent ? 'did' : 'will'} not receive any more guild packets` +
this.expectedGuilds.size `${hasGuildsIntent ? ` in ${waitGuildTimeout} ms` : ''}.\nUnavailable guild count: ${
}`, this.expectedGuilds.size
); }`,
);
this.readyTimeout = null; this.readyTimeout = null;
this.status = Status.READY; this.status = Status.READY;
this.emit(ShardEvents.ALL_READY, this.expectedGuilds); this.emit(ShardEvents.ALL_READY, this.expectedGuilds);
// }, hasGuildsIntent && waitGuildTimeout).unref(); },
}, waitGuildTimeout).unref(); hasGuildsIntent ? waitGuildTimeout : 0,
).unref();
} }
/** /**
@ -593,6 +594,7 @@ class WebSocketShard extends EventEmitter {
} }
this.wsCloseTimeout = setTimeout(() => { this.wsCloseTimeout = setTimeout(() => {
this.setWsCloseTimeout(-1); this.setWsCloseTimeout(-1);
// Check if close event was emitted. // Check if close event was emitted.
if (this.closeEmitted) { if (this.closeEmitted) {
this.debug(`[WebSocket] close was already emitted, assuming the connection was closed properly.`); this.debug(`[WebSocket] close was already emitted, assuming the connection was closed properly.`);
@ -656,7 +658,6 @@ class WebSocketShard extends EventEmitter {
Sequence : ${this.sequence} Sequence : ${this.sequence}
Connection State: ${this.connection ? CONNECTION_STATE[this.connection.readyState] : 'No Connection??'}`, Connection State: ${this.connection ? CONNECTION_STATE[this.connection.readyState] : 'No Connection??'}`,
); );
this.destroy({ reset: true, closeCode: 4009 }); this.destroy({ reset: true, closeCode: 4009 });
return; return;
} }
@ -700,25 +701,24 @@ class WebSocketShard extends EventEmitter {
this.status = Status.IDENTIFYING; this.status = Status.IDENTIFYING;
// Patch something
Object.keys(client.options.ws.properties) Object.keys(client.options.ws.properties)
.filter(k => k.startsWith('$')) .filter(k => k.startsWith('$'))
.forEach(k => { .forEach(k => {
client.options.ws.properties[k.slice(1)] = client.options.ws.properties[k]; client.options.ws.properties[k.slice(1)] = client.options.ws.properties[k];
delete client.options.ws.properties[k]; delete client.options.ws.properties[k];
}); });
// Clone the identify payload and assign the token and shard info
const d = { const d = {
presence: client.options.presence,
...client.options.ws, ...client.options.ws,
token: client.token, token: client.token,
}; };
delete d.version; delete d.version;
delete d.agent;
this.debug( this.debug(`[IDENTIFY] Shard ${this.id}/${client.options.shardCount} with intents: ${d.intents}`);
`[IDENTIFY] Shard ${this.id}/${client.options.shardCount} with intents: ${Intents.resolve(
client.options.intents,
)} 😊`,
);
this.send({ op: Opcodes.IDENTIFY, d }, true); this.send({ op: Opcodes.IDENTIFY, d }, true);
} }
@ -815,6 +815,7 @@ class WebSocketShard extends EventEmitter {
// Step 0: Remove all timers // Step 0: Remove all timers
this.setHeartbeatTimer(-1); this.setHeartbeatTimer(-1);
this.setHelloTimeout(-1); this.setHelloTimeout(-1);
this.debug( this.debug(
`[WebSocket] Destroy: Attempting to close the WebSocket. | WS State: ${ `[WebSocket] Destroy: Attempting to close the WebSocket. | WS State: ${
CONNECTION_STATE[this.connection?.readyState ?? WebSocket.CLOSED] CONNECTION_STATE[this.connection?.readyState ?? WebSocket.CLOSED]