bridgecord/GatewayServer.js

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;