Websocket server implementation

[filipef101]

I wrote this websocket server implementation, if someone finds it useful Buffer polyfil and crypto required, the buffer polyfil could benefit from some native speed ups but it works perfectly

import QuickCrypto from "react-native-quick-crypto"
import { EventEmitter } from "events"
var Buffer: Buffer = require("buffer/").Buffer

interface WebSocketFrame {
  fin: boolean
  opcode: number
  payload: Buffer
  payloadString: string
}

const enum WebSocketOpCode {
  Continuation = 0x0,
  Text = 0x1,
  Binary = 0x2,
  Close = 0x8,
  Ping = 0x9,
  Pong = 0xa,
}

export class WebSocketServer extends EventEmitter {
  private clients: Set<any> = new Set()

  handleUpgrade(socket: any, headers: Record<string, string>) {
    console.log("Handling WebSocket upgrade")
    try {
      const key = headers["sec-websocket-key"]
      if (!key) {
        console.error("No WebSocket key provided")
        socket.end()
        return
      }

      const acceptKey = this.generateAcceptValue(key)
      const responseHeaders = [
        "HTTP/1.1 101 Switching Protocols",
        "Upgrade: websocket",
        "Connection: Upgrade",
        `Sec-WebSocket-Accept: ${acceptKey}`,
        "",
        "",
      ].join("\r\n")

      socket.write(responseHeaders)
      this.clients.add(socket)

      socket.on("data", (data: Buffer) => {
        try {
          const frame = this.decodeWebSocketFrame(data)
          this.processFrame(socket, frame)
        } catch (error) {
          console.error("Error processing WebSocket data:", error)
        }
      })

      socket.on("end", () => {
        console.log("Client disconnected:", socket.address())
        this.clients.delete(socket)
        this.emit("client.disconnect", socket.address())
      })

    } catch (error) {
      console.error("Error during WebSocket upgrade:", error)
      socket.end()
    }
  }

  private processFrame(socket: any, frame: WebSocketFrame) {
    try {
      // Handle control frames
      if (frame.opcode >= 0x8) {
        switch (frame.opcode) {
          case WebSocketOpCode.Close:
            socket.end()
            return
          case WebSocketOpCode.Ping:
            const pongFrame = this.encodeWebSocketFrame(frame.payloadString, WebSocketOpCode.Pong)
            socket.write(pongFrame)
            return
          case WebSocketOpCode.Pong:
            return
        }
        return
      }

      // Handle data frames
      if (frame.opcode === WebSocketOpCode.Text) {
        console.log("Processing frame payload:", frame.payloadString.slice(0, 200))
        const command = JSON.parse(frame.payloadString)
        this.emit("command", command)
      }
    } catch (error) {
      console.error("Error processing frame:", error)
    }
  }

  private generateAcceptValue(key: string): string {
    const GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
    const hash = QuickCrypto.createHash("sha1")
    hash.update(key + GUID)
    return hash.digest("base64")
  }

  private decodeWebSocketFrame(buffer: Buffer): WebSocketFrame {
    let offset = 0
    const firstByte = buffer[offset++]
    const secondByte = buffer[offset++]

    const fin = Boolean(firstByte & 0x80)
    const opcode = firstByte & 0x0f
    const masked = Boolean(secondByte & 0x80)
    let payloadLength = secondByte & 0x7f

    if (payloadLength === 126) {
      payloadLength = buffer.readUInt16BE(offset)
      offset += 2
    } else if (payloadLength === 127) {
      payloadLength = Number(buffer.readBigUInt64BE(offset))
      offset += 8
    }

    const maskingKey = masked ? buffer.slice(offset, offset + 4) : null
    offset += masked ? 4 : 0

    let payload = buffer.slice(offset, offset + payloadLength)

    if (masked && maskingKey) {
      payload = this.unmaskPayload(payload, maskingKey)
    }

    return {
      fin,
      opcode,
      payload,
      payloadString: payload.toString("utf8"),
    }
  }

  private unmaskPayload(payload: Buffer, maskingKey: Buffer): Buffer {
    const result = Buffer.alloc(payload.length)
    for (let i = 0; i < payload.length; i++) {
      result[i] = payload[i] ^ maskingKey[i % 4]
    }
    return result
  }

  private encodeWebSocketFrame(data: string, opcode = WebSocketOpCode.Text): Buffer {
    const payload = Buffer.from(data)
    const payloadLength = payload.length

    let headerLength = 2
    if (payloadLength > 65535) {
      headerLength += 8
    } else if (payloadLength > 125) {
      headerLength += 2
    }

    const frame = Buffer.alloc(headerLength + payloadLength)
    frame[0] = 0x80 | opcode // FIN + opcode

    if (payloadLength > 65535) {
      frame[1] = 127
      frame.writeBigUInt64BE(BigInt(payloadLength), 2)
    } else if (payloadLength > 125) {
      frame[1] = 126
      frame.writeUInt16BE(payloadLength, 2)
    } else {
      frame[1] = payloadLength
    }

    payload.copy(frame, headerLength)
    return frame
  }

  broadcast(type: string, payload: any) {
    const message = JSON.stringify({
      type,
      payload,
      date: new Date().toISOString(),
    })

    const frame = this.encodeWebSocketFrame(message)
    this.clients.forEach((client) => {
      try {
        client.write(frame)
      } catch (error) {
        console.error("Error sending message:", error)
      }
    })
  }
}

// Create a singleton instance
export const wsServer = new WebSocketServer()