brainlet/api/v1/gateway/index.js

const User = require('../../../models/User');
const secret = require('../../../secret');
const config = require('../../../config');
const Category = require('../../../models/Category');
const RateLimiter = require('./ratelimiter');

const jwt = require('jsonwebtoken');
const siolib = require('socket.io');
const uuid = require('uuid');

class GatewayServer {
    constructor(httpServer) {
        this._io = siolib(httpServer);
        this._gateway = this._io.of('/gateway');
        this.rateLimiter = new RateLimiter({
            points: 5,
            time: 1000,
            minPoints: 0
        });
        this.eventSetup();

        this._commandPrefix = '/';
    }
}

GatewayServer.prototype._sendSystemMessage = function(socket, message, category) {
    const messageObject = {
        author: {
            username: '__SYSTEM',
            _id: '5fc69864f15a7c5e504c9a1f'
        },
        category: {
            title: category.title,
            _id: category._id
        },
        content: message,
        _id: uuid.v4()
    };

    socket.emit('message', messageObject);
};

GatewayServer.prototype.notifyClientsOfUpdate = function(reason) {
    this._gateway.emit('refreshClient', { reason: reason || 'REFRESH' });
};

GatewayServer.prototype._processCommand = function(socket, message) {
    const content = message.content;
    const fullCommandString = content.slice(this._commandPrefix.length, content.length);
    const fullCommand = fullCommandString.split(' ');
    const command = fullCommand[0] || 'INVALID_COMMAND';
    const args = fullCommand.length - 1;

    switch (command) {
        case 'INVALID_COMMAND': {
            this._sendSystemMessage(socket, 'Invalid command.', message.category);
            break;
        }
        case 'admin/fr': {
            if (args === 1) {
                if (socket.user.permissionLevel >= config.roleMap.ADMIN) {
                    this._gateway.emit('refreshClient', { reason: fullCommand[1] || 'REFRESH' });
                } else {
                    this._sendSystemMessage(socket, 'how about no', message.category);
                }
            } else {
                this._sendSystemMessage(socket, 'Invalid number of arguments.', message.category);
            }
            break;
        }
        default: {
            this._sendSystemMessage(socket, 'That command does not exist.', message.category);
            break;
        }
    }
};

GatewayServer.prototype.authDisconnect = function(socket, callback) {
    console.log('[E] [gateway] [handshake] User disconnected due to failed authentication');
    socket.isConnected = false;
    socket.disconnect();
    socket.disconnect(true);
    callback(new Error('ERR_GATEWAY_AUTH_FAIL'));
};

