From b56909dc840e589dffa67aff4f84cd3800c51493 Mon Sep 17 00:00:00 2001
From: Matteo Collina <hello@matteocollina.com>
Date: Sun, 26 Jan 2020 10:36:27 +0100
Subject: [PATCH] Send a close to all websocket connections when server closes

---
 index.js     | 15 ++++++++++++++-
 test/base.js | 43 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/index.js b/index.js
index 888d88f..caf9ecd 100644
--- a/index.js
+++ b/index.js
@@ -61,6 +61,18 @@ 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.
+  const oldClose = fastify.server.close
+  fastify.server.close = function (cb) {
+    const server = fastify.websocketServer
+    for (const client of server.clients) {
+      client.close()
+    }
+    oldClose.call(this, cb)
+  }
+
   function handleRouting (connection, request) {
     const response = new ServerResponse(request)
     request[kWs] = WebSocket.createWebSocketStream(connection)
@@ -72,7 +84,8 @@ function fastifyWebsocket (fastify, opts, next) {
 }
 
 function close (fastify, done) {
-  fastify.websocketServer.close(done)
+  const server = fastify.websocketServer
+  server.close(done)
 }
 
 module.exports = fp(fastifyWebsocket, {
diff --git a/test/base.js b/test/base.js
index b282039..238adcf 100644
--- a/test/base.js
+++ b/test/base.js
@@ -144,3 +144,46 @@ test('Should throw on an invalid handle parameter', (t) => {
     t.equal(err.message, 'invalid handle function')
   })
 })
+
+test('Should gracefully close with a connected client', (t) => {
+  t.plan(6)
+
+  const fastify = Fastify()
+
+  fastify.register(fastifyWebsocket, { handle })
+
+  function handle (connection) {
+    connection.setEncoding('utf8')
+    connection.write('hello client')
+
+    connection.once('data', (chunk) => {
+      t.equal(chunk, 'hello server')
+    })
+
+    connection.on('end', () => {
+      t.pass('end emitted on server side')
+    })
+    // this connection stays alive untile we close the server
+  }
+
+  fastify.listen(0, (err) => {
+    t.error(err)
+
+    const ws = new WebSocket('ws://localhost:' + fastify.server.address().port)
+    const client = WebSocket.createWebSocketStream(ws, { encoding: 'utf8' })
+
+    client.setEncoding('utf8')
+    client.write('hello server')
+
+    client.on('end', () => {
+      t.pass('end emitted on client side')
+    })
+
+    client.once('data', (chunk) => {
+      t.equal(chunk, 'hello client')
+      fastify.close(function (err) {
+        t.error(err)
+      })
+    })
+  })
+})