WebSocketHook.js

window.WebSocket = class WebSocket extends window.WebSocket {
  static event_listener = {}

  static on(eventName, fn) {
    WebSocket.event_listener[eventName] = WebSocket.event_listener[eventName] || []
    WebSocket.event_listener[eventName].push(fn)
  }
  static on(eventName, fn) {
    WebSocket.event_listener[eventName] = WebSocket.event_listener[eventName] || []
    WebSocket.event_listener[eventName].push(fn)
  }

  static emit(eventName, ...args) {
    WebSocket.event_listener[eventName] = WebSocket.event_listener[eventName] || []
    WebSocket.event_listener[eventName].forEach(fn => fn(...args))
  }

  constructor(url, protocols) {
    super(url, protocols)
    const { origin } = new URL(url)
    this.origin = origin
    this.onmessage_listener = []

    this.addEventListener('open', () => WebSocket.emit('open', this))
    this.addEventListener('close', () => WebSocket.emit('close', this))
    super.onmessage = this.handleMessageEvent.bind(this)

    WebSocket.emit('init', this)
  }

  rawSend(data) {
    super.send(data)
  }

  rawRecv(data, target=this) {
    if (!(data instanceof ArrayBuffer) && ArrayBuffer.isView(data) && data.buffer) {
      data = data.buffer
    }

    const event = new MessageEvent('message', {
      data,
      origin: this.origin
    })
    Object.defineProperty(event, 'target', {
      value: target,
      writable: false,
    })

    if (this.original_onmessage) this.original_onmessage(event)
    this.onmessage_listener.forEach(fn => fn(event))
  }

  send(data) {
    const message = { socket: this, data }
    WebSocket.emit('beforeSend', message)

    if (message.data) this.rawSend(message.data)

    WebSocket.emit('afterSend', message)
  }

  handleMessageEvent(event) {
    const message = { socket: this, data: event.data, event }
    WebSocket.emit('beforeRecv', message)

    if (message.data) this.rawRecv(message.data, event.target)

    WebSocket.emit('afterRecv', message)
  }

  addEventListener(name, listener) {
    if (name === 'message') return this.onmessage_listener.push(listener)

    return super.addEventListener(name, listener)
  }

  removeEventListener(name, listener) {
    if (name === 'message') {
      for (let i=this.onmessage_listener.length-1; i>=0; i--) {
        if (this.onmessage_listener[i] === listener) this.onmessage_listener.splice(i, 1)
      }
      return
    }
    return super.removeEventListener(name, listener)
  }

  set onmessage(listener) {
    this.original_onmessage = listener
  }
}