Skip to content

Instantly share code, notes, and snippets.

@Xavier577
Last active April 22, 2023 16:29
Show Gist options
  • Star 5 You must be signed in to star a gist
  • Fork 1 You must be signed in to fork a gist
  • Save Xavier577/e6be4fedae6ca4278879a025f450724c to your computer and use it in GitHub Desktop.
Save Xavier577/e6be4fedae6ca4278879a025f450724c to your computer and use it in GitHub Desktop.
Implementing a websocket server without any libraries with raw nodejs

Code snippet

import { createServer } from "http";
import crypto from "crypto";

const PORT = 8001;

// this is from the web-socket specification and not something that is generated
const WEBSOCKET_MAGIC_STRING_KEY = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

const SEVEN_BITS_INTEGER_MARKER = 125; // as byte: 01111101
const SIXTEEN_BITS_INTEGER_MARKER = 126; // as byte: 01111110
// const SIXTYFOUR_BITS_INTEGER_MARKER = 127; // as byte: 01111111
const MAXIMUM_SIXTEEN_BITS_INTEGER = 2 ** 16; // 2 ** 16 is 0 to 65536

const MASK_KEY_BYTES_LENGTH = 4;
const FIRST_BIT = 128;
const OPCODE_TEXT = 0x01;

function createSocketAcceptHeaderValue(webSocketSecKey: string) {
  const hash = crypto
    .createHash("sha1")
    .update(webSocketSecKey.concat(WEBSOCKET_MAGIC_STRING_KEY))
    .digest("base64");

  return hash;
}

function prepareHandshakeResponse(webSocketSecKey: string) {
  const acceptKey = createSocketAcceptHeaderValue(webSocketSecKey);

  const headerResponse = [
    "HTTP/1.1 101 Switching Protocols",
    "Upgrade: websocket",
    "Connection: Upgrade",
    `sec-webSocket-accept: ${acceptKey}`,
    // This empty line MUST be present for the response to be valid
    "",
  ]
    .map((line) => line.concat("\r\n"))
    .join("");

  return headerResponse;
}

function unmask(encodedBuffer: Buffer, maskKey: Buffer) {
  // helper funcations to help log process of unmasking (the XOR operation part)
  const fillWithZeroes = (t: string) => t.padStart(8, "0");
  const toBinary = (t: number) => fillWithZeroes(t.toString(2));
  const fromBinaryToDecimal = (t: number) => parseInt(toBinary(t), 2);
  const getCharFromBinary = (t: number) =>
    String.fromCharCode(fromBinaryToDecimal(t));

  const decoded = Uint8Array.from(encodedBuffer, (element, index) => {
    const decodedElement = element ^ maskKey[index % 4];

    console.log({
      unmakingCalc: `${toBinary(element)} ^ ${toBinary(
        maskKey[index % 4]
      )} = ${toBinary(decodedElement)}`,
      decodedElement: getCharFromBinary(decodedElement),
    });

    return decodedElement;
  });

  return Buffer.from(decoded);
}

function encodeWebsocketMsg(message: string): Buffer {
  const msg = Buffer.from(message);

  const messageSize = msg.length;

  // this would contain specify information defined on the websocket protocol
  let dataFrameBuffer: Buffer;

  const firstByte = 0x80 | OPCODE_TEXT;

  if (messageSize <= SEVEN_BITS_INTEGER_MARKER) {
    const bytes = [firstByte];

    dataFrameBuffer = Buffer.from(bytes.concat(messageSize));
  } else if (messageSize <= MAXIMUM_SIXTEEN_BITS_INTEGER) {
    const offsetFourBytes = 4;

    const target = Buffer.allocUnsafe(offsetFourBytes);

    target[0] = firstByte;

    // this is the mask indicator (0 means unmasked)
    target[1] = SIXTEEN_BITS_INTEGER_MARKER | 0x00;

    target.writeUint16BE(messageSize, 2);

    dataFrameBuffer = target;
  } else {
    throw new Error("message is too long :(");
  }

  const totalLength = dataFrameBuffer.byteLength + messageSize;

  const target = Buffer.allocUnsafe(totalLength);
  let offset = 0;

  for (const buffer of [dataFrameBuffer, msg]) {
    target.set(buffer, offset);
    offset += buffer.length;
  }

  return target;

  //  callBack(dataFrameBuffer);
}

const server = createServer((_request, response) => {
  response.writeHead(200);
  response.end("hey there");
}).listen(PORT, () => console.log("Server listening on port", PORT));

server.on("upgrade", (req, socket, _head) => {
  const { "sec-websocket-key": webSocketSecKey } = req.headers;

  console.log({ webSocketClientKey: webSocketSecKey });

  const response = prepareHandshakeResponse(webSocketSecKey as string);

  console.log({ headerResponse: response });

  socket.write(response, (err) => {
    if (err != null) {
      console.error(err);
    }
  });

  socket.on("readable", () => {
    // read the first byte this  (this contain the first to last fragment, we are not doing anything with it)
    socket.read(1);

    // read second byte and store it in a variable (this contains the payload length)
    const [markerAndPayloadLength] = socket.read(1);

    console.log({ markerAndPayloadLength });

    // Add these next lines
    const lengthIndicatorInBits = markerAndPayloadLength - FIRST_BIT;

    let messageLength = 0;

    console.log({ lengthIndicatorInBits, messageLength });

    if (lengthIndicatorInBits <= SEVEN_BITS_INTEGER_MARKER) {
      messageLength = lengthIndicatorInBits;
    } else if (lengthIndicatorInBits === SIXTEEN_BITS_INTEGER_MARKER) {
      // unsigned, big-endian 16-bit integer [0 - 65k] - 2 ** 16
      messageLength = socket.read(2).readUint16BE(0);
    } else {
      throw new Error(
        "your message is too long! we don't handle more than 125 characters in the payload"
      );
    }

    const maskKey = socket.read(MASK_KEY_BYTES_LENGTH);
    const encoded = socket.read(messageLength);
    const decoded = unmask(encoded, maskKey);
    const receivedData = decoded.toString("utf-8");

    const data = JSON.parse(receivedData);

    console.log({ maskKey, encoded, decoded, receivedData, data });

    const msg = JSON.stringify({
      message: data,
      at: new Date().toISOString(),
    });

    const encodedMsg = encodeWebsocketMsg(msg);

    socket.write(encodedMsg);
  });
});

const handleUncaughtExceptions = (err: any) => {
  console.error(err);
};

process.on("uncaughtException", handleUncaughtExceptions);

process.on("unhandledRejection", handleUncaughtExceptions);

Resources

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment