WsHandler.js

const JSONStream = require('JSONStream');
const path = require('path');
const { pipeline } = require('stream');
const url = require('url');
const WebSocket = require('ws');

const ClientMessageHandler = require('./ClientMessageHandler');
const DemuxSink = require('./demux');
const logger = require('./logging');
const PromCollector = require('./metrics/PromCollector');
const { getStatsFormat } = require('./utils/stats-detection');
const { RequestType } = require('./utils/utils');
const cwd = process.cwd();

/**
 *
 */
class WsHandler {

    /**
     *
     */
    constructor({ tempPath, reconnectTimeout, sequenceNumberSendingInterval, workerPool, config }) {
        this.sessionTimeoutIds = {};
        this.tempPath = tempPath;
        this.reconnectTimeout = reconnectTimeout;
        this.sequenceNumberSendingInterval = sequenceNumberSendingInterval;
        this.processData = this.processData.bind(this);
        this.workerPool = workerPool;
        this.config = config;
        this.dumpFolder = './temp';
    }

    /**
     *
     * @param {*} wsServer
     */
    setupWebSocketsServer(wsServer) {
        const wss = new WebSocket.Server({ server: wsServer });

        wss.on('connection', this._handle.bind(this));

        return wss;
    }

    /**
     *
     * @param {*} meta
     * @param {*} connectionInfo
     */
    processData(id, dumpPath, connectionInfo) {
        logger.info('[WsHandler] Queue for processing id %s', id);

        // Metadata associated with a dump can get large so just select the necessary fields.
        const dumpData = {
            clientId: id,
            dumpPath,
            endDate: Date.now(),
            startDate: connectionInfo.startDate
        };

        // Don't process dumps generated by JVB & Jigasi, there should be a more formal process to
        if (this.config.features.disableFeatExtraction
            || connectionInfo.clientProtocol?.includes('JVB') || connectionInfo.clientProtocol?.includes('JIGASI')) {
            this.persistDumpData(dumpData);
        } else {
            // Add the clientId in the worker pool so it can process the associated dump file.
            this.workerPool.addTask({
                type: RequestType.PROCESS,
                body: dumpData
            });
        }
    }

    /**
     * Main handler for web socket connections.
     * Messages are sent through a node stream which saves them to a dump file.
     * After the websocket is closed the session is considered as terminated and the associated dump
     * is queued up for feature extraction through the {@code WorkerPool} implementation.
     *
     * @param {*} client
     * @param {*} upgradeReq
     */
    _handle(client, upgradeReq) {
        PromCollector.connected.inc();

        // the url the client is coming from
        const referer = upgradeReq.headers.origin + upgradeReq.url;
        const ua = upgradeReq.headers['user-agent'];
        const queryObject = url.parse(referer, true).query;
        const statsSessionId = queryObject?.statsSessionId;
        const dumpPath = this._getDumpPath(statsSessionId);

        this._clearConnectionTimeout(statsSessionId);
        const connectionInfo = this._createConnectionInfo(upgradeReq, referer, ua, client);
        const demuxSink = this._createDemuxSink(dumpPath, connectionInfo);

        logger.info('[WsHandler] Client connected: %s', statsSessionId);

        const clientMessageHandler = this._createClientMessageHandler(statsSessionId, demuxSink, client);


        clientMessageHandler.sendLastSequenceNumber(true);

        demuxSink.on('close-sink', ({ id }) => {

            logger.info(
                '[WsHandler] Websocket disconnected waiting for processing the data %s in %d ms',
                id,
                this.reconnectTimeout
            );

            const timeoutId = setTimeout(this.processData,
                this.reconnectTimeout,
                id, dumpPath, connectionInfo
            );

            this.sessionTimeoutIds[id] = timeoutId;
        });

        const connectionPipeline = pipeline(
            WebSocket.createWebSocketStream(client),
            JSONStream.parse(),
            demuxSink,
            err => {
                if (err) {
                    // A pipeline can multiplex multiple sessions however if one fails
                    // the whole pipeline does as well,
                    PromCollector.sessionErrorCount.inc();

                    logger.error('[WsHandler] Connection pipeline: %o;  error: %o', connectionInfo, err);
                }
            });

        connectionPipeline.on('finish', () => {
            logger.info('[WsHandler] Connection pipeline successfully finished %o', connectionInfo);

            // We need to explicity close the ws, you might notice that we don't do the same in case of an error
            // that's because in that case the error will propagate up the pipeline chain and the ws stream will also
            // close the ws.
            client.close();
        });

        logger.info(
            '[WsHandler] New app connected: ua: %s, protocol: %s, referer: %s',
            ua,
            client.protocol,
            referer
        );

        client.on('error', e => {
            logger.error('[WsHandler] Websocket error: %s', e);
            PromCollector.connectionError.inc();
        });

        client.on('close', () => {
            PromCollector.connected.dec();
        });
    }

    /**
     *
     */
    _createDemuxSink(dumpPath, connectionInfo) {
        const demuxSinkOptions = {
            tempPath: this.tempPath,
            connectionInfo,
            dumpPath,
            dumpFolder: this.dumpFolder,
            log: logger
        };

        return new DemuxSink(demuxSinkOptions);
    }

    /**
     *
     */
    _getDumpPath(statsSessionId) {
        return path.resolve(cwd, this.dumpFolder, statsSessionId);
    }

    /**
     *
     */
    _createClientMessageHandler(statsSessionId, demuxSink, client) {
        const clientMessageHandlerOptions = {
            statsSessionId,
            tempPath: this.tempPath,
            sequenceNumberSendingInterval: this.sequenceNumberSendingInterval,
            demuxSink,
            client
        };

        return new ClientMessageHandler(clientMessageHandlerOptions);
    }

    /**
     *
     * @returns
     */
    _createConnectionInfo(upgradeReq, referer, ua, client) {
        // During feature extraction we need information about the browser in order to decide which algorithms use.
        const connectionInfo = {
            path: upgradeReq.url,
            origin: upgradeReq.headers.origin,
            url: referer,
            userAgent: ua,
            clientProtocol: client.protocol,
            startDate: Date.now()
        };

        connectionInfo.statsFormat = getStatsFormat(connectionInfo);

        return connectionInfo;
    }

    /**
     * Clear the connection timeout if the user is reconnected/
     *
     * @param {*} timeoutId
     */
    _clearConnectionTimeout(sessionId) {
        const timeoutId = this.sessionTimeoutIds[sessionId];

        if (timeoutId) {
            clearTimeout(timeoutId);
            delete this.sessionTimeoutIds[timeoutId];
            logger.info('[WsHandler] Client reconnected. Clear timeout for connectionId: %s', sessionId);
            PromCollector.clientReconnectedCount.inc();
        }
    }
}

module.exports = WsHandler;