import { use } from "express/lib/application"; import { WebSocketServer } from "ws"; import { guildMap } from "./common.js"; import { decodeToken } from "./tokens.js"; const messageSchema = { t: "number", d: "object" }; const authenticationTimeoutMs = 15000; const messageTypes = { HELLO: 0, YOO: 1, READY: 2, EVENT: 3 }; class GatewayServer { constructor(server, extraWebsocketServerConfig={}) { this.wss = new WebSocketServer({ server, ...extraWebsocketServerConfig }); setInterval(this.onAliveCheck.bind(this), 30000); this.wss.on("connection", (ws) => { this.onConnection(ws); ws.on("message", (data, isBinary) => { this.onMessage(ws, data, isBinary); }); ws.on("close", (code, reason) => { this.onDisconnect(ws, code, reason.toString()); }); ws.on("pong", () => { this.onPong(ws); }); }); } _clientDispatch(ws, e) { switch (e.type) { case "INIT": { ws.state = ws.state || { authenticated: false, user: null, token: null, alive: true, handlers: new Map() }; break; } case "AUTHENTICATE_AS": { ws.state.user = e.user; ws.state.authenticated = true; break; } case "SET_ALIVE": { ws.state.alive = e.alive; break; } case "SET_TOKEN": { ws.state.token = e.token; break; } case "ADD_HANDLER": { ws.state.handlers.set(e.guildId, e.handler); break; } } } _checkMessageSchema(message) { for (const [key, value] of Object.entries(message)) { if (!messageSchema[key]) return false; if (typeof value !== messageSchema[key]) return false; } return true; } onConnection(ws) { this._clientDispatch(ws, { type: "INIT" }); setTimeout(() => { if (!ws.state.authenticated) { ws.close(4001, "Authentication timeout."); } }, authenticationTimeoutMs); ws.send(JSON.stringify({ t: messageTypes.HELLO, d: null })); } async onMessage(ws, message, isBinary) { if (isBinary) return ws.close(4000, "Binary payload not supported."); try { message = JSON.parse(message.toString()); } catch (e) { console.error("GatewayServer: payload decode error", e); return ws.close(4000, "Payload error."); } if (!this._checkMessageSchema(message)) return ws.close(4000, "JSON payload does not match schema."); switch (message.t) { case messageTypes.YOO: { if (message.d.token) { let user; try { user = await decodeToken(message.d.token); } catch(e) { console.error(e); ws.close(4001, "Bad token."); break; } if (user && user.username) { if (!user.guildAccess || user.guildAccess.length < 1) { ws.close(4002, "No possible events: no guilds."); break; } this._clientDispatch(ws, { type: "SET_TOKEN", token: message.token }); this._clientDispatch(ws, { type: "AUTHENTICATE_AS", user }); // TODO: it might actually be more efficient to have a single listener // for each guild and broadcast the relevant events that way user.guildAccess.forEach((guildId) => { const guild = guildMap.get(guildId); if (!guild) { ws.close(4003, "User is in a guild that does not exist."); return; } const handle = (ev) => { ws.send(JSON.stringify({ t: messageTypes.EVENT, d: ev })); }; guild.on("pushed", handle); this._clientDispatch(ws, { type: "ADD_HANDLER", handler: handle, guildId }); }); ws.send(JSON.stringify({ t: messageTypes.READY, d: { user: { username: user.username, guildAccess: user.guildAccess, discordID: user.discordID, avatarURL: user.avatarURL, isSuperToken: user.isSuperToken } } })); } else { ws.close(4001, "Bad token."); break; } } else { ws.close(4001, "No token."); break; } break; } default: { return ws.close(4000, "Invalid payload type."); } } } onDisconnect(ws, code, reason) { if (ws.state && ws.state.handlers) { for (const [guildId, handler] of ws.state.handlers.entries()) { const guild = guildMap.get(guildId); if (guild) { guild.removeListener("pushed", handler); } } } } onPong(ws) { this._clientDispatch(ws, { type: "SET_ALIVE", alive: true }); } onAliveCheck() { this.wss.clients.forEach(ws => { if (ws.isAlive === false) return ws.terminate(); this._clientDispatch(ws, { type: "SET_ALIVE", alive: false }); ws.ping(); }); } } export default GatewayServer;