GatewayServer.prototype.eventSetup = function() {
    this._gateway.use((socket, callback) => {
        console.log('[*] [gateway] [handshake] User authentication attempt');
        socket.isConnected = false;

        setTimeout(() => {
            if (socket.isConnected) return;
            console.log('[E] [gateway] [handshake] User still not connected after timeout, removing...');
            socket.disconnect();
            socket.disconnect(true);
        }, config.gatewayStillNotConnectedTimeoutMS);

        // TODO: Maybe passing the token in the query is not the best idea?
        const token = socket.handshake.query.token;

        if (!token) return this.authDisconnect(socket, callback);
        if (!(typeof token === 'string')) return this.authDisconnect(socket, callback);

        const allSockets = this._gateway.sockets;
        for (let [_, e] of allSockets) {
            if (e.user && e.user.token === token) {
                console.log(`[E] [gateway] [handshake] User ${e.user.username} tried to connect more than once, rejecting connection...`);
                return this.authDisconnect(socket, callback);
            }
        }

        jwt.verify(token, secret.jwtPrivateKey, {}, async (err, data) => {
            if (err) return this.authDisconnect(socket, callback);
            if (!data) return this.authDisconnect(socket, callback);
            if (!data.username) return this.authDisconnect(socket, callback);

            const user = await User.findByUsername(data.username);

            if (!user) return this.authDisconnect(socket, callback);

            let permissionLevel = config.roleMap[user.role];
            if (!permissionLevel) {
                permissionLevel = 0;
            }

            if (permissionLevel < config.roleMap.USER) return this.authDisconnect(socket, callback);

            socket.user = {
                username: data.username,
                _id: user._id,
                token, // Maybe not secure
                permissionLevel
            };
            console.log(`[*] [gateway] [handshake] User ${data.username} has successfully authenticated`);
            return callback();
        });
    });

    this._gateway.on('connection', (socket) => {
        console.log(`[*] [gateway] [handshake] User ${socket.user.username} connected, sending hello and waiting for yoo...`);
        
        socket.emit('hello', {
            gatewayStillNotConnectedTimeoutMS: config.gatewayStillNotConnectedTimeoutMS,
            resolvedUser: {
                username: socket.user.username,
                _id: socket.user._id
            }
        });
        
        socket.once('yoo', () => {
            console.log(`[*] [gateway] [handshake] Got yoo from ${socket.user.username}, connection is finally completed!`);
            socket.isConnected = true;

            socket.on('message', ({ category, content }) => {
                if (!category || !content || !socket.joinedCategories || !socket.isConnected || !socket.user || !(typeof content === 'string') || !(typeof category._id === 'string')) return;
                content = content.trim();
                if (!content || content === '' || content === ' ' || content.length >= 2000) return;
                if (!this.rateLimiter.consoom(socket.user.token)) { // TODO: maybe user ip instead of token?
                    console.log(`[E] [gateway] Rate limiting ${socket.user.username}`);
                    return;
                }


                // TODO: When/if category permissions are added, check if the user has permissions for that category
                const categoryTitle = socket.joinedCategories[category._id];
                if (!categoryTitle || !(typeof categoryTitle === 'string')) return;

                const messageObject = {
                    author: {
                        username: socket.user.username,
                        _id: socket.user._id
                    },
                    category: {
                        title: categoryTitle,
                        _id: category._id
                    },
                    content: content,
                    _id: uuid.v4()
                };

                if (messageObject.content.startsWith(this._commandPrefix)) {
                    this._processCommand(socket, messageObject);
                    return;
                }

                this._gateway.in(category._id).emit('message', messageObject);
            });

            socket.on('subscribe', async (categories) => {
                if ( !socket.isConnected || !socket.user || !categories || !Array.isArray(categories) || categories === []) return;
                for (const v of categories) {
                    if (!v && !(typeof v === 'string')) continue;
                    // TODO: When/if category permissions are added, check if the user has permissions for that category
                    const category = await Category.findById(v);
                    if (category && category.title && category._id) {
                        if (!socket.joinedCategories) socket.joinedCategories = {};
                        if (socket.joinedCategories[v]) continue;
                        socket.joinedCategories[v] = category.title;
                        await socket.join(v);

                        console.log(`[*] [gateway] User ${socket.user.username} subscribed to room ${v} (${category.title}), sending updated user list to all members of that room...`);

                        const upd = await this._generateClientListUpdateObject(v, category.title);
                        this._gateway.in(v).emit('clientListUpdate', upd);
                    }
                }    
            });

            socket.on('disconnecting', async () => {
                console.log(`[*] [gateway] User ${socket.user.username} is disconnecting, broadcasting updated user list to all of the rooms they have been in...`);
                const rooms = socket.rooms;
                rooms.forEach(async (room) => {
                    // Socket io automatically adds a user to a room with their own id
                    if (room === socket.id) return;

                    const categoryTitle = socket.joinedCategories[room] || 'UNKNOWN';
                    await socket.leave(room);

                    const upd = await this._generateClientListUpdateObject(room, categoryTitle);
                    socket.in(room).emit('clientListUpdate', upd);
                });
            });
        });
    });
};

GatewayServer.prototype._getSocketsInRoom = async function(room) {
    // NOTE: I have no idea why i have to do this dumb thing, why can't socket io just let you simply get the sockets from a room? idk
    // There kinda was a way in the previous version, but they want to change the api for the worse each version, i'm guessing
    const clients = await this._gateway.in(room).allSockets();
    const updatedClientList = [];

    clients.forEach((sid) => {
        const client = this._gateway.sockets.get(sid); // lol they also used dumb ass maps for the socket list, can you fucking not?
        if (!client || !client.isConnected || !client.user) return;
        updatedClientList.push({
            user: {
                username: client.user.username,
                _id: client.user._id
            }
        });
    });
    return updatedClientList;
};

GatewayServer.prototype._generateClientListUpdateObject = async function(room, categoryTitle='UNKNOWN') {
    const clientList = await this._getSocketsInRoom(room);
    return {
        category: {
            title: categoryTitle,
            _id: room
        },
        clientList
    };
};


module.exports = GatewayServer;