index.js

/* global Blob */

module.exports = Socket

var debug = require('debug')('simple-websocket')
var inherits = require('inherits')
var isTypedArray = require('is-typedarray')
var stream = require('stream')
var toBuffer = require('typedarray-to-buffer')
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
  stream.Duplex.call(self, opts)

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

  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 = self._onError.bind(self)

  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.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)

  self.connected = false
  self.destroyed = true

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

  if (self._ws) {
    try {
      self._ws.close()
    } catch (err) {}

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

  this.readable = this.writable = false

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

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

Socket.prototype._read = function () {}

/**
 * Send text/binary data to the WebSocket server.
 * @param {string|Buffer|TypedArrayView|ArrayBuffer|Blob} chunk
 * @param {string} encoding
 * @param {function} cb
 */
Socket.prototype._write = function (chunk, encoding, cb) {
  var self = this
  if (self.destroyed) return cb(new Error('cannot write after socket is destroyed'))

  if (!isTypedArray.strict(chunk) && !(chunk instanceof ArrayBuffer) &&
    !Buffer.isBuffer(chunk) && typeof chunk !== 'string' &&
    (typeof Blob === 'undefined' || !(chunk instanceof Blob))) {
    chunk = JSON.stringify(chunk)
  }

  if (self.connected) {
    self.send(chunk)
    if (typeof ws !== 'function' && self._ws.bufferedAmount) {
      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 = toBuffer(new Uint8Array(data))
    self.push(data)
  } else if (Buffer.isBuffer(data)) {
    self.push(data)
  } else {
    try {
      data = JSON.parse(data)
    } catch (err) {}
    self.emit('data', data)
  }
}

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

  if (self._chunk) {
    self.send(self._chunk)
    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) 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 () {
  var self = this
  if (self.destroyed) return
  var err = new Error('connection error to ' + self.url)
  debug('error: %s', err.message || err)
  self._destroy(err)
}