Minimal WebSocket Server in NodeJS/TypeScript, no dependencies

minimalWsServer.ts

import { createHash } from "crypto";
import { createServer, IncomingMessage } from "http";
import { Duplex } from "stream";

type WsFn = (val: string | Uint8Array) => void;

// Create this object whenever a websocket connection has been established
type WebSocket = {
  write: WsFn;
  addListener: (listener: WsFn) => void;
  removeListener: (listener: WsFn) => void;
  addCloseListener: (listener: () => void) => void;
  close: () => void;
};

//  creates a server that simply echos what is sent to it
makeWsServer((ws: WebSocket) => {
  ws.addListener((val: string | Uint8Array) => {
    console.log("Got message!!!", val);
    ws.write(val);
  });
});

/**
 * Data Frame Parsing
 *
 * Data passed in WebSockets takes place within a data frame containing
 * - fin: a boolean indicating that the data frameis the end of a complete message
 * - opCode: one of 5 reserved operations
 * - hasMask: a boolean containing whether the payload should be masked by being
 *      XOR'd with a 4-byte repeating mask. The server does not need to transmit
 *      masked payloads, and the client is required to do so.
 * - payloadLength: 7, 16, or 64-bit number indicating the length of the payload
 * - mask: 4-byte number for masking payload, provided by client
 * - payload: the good stuff
 *
 * This section deals with reading and writing the data frames
 *
 */
type DataFrame = {
  fin: boolean;
  opCode: WsOpCode;
  hasMask?: boolean;
  payload: Uint8Array;
};
enum WsOpCode {
  CONTINUATION = 0,
  TEXT = 1,
  BINARY = 2,
  CONNECTION_CLOSE = 8,
  PING = 9,
  PONG = 10,
  UNDEFINED = 0xff,
}

function findPayloadLen(data: Buffer, offset: number): [number, number] {
  const first = data[offset] & 0x7f;
  if (first === 126) {
    return [data.readUInt16BE(offset + 1), 2];
  } else if (first === 127) {
    const val = data.readBigUInt64BE(offset + 1);
    if (val > Number.MAX_SAFE_INTEGER) {
      throw new Error("Can't handle that kind of payload");
    }
    return [Number(val), 4];
  }
  return [first, 1];
}

function parseDataFrame(data: Buffer) {
  let offset = 0;
  const fin = Boolean((data[offset] >> 7) & 1);
  const opCodeInt = data[offset] & 0xf;
  const opCode: WsOpCode | null =
    opCodeInt in WsOpCode ? opCodeInt : WsOpCode.UNDEFINED;

  offset++;

  const hasMask = Boolean((data[offset] >> 7) & 1);

  const [payloadLen, payloadWidth] = findPayloadLen(data, offset);
  offset += payloadWidth;
  const mask: Uint8Array | null = hasMask
    ? data.subarray(offset, offset + 4)
    : null;
  offset += hasMask ? 4 : 0;

  const payload =
    mask !== null
      ? Uint8Array.from(
          data.subarray(offset, offset + payloadLen),
          (maskedByte, index) => maskedByte ^ mask[index % 4]
        )
      : Uint8Array.from(data.subarray(offset, offset + payloadLen));

  return { fin, opCode, hasMask, payloadLen, mask, payload };
}

function sendDataFrame(df: DataFrame, socket: Duplex) {
  // payload length is encoded in 1, 2, or 4 byte values depending on size
  const sizeBytes =
    df.payload.length < 126
      ? [df.payload.length]
      : df.payload.length < 65536
      ? [126, df.payload.length >> 8, df.payload.length & 0xff]
      : [
          127,
          df.payload.length >> 24,
          (df.payload.length >> 16) & 0xff,
          (df.payload.length >> 8) & 0xff,
        ];
  socket.write(
    Uint8Array.from([
      (Number(df.fin) << 7) | df.opCode,
      ...sizeBytes,
      ...df.payload,
    ])
  );
}

// the PONG message is the PING message with a PONG opcode
function sendPong(ping: DataFrame, socket: Duplex) {
  sendDataFrame({ ...ping, opCode: WsOpCode.PONG }, socket);
}

