tiny 1-way websocket for prototyping

ws.js

// MIT License
//
// Copyright (c) 2021 Daniel Cousens
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
//
//   see RFC6455
import { createHash } from 'node:crypto'
import { EventEmitter } from 'node:events'

function sha1 (str) {
  return createHash('sha1').update(str).digest('base64')
}

function wsEncode (message) {
  if (typeof message !== 'string') {
    message = JSON.stringify(message)
  }

  const { length } = message
  if (length > 65536) throw new Error('Message length exceeds 16-bits')

  const offset = length < 126 ? 2 : 4
  const buffer = Buffer.alloc(offset + length)
  buffer.writeUInt8(0b10000001, 0) // FIN ... TEXT

  if (length < 126) {
    buffer.writeUInt8(length & 127, 1)
  } else {
    buffer.writeUInt8(126, 1) // 16-bit flag
    buffer.writeUInt16BE(length, 2)
  }

  buffer.write(message, offset)
  return buffer
}

export function handleUpgrade (req, res, socket) {
  const {
    upgrade,
    'sec-websocket-key': key
  } = req.headers

  if (upgrade !== 'websocket') return null
  if (!key) return null
  if (key.length > 24) return null

  const accept = sha1(`${key}258EAFA5-E914-47DA-95CA-C5AB0DC85B11`)
  res.removeHeader('Date')
  res.removeHeader('Transfer-Encoding')
  res.writeHead(101, {
    'Connection': 'Upgrade',
    'Upgrade': upgrade,
    'Sec-WebSocket-Accept': accept,
  })
  res.end()

  const emitter = new EventEmitter()
  function fin () {
    emitter.emit('end')
  }

  function send (message) {
    try {
      socket.write(wsEncode(message))
    } catch (e) {
      fin()
    }
  }

  socket.on('close', fin)
  socket.on('end', fin)
  socket.on('error', fin)
  emitter.send = send
  return emitter
}