Writing a WebSocket Server

Writing a WebSocket Server

·

7 min read

WebSocket Server

The WebSocket Protocol enables two-way communication between a client and server. It is built upon HTTP protocol and defined in rfc6455.

In this article, let's explore some details about websocket server, and try to implement it based on Node.js http package.

Client

Since there is built-in websocket client in browser, so we use it to test server code. The logic is simple, establish connection, send message hello.

const ws = new WebSocket("ws://localhost:1988/ws");
ws.onopen = () => {
    console.log("onopen");
    ws.send("hello");
};
ws.onerror = e => {
    console.error("onerror", e);
};
ws.onclose = e => {
    console.log("onclose", e);
};
ws.onmessage = e => {
    console.log("receive", e.data);
}

Handshake

To establish a websocket connection, the client should send a http request first. This request is required to provide some special headers. An example header is like below.

GET /chat HTTP/1.1
Host: server.example.com
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==
Origin: http://example.com
Sec-WebSocket-Version: 13

Once the client sends this request, the server should handle it properly.

In node.js, we can write a simple http server like below.

const http = require('http');

const server = http.createServer(function (req, res) {
    res.writeHead(200, { 'Content-Type': 'text/html' });
    res.end('<p>Hello World</p>');
});

server.listen(1988);

But we can't handle this websocket handshake request like this, because we want to reuse the same tcp connection after handshake.

Node.js provides a upgrade event to handle cases like this. we can listen to this event then we are able to access the raw tcp connection object to do further operations.

server.on('upgrade', (req, socket) => {
    // ...
});

Then the server should send a response like below to complete this handshake.

HTTP/1.1 101 Switching Protocols
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=
Sec-WebSocket-Protocol: chat

Note that there is a Sec-WebSocket-Accept header, this value is calculate based on the header Sec-WebSocket-Key from request. The calculation process is defined in the protocol, and code is like below.

function calculateSWA(swk) {
    return crypto
        .createHash('sha1')
        .update(swk + '258EAFA5-E914-47DA-95CA-C5AB0DC85B11', 'binary')
        .digest('base64');
}

Then handshake is completed.

server.on('upgrade', (req, socket) => {
    console.log("upgrade", req.url);

    if (req.headers['upgrade'] !== 'websocket') {
        socket.end('HTTP/1.1 400 Bad Request');
        return;
    }

    const swk = req.headers['sec-websocket-key'];
    const swa = calculateSWA(swk);
    const responseHeaders = [
        'HTTP/1.1 101 Switching Protocols',
        'Upgrade: websocket',
        'Connection: Upgrade',
        `Sec-WebSocket-Accept: ${swa}`,
    ].join("\r\n") + "\r\n\r\n";

    socket.write(responseHeaders);
});

Messaging

After handshake, the websocket connection is established. The next step is to send/receive message between client and server.

// recieve data
socket.on('data', buffer => {
    // send data
    socket.write("...");
});

According to the websocket protocol, data format is like below.

      0                   1                   2                   3
      0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
     +-+-+-+-+-------+-+-------------+-------------------------------+
     |F|R|R|R| opcode|M| Payload len |    Extended payload length    |
     |I|S|S|S|  (4)  |A|     (7)     |             (16/64)           |
     |N|V|V|V|       |S|             |   (if payload len==126/127)   |
     | |1|2|3|       |K|             |                               |
     +-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - +
     |     Extended payload length continued, if payload len == 127  |
     + - - - - - - - - - - - - - - - +-------------------------------+
     |                               |Masking-key, if MASK set to 1  |
     +-------------------------------+-------------------------------+
     | Masking-key (continued)       |          Payload Data         |
     +-------------------------------- - - - - - - - - - - - - - - - +
     :                     Payload Data continued ...                :
     + - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - +
     |                     Payload Data continued ...                |
     +---------------------------------------------------------------+

So first thing is to parse this data to get the payload data.

You can check the meaning of each fields from the protocol. I will not repeat in here. The parsing code is below.

function recvMsg(buffer) {
    let offset = 0;

    const FIN = (buffer[offset] >> 7) & 1;
    console.log({ FIN });

    const Opcode = buffer[offset] & 0b1111;
    console.log({ Opcode });
    switch (Opcode) {
        case 0x0:
            console.log("continuation frame");
            break;
        case 0x1:
            console.log("text frame");
            break;
        case 0x2:
            console.log("binary frame");
            break;
        case 0x8:
            console.log("connection close");
            break;
        case 0x9:
            console.log("ping");
            break;
        case 0xA:
            console.log("pong");
            break;
        default:
            console.log("invalid Opcode: ", Opcode);
    }

    offset += 1;

    const Mask = (buffer[offset] >> 7) & 1;
    console.log({ Mask });

    // 7 bits
    let PayloadLength = buffer[offset] & 0b01111111;

    offset += 1;

    // next 16 bits
    if (PayloadLength === 126) {
        PayloadLength = buffer.readUInt16BE(offset);
        offset += 2;
    }

    // next 64 bits
    if (PayloadLength === 127) {
        // ...
        offset += 8;
    }
    console.log({ PayloadLength });

    const MaskingKey = buffer.subarray(offset, offset + 4);
    console.log({ MaskingKey });

    offset += 4;

    const PayloadData = buffer.subarray(offset);
    console.log({ PayloadData }, PayloadData.length);

    let message;
    if (Mask > 0) {
        message = unmask(MaskingKey, PayloadData).toString();
    } else {
        message = PayloadData.toString();
    }

    return message;
}

