wormhole/lib/Socket.js

const EventEmitter = require('events');

const constants = require("./constants");
const parser = require('./Parser');

class Socket extends EventEmitter {
    constructor({ initialState='CONNECTING', socket }) {
        super();

        this._state = initialState;
        this._socket = socket;
        this._accepted = false;
        this._fateDecided = false; // Wether the decision to accept or reject the socket was made

        this._buffering = undefined;

        this._socket.on('data', (e) => {
            if (!this._isConnected()) return;

            const packet = this._decodePayload(e.buffer);

            console.log(packet);

            if (this._buffering && packet.FIN === 1) {
                this.emit('binary', Buffer.concat([this._buffering, packet.PayloadData]));
                this._buffering = undefined;
            }

            if (packet.FIN === 0) {
                if (!this._buffering) {
                    this._buffering = Buffer.from(packet.PayloadData);
                } else {
                    Buffer.concat([this._buffering, Buffer.from(packet.PayloadData)]);
                }
                
                console.log('Buffering', packet);
                return;
            }

            if (packet.Opcode === constants.opcodes.TEXT_FRAME) {
                this.emit('text', packet);
            } else if (packet.Opcode === constants.opcodes.BINARY_FRAME) {
                this.emit('binary', packet);
            }
        });
    }
}

Socket.prototype.send = function(payload) {
    if (typeof payload === 'string') {
        this._sendBuffer(Buffer.from(payload));
    } else if (payload instanceof Buffer) {
        this._sendBuffer(payload);
    } else {
        throw new Error(`Unsupported type ${typeof payload}, supported types are: string, Buffer`);
    }
};

Socket.prototype._sendBuffer = function(buffer) {
    const frame = new parser.WebsocketFrame();
    frame.FIN = 1;
    frame.RSVx = 0;
    frame.Opcode = 0x01;

    frame.MASK = 0;
    const length = buffer.byteLength;
    if (length > 125) {
        frame.PayloadLen = 126;
    } else if (length >= 65536) {
        frame.PayloadLen = 127;
    } else if (length <= 125) {
        frame.PayloadLen = length;
    }
    frame.PayloadLenEx = length;
    frame.PayloadData = buffer;

    this._socket.write(frame.toBuffer());
};

Socket.prototype._decodePayload = function(payload) {
    return parser.parseWebsocketFrame(new DataView(payload));
};

Socket.prototype._setConnectionState = function(state) {
    this._state = state;
};

Socket.prototype._setAccepted = function(state) {
    if (this._fateDecided) throw new Error('Tried to decide fate (wether socket is accepted or not) more than 1 time');

    this._fateDecided = true;
    this._accepted = state;
};

Socket.prototype._isConnected = function() {
    return (this._fateDecided && this._isConnected && this._accepted === true);
}

module.exports = Socket;