brainlet/brainlet/api/v2/gateway/index.js

const websockets = require("ws");
const EventEmitter = require("events");
const uuid = require("uuid");

const User = require("../../../models/User");
const Channel = require("../../../models/Channel");
const { parseMessage, opcodeSeparator, getOpcodeByName } = require("./messageparser");
const { checkToken } = require("../../../common/auth/authfunctions");

const pingCheckDelay = 10000;

class GatewayServer extends EventEmitter {
    constructor({ server }) {
        super();
        this.wss = new websockets.Server({ server: server, path: "/gateway" });

        this.pingInterval = setInterval(() => {
            this.wss.clients.forEach((client) => {
                if (!client.alive) {
                    client.terminate();
                }
                client.alive = false;
                client.ping(() => {});
            });
        }, pingCheckDelay);

        this.wss.on("close", () => {
            clearInterval(this.pingInterval);
            console.log("gateway: websocket server closed");
        });

        this.wss.on("connection", (ws) => {
            // Send HELLO message as soon as the client connects
            ws.send(this.packet("HELLO", {}));
            ws.session = {
                authenticated: false,
                user: null,
                sessionId: uuid.v4()
            };
            ws.alive = true;

            ws.on("pong", () => {
                ws.alive = true;
            });
            ws.on("message", async (data) => {
                try {
                    const message = parseMessage(data);
                    switch (message.opcodeType) {
                    case "YOO": {
                        // The client has responded to our HELLO with a YOO packet
                        try {
                            const user = await checkToken(message.data.token);
                            if (!user) return ws.close(4006, "Authentication failed.");
                            ws.session.user = user;
                            ws.session.authenticated = true;

                            // The user is now successfully authenticated, send the YOO_ACK packet
                            // TODO: This is probably not efficient

                            let channels = await Channel.find().lean().sort({ _id: -1 }).limit(50).select("-posts -__v").populate("creator", User.getPulicFields(true));
                            if (!channels) channels = [];
                            channels = channels.map(x => ({ ...x, _id: x._id.toString() }));

                            ws.channels = channels.map(x => x._id);

                            ws.send(this.packet("YOO_ACK", { session_id: ws.session.sessionId, channels, user: { username: user.username, _id: user._id } }));
                            console.log(`gateway: user ${user.username}: handshake complete`);
                        } catch (e) {
                            console.log("gateway:", e);
                            return ws.close(4006, "Authentication failed.");
                        }
                        break;
                    }
                    case "ACTION_CREATE_MESSAGE": {
                        if (!this.authMessage(ws)) return;
                        if (typeof message.data.content !== "string" || typeof message.data.channel !== "object" || typeof message.data.channel._id !== "string") throw new Error("msg: invalid fields in json payload");
                        
                        const messageContent = message.data.content.trim();
                        if (messageContent.length > 2000) return;
                        if (messageContent === "") return;
                        if (message.data.channel._id.length !== 24) throw new Error("msg: payload has invalid id"); // MONGODB ONLY!!

                        // Check if the user is in that channel before broadcasting the message
                        if (!ws.channels.includes(message.data.channel._id)) return ws.close(4008, "Not authorized to perform action.");

                        this.broadcast(message.data.channel._id, this.packet("EVENT_CREATE_MESSAGE", {
                            content: messageContent,
                            channel: {
                                _id: message.data.channel._id
                            },
                            author: {
                                _id: ws.session.user._id,
                                username: ws.session.user.username
                            },
                            _id: uuid.v4()
                        }));

                        break;
                    }
                    }
                } catch(e) {
                    console.error("gateway:", e);
                    return ws.close(4000, "Error while handling payload.");
                }
            });
        });
    }
}

GatewayServer.prototype.broadcast = function(channelId, data) {
    this.wss.clients.forEach((client) => {
        if (this.clientReady(client) && client.channels.includes(channelId)) client.send(data);
    });
};

GatewayServer.prototype.clientReady = function(ws) {
    return ws.readyState === websockets.OPEN && ws.session && ws.session.authenticated;
};

GatewayServer.prototype.authMessage = function(ws) {
    if (!this.clientReady(ws)) {
        ws.close(4007, "Not authenticated.");
        return false;
    }
    return true;
};

GatewayServer.prototype.packet = function(op, data) {
    if (typeof op === "string") op = getOpcodeByName(op);
    return `${op}${opcodeSeparator}${JSON.stringify(data)}`;
};

module.exports = GatewayServer;