One thing needed to be noted is that data can be masked according to the Mask field. If the data is masked, we need to unmask it. The key idea of unmasking is to use xor operations between MaskingKey and PayloadData, its logic is like below.

function unmask(MaskingKey, PayloadData) {
    for (let i = 0; i < PayloadData.length; i++) {
        PayloadData[i] = MaskingKey[i % 4] ^ PayloadData[i];
    }
    return PayloadData;
}

After receiving message, we also want to send message to the client. Because the data format is the same, we only need to construct it according to the format. It's simpler than parsing.

function sendMsg(message) {
    const PayloadData = Buffer.from(message);
    // todo: handle message bigger than length 125
    const header = Buffer.from([0b10000001, PayloadData.length]);
    const buffer = Buffer.concat([header, PayloadData]);
    return buffer;
}

OK, now we can send/receive messages.

socket.on('data', buffer => {
    const message = recvMsg(buffer);
    console.log("receive: ", message);

    const msgBuffer = sendMsg("World");
    socket.write(msgBuffer);
});

In the end

With all above code, we implement the handshake and messaging in websocket protocol. Although these 2 part are the key of the protocol, since this is a pretty complicated protocol, there are still many contents needed to be considered. Check the protocol by yourself if you're interested. Below I share the whole server code in case if you want to run by yourself.

const http = require('http');
const crypto = require('crypto');

function calculateSWA(swk) {
    return crypto
        .createHash('sha1')
        .update(swk + '258EAFA5-E914-47DA-95CA-C5AB0DC85B11', 'binary')
        .digest('base64');
}

const server = http.createServer(function (req, res) {
    res.writeHead(200, { 'Content-Type': 'text/html' });
    res.end('<p>Hello World</p>');
});


server.on('upgrade', (req, socket) => {
    console.log("upgrade", req.url);

    if (req.headers['upgrade'] !== 'websocket') {
        socket.end('HTTP/1.1 400 Bad Request');
        return;
    }

    const swk = req.headers['sec-websocket-key'];
    const swa = calculateSWA(swk);
    const responseHeaders = [
        'HTTP/1.1 101 Switching Protocols',
        'Upgrade: websocket',
        'Connection: Upgrade',
        `Sec-WebSocket-Accept: ${swa}`,
    ].join("\r\n") + "\r\n\r\n";

    socket.write(responseHeaders);


    socket.on('data', buffer => {
        const message = recvMsg(buffer);
        console.log("receive: ", message);

        const msgBuffer = sendMsg("World");
        socket.write(msgBuffer);
    });
});


function recvMsg(buffer) {
    let offset = 0;

    const FIN = (buffer[offset] >> 7) & 1;
    console.log({ FIN });

    const Opcode = buffer[offset] & 0b1111;
    console.log({ Opcode });
    switch (Opcode) {
        case 0x0:
            console.log("continuation frame");
            break;
        case 0x1:
            console.log("text frame");
            break;
        case 0x2:
            console.log("binary frame");
            break;
        case 0x8:
            console.log("connection close");
            break;
        case 0x9:
            console.log("ping");
            break;
        case 0xA:
            console.log("pong");
            break;
        default:
            console.log("invalid Opcode: ", Opcode);
    }

    offset += 1;

    const Mask = (buffer[offset] >> 7) & 1;
    console.log({ Mask });

    // 7 bits
    let PayloadLength = buffer[offset] & 0b01111111;

    offset += 1;

    // next 16 bits
    if (PayloadLength === 126) {
        PayloadLength = buffer.readUInt16BE(offset);
        offset += 2;
    }

    // next 64 bits
    if (PayloadLength === 127) {
        // ...
        offset += 8;
    }
    console.log({ PayloadLength });

    const MaskingKey = buffer.subarray(offset, offset + 4);
    console.log({ MaskingKey });

    offset += 4;

    const PayloadData = buffer.subarray(offset);
    console.log({ PayloadData }, PayloadData.length);

    let message;
    if (Mask > 0) {
        message = unmask(MaskingKey, PayloadData).toString();
    } else {
        message = PayloadData.toString();
    }

    return message;
}


function unmask(MaskingKey, PayloadData) {
    for (let i = 0; i < PayloadData.length; i++) {
        PayloadData[i] = MaskingKey[i % 4] ^ PayloadData[i];
    }
    return PayloadData;
}

function sendMsg(message) {
    const PayloadData = Buffer.from(message);
    // todo: handle message bigger than length 125
    const header = Buffer.from([0b10000001, PayloadData.length]);
    const buffer = Buffer.concat([header, PayloadData]);
    return buffer;
}

server.listen(1988);