Adds support for routing (#26)

Author: fox1t

.gitignore

@@ -58,3 +58,4 @@ typings/
 .env
 
 package-lock.json
+.vscode

README.md

@@ -1,5 +1,6 @@
 # fastify-websocket
-[![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/)
+
+[![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).
@@ -12,45 +13,146 @@ npm install fastify-websocket --save
 
 ## Usage
 
-All you need to do is to add it to your project with `register` and you are done!
+There are two possible ways of using this plugin: with a **global handler** or with a **per route handler**.
+
+### Global handler
 
-### Example
+All you need to do is to add it to your project with `register` and pass handle function. You are done!
 
 ```js
 'use strict'
 
 const fastify = require('fastify')()
 
-const wssOptions = {
-  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
+fastify.register(require('fastify-websocket'), {
+  handle,
+  options: {
+    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
     }
+  }
+})
+
+function handle (conn) {
+  conn.pipe(conn) // creates an echo server
+}
+
+fastify.listen(3000, err => {
+  if (err) {
+    fastify.log.error(err)
+    process.exit(1)
+  }
+})
+```
 
-    next(true) // the connection is allowed
+### Per route handler
+
+After registring this plugin, you can choose on which routes the WS server will respond. This could be achieved adding `websocket: true` property to `routeOptions` on a fastify's `.get` route. In this case two arguments will be passed to the handler: socket connection and the original `http.IncomingMessage` (instead of the usual fastify's request and reply objects).
+
+```js
+'use strict'
+
+const fastify = Fastify()
+
+fastify.register(require('fastify-websocket'))
+
+fastify.get('/', { websocket: true }, (connection, req) => {
+  connection.socket.on('message', message => {
+    // message === 'hi from client'
+    connection.socket.send('hi from server')
+  })
+})
+
+fastify.listen(3000, err => {
+  if (err) {
+    fastify.log.error(err)
+    process.exit(1)
   }
+})
+```
+
+In this case there won't be any global handler, so it will respond with a 404 error on every unregistered route, closing the incoming upgrade connection requests.
+
+However you can still pass a default global handler, that will be used as default handler.
+
+```js
+'use strict'
+
+const fastify = require('fastify')()
+
+function handle (conn) {
+  conn.pipe(conn) // creates an echo server
 }
 
 fastify.register(require('fastify-websocket'), {
   handle,
-  options: wssOptions
+  options: { maxPayload: 1048576 }
 })
 
+fastify.get('/', { websocket: true }, (connection, req) => {
+  connection.socket.on('message', message => {
+    // message === 'hi from client'
+    connection.socket.send('hi from server')
+  })
+})
+
+fastify.listen(3000, err => {
+  if (err) {
+    fastify.log.error(err)
+    process.exit(1)
+  }
+})
+```
+
+If you need to handle both HTTP requests and incoming socket connections on the same route, you can still do it using the [full declaration syntax](https://www.fastify.io/docs/latest/Routes/#full-declaration), adding a `wsHandler` property.
+
+```js
+'use strict'
+
+const fastify = require('fastify')()
+
 function handle (conn) {
   conn.pipe(conn) // creates an echo server
 }
 
-fastify.listen(3000, (err) => {
+fastify.register(require('fastify-websocket'), {
+  handle,
+  options: { maxPayload: 1048576 }
+})
+
+fastify.route({
+  method: 'GET',
+  url: '/hello',
+  handler: (req, reply) => {
+    // this will handle http requests
+    reply.send({ hello: 'world' })
+  },
+  wsHandler: (conn, req) => {
+    // this will handle websockets connections
+    conn.setEncoding('utf8')
+    conn.write('hello client')
+
+    conn.once('data', chunk => {
+      conn.end()
+    })
+  }
+})
+
+fastify.listen(3000, err => {
   if (err) {
-    fastify.log.error(err);
-    process.exit(1);
+    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.
@@ -70,10 +172,6 @@ For more informations you can check [`ws` options documentation](https://github.
 
 _**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?
-
 ## Acknowledgements
 
 This project is kindly sponsored by [nearForm](http://nearform.com).

index.js

@@ -1,22 +1,71 @@
 'use strict'
 
+const { ServerResponse } = require('http')
 const fp = require('fastify-plugin')
 const websocket = require('websocket-stream')
+const findMyWay = require('find-my-way')
+
+const kWs = Symbol('ws')
 
 function fastifyWebsocket (fastify, opts, next) {
+  if (opts.handle && typeof opts.handle !== 'function') {
+    return next(new Error('invalid handle function'))
+  }
   const handle = opts.handle
+    ? (req, res) => opts.handle(req[kWs], req)
+    : (req, res) => { req[kWs].socket.close() }
+
   const options = Object.assign({ server: fastify.server }, opts.options)
 
-  if (typeof handle !== 'function') {
-    return next(new Error('invalid handle function'))
-  }
+  const router = findMyWay({
+    ignoreTrailingSlash: true,
+    defaultRoute: handle
+  })
 
-  const wss = websocket.createServer(options, handle)
+  const wss = websocket.createServer(options, handleRouting)
 
   fastify.decorate('websocketServer', wss)
 
+  fastify.addHook('onRoute', routeOptions => {
+    if (routeOptions.websocket || routeOptions.wsHandler) {
+      if (routeOptions.method !== 'GET') {
+        throw new Error('websocket handler can only be declared in GET method')
+      }
+
+      let wsHandler = routeOptions.wsHandler
+      let handler = routeOptions.handler
+
+      if (routeOptions.websocket) {
+        wsHandler = routeOptions.handler
+        handler = function (request, reply) {
+          reply.code(404).send()
+        }
+      }
+
+      if (typeof wsHandler !== 'function') {
+        throw new Error('invalid wsHandler function')
+      }
+
+      router.on('GET', routeOptions.path, (req, _) => {
+        const result = wsHandler(req[kWs], req)
+
+        if (result && typeof result.catch === 'function') {
+          result.catch((err) => req[kWs].destroy(err))
+        }
+      })
+
+      routeOptions.handler = handler
+    }
+  })
+
   fastify.addHook('onClose', close)
 
+  function handleRouting (connection, request) {
+    const response = new ServerResponse(request)
+    request[kWs] = connection
+    router.lookup(request, response)
+  }
+
   next()
 }
 
@@ -25,6 +74,6 @@ function close (fastify, done) {
 }
 
 module.exports = fp(fastifyWebsocket, {
-  fastify: '>=0.39.0',
+  fastify: '>=2.4.1',
   name: 'fastify-websocket'
 })

package.json

@@ -4,7 +4,7 @@
   "description": "basic websocket support for fastify",
   "main": "index.js",
   "scripts": {
-    "test": "standard | snazzy && tap test.js"
+    "test": "standard | snazzy && tap \"test/*.js\""
   },
   "repository": {
     "type": "git",
@@ -21,14 +21,15 @@
   },
   "homepage": "https://github.com/fastify/fastify-websocket#readme",
   "devDependencies": {
-    "fastify": "^2.3.0",
+    "fastify": "^2.4.1",
     "pre-commit": "^1.2.2",
     "snazzy": "^8.0.0",
     "standard": "^12.0.1",
     "tap": "^12.7.0"
   },
   "dependencies": {
     "fastify-plugin": "^1.5.0",
+    "find-my-way": "^2.0.1",
     "websocket-stream": "^5.5.0"
   },
   "greenkeeper": {

test.js renamed to test/base.js

@@ -3,7 +3,7 @@
 const http = require('http')
 const test = require('tap').test
 const Fastify = require('fastify')
-const fastifyWebsocket = require('.')
+const fastifyWebsocket = require('../')
 const websocket = require('websocket-stream')
 
 test('Should expose a websocket', (t) => {