/**
 * Upgrade Response
 *
 * The Http Server must respond to a Connection: Upgrade request with the following HTTPish headers
 * in order to signal to the client that it can handle websocket connections
 *
 * Technically there are extensions and subprotocols that could also be referenced here
 */
function upgradeConnection(req: IncomingMessage, socket: Duplex) {
  const clientKey = req.headers["sec-websocket-key"];
  const reponseHeaders = [
    "HTTP/1.1 101 Switching Protocols",
    "Upgrade: websocket",
    `Sec-WebSocket-Accept: ${createHash("sha1")
      .update(clientKey + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11")
      .digest("base64")}`,
    "Connection: Upgrade",
    "\r\n",
  ].join("\r\n");
  socket.write(reponseHeaders);
}

const encoder = new TextEncoder();
const decoder = new TextDecoder();

function listenToSocket(socket: Duplex, listeners: WsFn[]) {
  let strBuffer = "";

  let bufOffset = 0;
  let buffer = new Uint8Array(2048);

  socket.on("data", (data: Buffer) => {
    const parsedFrame = parseDataFrame(data);

    if (parsedFrame.opCode === WsOpCode.PING) {
      sendPong(parsedFrame, socket);
      return;
    }
    if (parsedFrame.opCode === WsOpCode.TEXT) {
      strBuffer += decoder.decode(parsedFrame.payload);
      if (parsedFrame.fin) {
        listeners.forEach((listener) => listener(strBuffer));
        strBuffer = "";
      }
    } else if (parsedFrame.opCode === WsOpCode.BINARY) {
      // somewhat efficiently grow the underlying buffer
      if (parsedFrame.payloadLen + bufOffset > buffer.length) {
        let newSize = buffer.length * 2;
        while (parsedFrame.payloadLen + bufOffset > newSize) {
          newSize *= 2;
        }

        const newBuffer = new Uint8Array(newSize);
        newBuffer.set(buffer);
        buffer = newBuffer;
      }
      buffer.set(parsedFrame.payload, bufOffset);
      bufOffset += parsedFrame.payloadLen;

      if (parsedFrame.fin) {
        listeners.forEach((listener) => listener(buffer));
        bufOffset = 0;
      }
    } else if (parsedFrame.opCode === WsOpCode.CONNECTION_CLOSE) {
      socket.end();
    }
  });
}

export function makeWsServer(handleUpgrade: (handle: WebSocket) => void) {
  const srv = createServer();

  srv.on("upgrade", (req: IncomingMessage, socket: Duplex, head: Buffer) => {
    console.log("upgrading");
    upgradeConnection(req, socket);

    const closeListeners: (() => void)[] = [];
    const addCloseListener = (listener: () => void) => listeners.push(listener);
    // note that if the socket is closed real quick none of the listeners will trigger
    socket.on("close", () => closeListeners.forEach((closer) => closer()));

    let listeners: WsFn[] = [];
    const addListener = (listener: WsFn) => listeners.push(listener);
    const removeListener = (removed: WsFn) =>
      (listeners = listeners.filter((listener) => listener !== removed));

    // write string/binary data via a single dataframe
    const write = (val: string | Uint8Array) => {
      // Don't bother fragmenting
      if (typeof val === "string") {
        const payload = encoder.encode(val);

        sendDataFrame(
          {
            fin: true,
            opCode: WsOpCode.TEXT,
            payload,
          },
          socket
        );
      } else {
        sendDataFrame(
          {
            fin: true,
            opCode: WsOpCode.BINARY,
            payload: val,
          },
          socket
        );
      }
    };

    // a WebSocket connection closing down should be preceeded by the closing side
    // first broadcasting one last CONNECTION_CLOSE dataframe
    const close = () => {
      sendDataFrame(
        {
          fin: true,
          opCode: WsOpCode.CONNECTION_CLOSE,
          hasMask: false,
          payload: Uint8Array.from([]),
        },
        socket
      );
    };

    listenToSocket(socket, listeners);

    handleUpgrade({
      write,
      addListener,
      removeListener,
      close,
      addCloseListener,
    });
  });

  srv.listen(1337, "127.0.0.1", () =>
    console.log("listening on 127.0.0.1:1337")
  );
}