Allow wss plugin to be customized (#22)

Author: darkgl0w

README.md

@@ -1,26 +1,75 @@
 # fastify-websocket
-
-[![Greenkeeper badge](https://badges.greenkeeper.io/fastify/fastify-websocket.svg)](https://greenkeeper.io/)
+[![js-standard-style](https://img.shields.io/badge/code%20style-standard-brightgreen.svg?style=flat)](http://standardjs.com/)  [![Build Status](https://travis-ci.org/fastify/fastify-websocket.svg?branch=master)](https://travis-ci.org/fastify/fastify-websocket) [![Greenkeeper badge](https://badges.greenkeeper.io/fastify/fastify-websocket.svg)](https://greenkeeper.io/)
 
 WebSocket support for [Fastify](https://github.com/fastify/fastify).
 Built upon [websocket-stream](http://npm.im/websocket-stream).
 
-## Example
+## Install
+
+```
+npm install fastify-websocket --save
+```
+
+## Usage
+
+All you need to do is to add it to your project with `register` and you are done!
+
+### Example
 
 ```js
 'use strict'
 
 const fastify = require('fastify')()
 
-fastify.register(require('fastify-websocket'), { handle })
+const wssOtions = {
+  maxPayload: 1048576, // we set the maximum allowed messages size to 1 MiB (1024 bytes * 1024 bytes)
+  path: '/fastify', // we accept only connections matching this path e.g.: ws://localhost:3000/fastify
+  verifyClient: function (info, next) {
+    if (info.req.headers['x-fastify-header'] !== 'fastify is awesome !') {
+      return next(false) // the connection is not allowed
+    }
+
+    next(true) // the connection is allowed
+  }
+}
+
+fastify.register(require('fastify-websocket'), {
+  handle,
+  options: wssOptions
+})
 
 function handle (conn) {
-  conn.pipe(conn) // creates a echo server
+  conn.pipe(conn) // creates an echo server
 }
 
-fastify.listen(0)
+fastify.listen(3000, (err) => {
+  if (err) {
+    fastify.log.error(err);
+    process.exit(1);
+  }
+})
 ```
 
+## Options :
+`fastify-websocket` accept the same options as [`websocket-stream`](https://github.com/maxogden/websocket-stream#options) and as [`ws`](https://github.com/websockets/ws/blob/master/doc/ws.md#new-websocketserveroptions-callback) :
+
+- `objectMode` - Send each chunk on its own, and do not try to pack them in a single websocket frame.
+- `host` - The hostname where to bind the server.
+- `port` - The port where to bind the server.
+- `backlog` - The maximum length of the queue of pending connections.
+- `server` - A pre-created Node.js HTTP/S server.
+- `verifyClient` - A function which can be used to validate incoming connections.
+- `handleProtocols` - A function which can be used to handle the WebSocket subprotocols.
+- `path` - Accept only connections matching this path.
+- `noServer` - Enable no server mode.
+- `clientTracking` - Specifies whether or not to track clients.
+- `perMessageDeflate` - Enable/disable permessage-deflate.
+- `maxPayload` - The maximum allowed message size in bytes.
+
+For more informations you can check [`ws` options documentation](https://github.com/websockets/ws/blob/master/doc/ws.md#new-websocketserveroptions-callback) and [`websocket-stream` options documentation](https://github.com/maxogden/websocket-stream#options).
+
+_**NB:** By default if you do not provide a `server` option `fastify-websocket` will bind your websocket server instance to the scoped `fastify` instance._
+
 ## TODO
 
 * [ ] Support Hooks?

index.js

@@ -3,21 +3,28 @@
 const fp = require('fastify-plugin')
 const websocket = require('websocket-stream')
 
-module.exports = fp(function (fastify, opts, next) {
-  var handle = opts.handle
+function fastifyWebsocket (fastify, opts, next) {
+  const handle = opts.handle
+  const options = Object.assign({ server: fastify.server }, opts.options)
 
   if (typeof handle !== 'function') {
     return next(new Error('invalid handle function'))
   }
 
-  var wss = websocket.createServer({
-    server: fastify.server
-  }, handle)
+  const wss = websocket.createServer(options, handle)
 
   fastify.decorate('websocketServer', wss)
 
+  fastify.addHook('onClose', close)
+
   next()
-}, {
+}
+
+function close (fastify, done) {
+  fastify.websocketServer.close(done)
+}
+
+module.exports = fp(fastifyWebsocket, {
   fastify: '>=0.39.0',
   name: 'fastify-websocket'
 })

test.js

@@ -1,34 +1,35 @@
 'use strict'
 
+const http = require('http')
 const test = require('tap').test
-const fastify = require('fastify')
+const Fastify = require('fastify')
 const fastifyWebsocket = require('.')
 const websocket = require('websocket-stream')
 
-test('expose a websocket', (t) => {
+test('Should expose a websocket', (t) => {
   t.plan(3)
 
-  const server = fastify()
-  server.register(fastifyWebsocket, { handle })
+  const fastify = Fastify()
+  t.tearDown(() => fastify.close())
 
-  function handle (conn) {
-    conn.setEncoding('utf8')
-    conn.write('hello client')
-    t.tearDown(conn.destroy.bind(conn))
+  fastify.register(fastifyWebsocket, { handle })
 
-    conn.once('data', (chunk) => {
+  function handle (connection) {
+    connection.setEncoding('utf8')
+    connection.write('hello client')
+    t.tearDown(() => connection.destroy())
+
+    connection.once('data', (chunk) => {
       t.equal(chunk, 'hello server')
-      conn.end()
+      connection.end()
     })
   }
 
-  t.tearDown(server.close.bind(server))
-
-  server.listen(0, (err) => {
+  fastify.listen(0, (err) => {
     t.error(err)
 
-    const client = websocket('ws://localhost:' + server.server.address().port)
-    t.tearDown(client.destroy.bind(client))
+    const client = websocket('ws://localhost:' + fastify.server.address().port)
+    t.tearDown(() => client.destroy())
 
     client.setEncoding('utf8')
     client.write('hello server')
@@ -39,3 +40,104 @@ test('expose a websocket', (t) => {
     })
   })
 })
+
+test('Should be able to pass custom options to websocket-stream', (t) => {
+  t.plan(3)
+
+  const fastify = Fastify()
+  t.tearDown(() => fastify.close())
+
+  const options = {
+    verifyClient: function (info) {
+      t.equal(info.req.headers['x-custom-header'], 'fastify is awesome !')
+
+      return true
+    }
+  }
+
+  fastify.register(fastifyWebsocket, { handle, options })
+
+  // this is all that's needed to create an echo server
+  function handle (connection) {
+    connection.pipe(connection)
+    t.tearDown(() => connection.destroy())
+  }
+
+  fastify.listen(0, (err) => {
+    t.error(err)
+
+    const clientOptions = { headers: { 'x-custom-header': 'fastify is awesome !' } }
+    const client = websocket('ws://localhost:' + fastify.server.address().port, clientOptions)
+    t.tearDown(() => client.destroy())
+
+    client.setEncoding('utf8')
+    client.write('hello')
+
+    client.once('data', (chunk) => {
+      t.equal(chunk, 'hello')
+      client.end()
+    })
+  })
+})
+
+test('Should be able to pass a custom server option to websocket-stream', (t) => {
+  t.plan(2)
+
+  // We create an external server
+  const externalServerPort = 3000
+  const externalServer = http
+    .createServer()
+    .on('connection', (socket) => {
+      socket.unref()
+    })
+    .listen(externalServerPort, 'localhost')
+
+  const fastify = Fastify()
+  t.tearDown(() => {
+    externalServer.close()
+    fastify.close()
+  })
+
+  const options = {
+    server: externalServer
+  }
+
+  fastify.register(fastifyWebsocket, { handle, options })
+
+  // this is all that's needed to create an echo server
+  function handle (connection) {
+    connection.pipe(connection)
+    t.tearDown(() => connection.destroy())
+  }
+
+  fastify.listen(0, (err) => {
+    t.error(err)
+
+    const client = websocket('ws://localhost:' + externalServerPort)
+    t.tearDown(() => client.destroy())
+
+    client.setEncoding('utf8')
+    client.write('hello')
+
+    client.once('data', (chunk) => {
+      t.equal(chunk, 'hello')
+      client.end()
+    })
+  })
+})
+
+test('Should throw on an invalid handle parameter', (t) => {
+  t.plan(2)
+
+  const fastify = Fastify()
+  t.tearDown(() => fastify.close())
+
+  const handle = 'handle must be a function'
+
+  fastify.register(fastifyWebsocket, { handle })
+
+  fastify.listen(0, (err) => {
+    t.ok(err)
+    t.equal(err.message, 'invalid handle function')
+  })
+})