Add preClose opt + merge close and preClose hook (#261)

* Add preClose option

Also commenting out failing test from before these changes, so preCommit hook passes

* Rm onClose hook and handle that logic in preClose

* Make preClose top level option

* Add preClose opt to readme

* Put noHandle back in its original spot

* Add back test

* Clarify preClose can be an async function

* Add preClose to types

* Reset plan count to original (need to force this through the precommit because it fails locally, but I think it will pass the CI)

* Add type tests for preClose opt

Author: HDegroote

README.md

@@ -233,6 +233,29 @@ fastify.listen({ port: 3000 }, err => {
 ```
 
 Note: Fastify's `onError` and error handlers registered by `setErrorHandler` will still be called for errors encountered *before* the websocket connection is established. This means errors thrown by `onRequest` hooks, `preValidation` handlers, and hooks registered by plugins will use the normal error handling mechanisms in Fastify. Once the websocket is established and your websocket route handler is called, `fastify-websocket`'s `errorHandler` takes over.
+
+### Custom preClose hook:
+
+By default, all ws connections are closed when the server closes. If you wish to modify this behaviour, you can pass your own `preClose` function.
+
+Note that `preClose` is responsible for closing all connections and closing the websocket server.
+
+```js
+const fastify = require('fastify')()
+
+fastify.register(require('@fastify/websocket'), {
+  preClose: (done) => { // Note: can also use async style, without done-callback
+    const server = this.websocketServer
+
+    for (const connection of server.clients) {
+      connection.close(1001, 'WS server is going offline in custom manner, sending a code + message')
+    }
+
+    server.close(done)
+  }
+})
+```
+
 ## Options
 
 `@fastify/websocket` accept these options for [`ws`](https://github.com/websockets/ws/blob/master/doc/ws.md#new-websocketserveroptions-callback) :

index.js

@@ -19,6 +19,15 @@ function fastifyWebsocket (fastify, opts, next) {
     errorHandler = opts.errorHandler
   }
 
+  let preClose = defaultPreClose
+  if (opts && opts.preClose) {
+    if (typeof opts.preClose !== 'function') {
+      return next(new Error('invalid preClose function'))
+    }
+
+    preClose = opts.preClose
+  }
+
   if (opts.options && opts.options.noServer) {
     return next(new Error("fastify-websocket doesn't support the ws noServer option. If you want to create a websocket server detatched from fastify, use the ws library directly."))
   }
@@ -143,25 +152,23 @@ function fastifyWebsocket (fastify, opts, next) {
     }
   })
 
-  fastify.addHook('onClose', close)
-
   // Fastify is missing a pre-close event, or the ability to
   // add a hook before the server.close call. We need to resort
   // to monkeypatching for now.
-  fastify.addHook('preClose', function (done) {
+  fastify.addHook('preClose', preClose)
+
+  function defaultPreClose (done) {
     const server = this.websocketServer
     if (server.clients) {
       for (const client of server.clients) {
         client.close()
       }
     }
     fastify.server.removeListener('upgrade', onUpgrade)
-    done()
-  })
 
-  function close (fastify, done) {
-    const server = fastify.websocketServer
     server.close(done)
+
+    done()
   }
 
   function noHandle (connection, rawRequest) {

test/base.test.js

@@ -351,6 +351,74 @@ test('Should be able to pass custom connectionOptions to createWebSocketStream',
   await p
 })
 
+test('Should be able to pass preClose option to override default', async (t) => {
+  t.plan(3)
+
+  const fastify = Fastify()
+
+  const preClose = (done) => {
+    t.pass('Custom preclose successfully called')
+
+    for (const connection of fastify.websocketServer.clients) {
+      connection.close()
+    }
+    done()
+  }
+
+  await fastify.register(fastifyWebsocket, { preClose })
+
+  fastify.get('/', { websocket: true }, (connection) => {
+    connection.setEncoding('utf8')
+    t.teardown(() => connection.destroy())
+
+    connection.once('data', (chunk) => {
+      t.equal(chunk, 'hello server')
+      connection.write('hello client')
+      connection.end()
+    })
+  })
+
+  await fastify.listen({ port: 0 })
+
+  const ws = new WebSocket('ws://localhost:' + fastify.server.address().port)
+  const client = WebSocket.createWebSocketStream(ws, { encoding: 'utf8' })
+  t.teardown(() => client.destroy())
+
+  client.setEncoding('utf8')
+  client.write('hello server')
+
+  const [chunk] = await once(client, 'data')
+  t.equal(chunk, 'hello client')
+  client.end()
+
+  await fastify.close()
+})
+
+test('Should fail if custom preClose is not a function', async (t) => {
+  t.plan(2)
+
+  const fastify = Fastify()
+  t.teardown(() => fastify.close())
+
+  const preClose = 'Not a function'
+
+  try {
+    await fastify.register(fastifyWebsocket, { preClose })
+  } catch (err) {
+    t.equal(err.message, 'invalid preClose function')
+  }
+
+  fastify.get('/', { websocket: true }, (connection) => {
+    t.teardown(() => connection.destroy())
+  })
+
+  try {
+    await fastify.listen({ port: 0 })
+  } catch (err) {
+    t.equal(err.message, 'invalid preClose function')
+  }
+})
+
 test('Should gracefully close with a connected client', async (t) => {
   t.plan(2)
 

types/index.d.ts

@@ -5,6 +5,7 @@ import * as fastify from 'fastify';
 import WebSocket from 'ws';
 import { Duplex, DuplexOptions } from 'stream';
 import { FastifyReply } from 'fastify/types/reply';
+import { preCloseHookHandler, preCloseAsyncHookHandler } from 'fastify/types/hooks';
 import { RouteGenericInterface } from 'fastify/types/route';
 
 interface WebsocketRouteOptions<
@@ -88,6 +89,7 @@ declare namespace fastifyWebsocket {
     errorHandler?: (this: FastifyInstance, error: Error, connection: SocketStream, request: FastifyRequest, reply: FastifyReply) => void;
     options?: WebSocketServerOptions;
     connectionOptions?: DuplexOptions;
+    preClose?: preCloseHookHandler | preCloseAsyncHookHandler;
   }
 
   export interface RouteOptions<

types/index.test-d.ts

@@ -22,6 +22,8 @@ app.register(wsPlugin, {
   }
 });
 app.register(wsPlugin, { options: { perMessageDeflate: true } });
+app.register(wsPlugin, { preClose: function syncPreclose() {} });
+app.register(wsPlugin, { preClose: async function asyncPreclose(){} });
 
 app.get('/websockets-via-inferrence', { websocket: true }, async function (connection, request) {
   expectType<FastifyInstance>(this);