simple-websocket/index.js

module.exports = Socket

var debug = require('debug')('simple-websocket')
var inherits = require('inherits')
var stream = require('readable-stream')
var ws = require('ws') // websockets in node - will be empty object in browser

var WebSocket = typeof window !== 'undefined' ? window.WebSocket : ws

inherits(Socket, stream.Duplex)

/**
 * WebSocket. Same API as node core `net.Socket`. Duplex stream.
 * @param {string} url websocket server url
 * @param {Object} opts options to stream.Duplex
 */
function Socket (url, opts) {
  var self = this
  if (!(self instanceof Socket)) return new Socket(url, opts)
  if (!opts) opts = {}
  debug('new websocket: %s %o', url, opts)

  opts.allowHalfOpen = false
  if (opts.highWaterMark == null) opts.highWaterMark = 1024 * 1024

  stream.Duplex.call(self, opts)

  self.url = url
  self.connected = false
  self.destroyed = false

  self._maxBufferedAmount = opts.highWaterMark
  self._chunk = null
  self._cb = null
  self._interval = null

  self._ws = new WebSocket(self.url)
  self._ws.binaryType = 'arraybuffer'
  self._ws.onopen = self._onOpen.bind(self)
  self._ws.onmessage = self._onMessage.bind(self)
  self._ws.onclose = self._onClose.bind(self)
  self._ws.onerror = function () {
    self._onError(new Error('connection error to ' + self.url))
  }

  self.on('finish', function () {
    if (self.connected) {
      // When stream is finished writing, close socket connection. Half open connections
      // are currently not supported.
      // Wait a bit before destroying so the socket flushes.
      // TODO: is there a more reliable way to accomplish this?
      setTimeout(function () {
        self._destroy()
      }, 100)
    } else {
      // If socket is not connected when stream is finished writing, wait until data is
      // flushed to network at "connect" event.
      // TODO: is there a more reliable way to accomplish this?
      self.once('connect', function () {
        setTimeout(function () {
          self._destroy()
        }, 100)
      })
    }
  })
}

Socket.WEBSOCKET_SUPPORT = !!WebSocket

/**
 * Send text/binary data to the WebSocket server.
 * @param {TypedArrayView|ArrayBuffer|Buffer|string|Blob|Object} chunk
 */
Socket.prototype.send = function (chunk) {
  var self = this

  var len = chunk.length || chunk.byteLength || chunk.size
  self._ws.send(chunk)
  debug('write: %d bytes', len)
}

Socket.prototype.destroy = function (onclose) {
  var self = this
  self._destroy(null, onclose)
}

Socket.prototype._destroy = function (err, onclose) {
  var self = this
  if (self.destroyed) return
  if (onclose) self.once('close', onclose)

  debug('destroy (error: %s)', err && err.message)

  this.readable = this.writable = false

  if (!self._readableState.ended) self.push(null)
  if (!self._writableState.finished) self.end()

  self.connected = false
  self.destroyed = true

  clearInterval(self._interval)
  self._interval = null
  self._chunk = null
  self._cb = null

  if (self._ws) {
    var ws = self._ws
    var onClose = function () {
      ws.onclose = null
      self.emit('close')
    }
    if (ws.readyState === WebSocket.CLOSED) {
      onClose()
    } else {
      try {
        ws.onclose = onClose
        ws.close()
      } catch (err) {
        onClose()
      }
    }

    ws.onopen = null
    ws.onmessage = null
    ws.onerror = null
  }
  self._ws = null

  if (err) self.emit('error', err)
}

Socket.prototype._read = function () {}

Socket.prototype._write = function (chunk, encoding, cb) {
  var self = this
  if (self.destroyed) return cb(new Error('cannot write after socket is destroyed'))

  if (self.connected) {
    try {
      self.send(chunk)
    } catch (err) {
      return self._onError(err)
    }
    if (typeof ws !== 'function' && self._ws.bufferedAmount > self._maxBufferedAmount) {
      debug('start backpressure: bufferedAmount %d', self._ws.bufferedAmount)
      self._cb = cb
    } else {
      cb(null)
    }
  } else {
    debug('write before connect')
    self._chunk = chunk
    self._cb = cb
  }
}

Socket.prototype._onMessage = function (event) {
  var self = this
  if (self.destroyed) return
  var data = event.data
  debug('read: %d bytes', data.byteLength || data.length)

  if (data instanceof ArrayBuffer) data = new Buffer(data)
  self.push(data)
}

Socket.prototype._onOpen = function () {
  var self = this
  if (self.connected || self.destroyed) return
  self.connected = true

  if (self._chunk) {
    try {
      self.send(self._chunk)
    } catch (err) {
      return self._onError(err)
    }
    self._chunk = null
    debug('sent chunk from "write before connect"')

    var cb = self._cb
    self._cb = null
    cb(null)
  }

  // No backpressure in node. The `ws` module has a buggy `bufferedAmount` property.
  // See: https://github.com/websockets/ws/issues/492
  if (typeof ws !== 'function') {
    self._interval = setInterval(function () {
      if (!self._cb || !self._ws || self._ws.bufferedAmount > self._maxBufferedAmount) {
        return
      }
      debug('ending backpressure: bufferedAmount %d', self._ws.bufferedAmount)
      var cb = self._cb
      self._cb = null
      cb(null)
    }, 150)
    if (self._interval.unref) self._interval.unref()
  }

  debug('connect')
  self.emit('connect')
}

Socket.prototype._onClose = function () {
  var self = this
  if (self.destroyed) return
  debug('on close')
  self._destroy()
}

Socket.prototype._onError = function (err) {
  var self = this
  if (self.destroyed) return
  debug('error: %s', err.message || err)
  self._destroy(err)
}