remove all pre-existing upgrade listeners (#253)

* remove all pre-existing upgrade listeners

Signed-off-by: Matteo Collina <[email protected]>

* rewrite

Signed-off-by: Matteo Collina <[email protected]>

* fixup

Signed-off-by: Matteo Collina <[email protected]>

* fixup

Signed-off-by: Matteo Collina <[email protected]>

---------

Signed-off-by: Matteo Collina <[email protected]>

Author: mcollina

index.js

@@ -38,21 +38,19 @@ function fastifyWebsocket (fastify, opts, next) {
   const wss = new WebSocket.Server(wssOptions)
   fastify.decorate('websocketServer', wss)
 
-  websocketListenServer.on('upgrade', (rawRequest, socket, head) => {
+  function onUpgrade (rawRequest, socket, head) {
     // Save a reference to the socket and then dispatch the request through the normal fastify router so that it will invoke hooks and then eventually a route handler that might upgrade the socket.
     rawRequest[kWs] = socket
     rawRequest[kWsHead] = head
-
-    if (closing) {
-      handleUpgrade(rawRequest, (connection) => {
-        connection.socket.close(1001)
-      })
-    } else {
-      const rawResponse = new ServerResponse(rawRequest)
+    const rawResponse = new ServerResponse(rawRequest)
+    try {
       rawResponse.assignSocket(socket)
       fastify.routing(rawRequest, rawResponse)
+    } catch (err) {
+      fastify.log.warn({ err }, 'websocket upgrade failed')
     }
-  })
+  }
+  websocketListenServer.on('upgrade', onUpgrade)
 
   const handleUpgrade = (rawRequest, callback) => {
     wss.handleUpgrade(rawRequest, rawRequest[kWs], rawRequest[kWsHead], (socket) => {
@@ -147,24 +145,23 @@ function fastifyWebsocket (fastify, opts, next) {
 
   fastify.addHook('onClose', close)
 
-  let closing = false
-
   // 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) {
-    closing = true
-
-    // Call oldClose first so that we stop listening. This ensures the
-    // server.clients list will be up to date when we start closing below.
-    oldClose.call(this, cb)
+  fastify.addHook('preClose', function (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
-    if (!server.clients) return
-    for (const client of server.clients) {
-      client.close()
-    }
+    server.close(done)
   }
 
   function noHandle (connection, rawRequest) {
@@ -185,13 +182,8 @@ function fastifyWebsocket (fastify, opts, next) {
   next()
 }
 
-function close (fastify, done) {
-  const server = fastify.websocketServer
-  server.close(done)
-}
-
 module.exports = fp(fastifyWebsocket, {
-  fastify: '>= 4.0.0',
+  fastify: '^4.16.0',
   name: '@fastify/websocket'
 })
 module.exports.default = fastifyWebsocket

package.json

@@ -30,7 +30,7 @@
     "@sinclair/typebox": "^0.25.10",
     "@types/node": "^18.15.0",
     "@types/ws": "^8.2.2",
-    "fastify": "^4.0.0",
+    "fastify": "^4.16.0",
     "split2": "^4.1.0",
     "standard": "^17.0.0",
     "tap": "^16.0.0",

test/base.test.js

@@ -398,11 +398,8 @@ test('Should gracefully close when clients attempt to connect after calling clos
   fastify.server.close = function (cb) {
     const ws = new WebSocket('ws://localhost:' + fastify.server.address().port)
 
-    p = once(ws, 'close').then(() => {
-      t.pass('client 2 closed')
-    })
-
-    ws.on('open', () => {
+    p = once(ws, 'close').catch((err) => {
+      t.equal(err.message, 'Unexpected server response: 503')
       oldClose.call(this, cb)
     })
   }
@@ -411,7 +408,7 @@ test('Should gracefully close when clients attempt to connect after calling clos
 
   fastify.get('/', { websocket: true }, (connection) => {
     t.pass('received client connection')
-    t.teardown(() => connection.destroy())
+    connection.destroy()
     // this connection stays alive until we close the server
   })
 
@@ -586,3 +583,38 @@ test('Should Handle WebSocket errors to avoid Node.js crashes', async t => {
 
   await fastify.close()
 })
+
+test('remove all others websocket handlers on close', async (t) => {
+  const fastify = Fastify()
+
+  await fastify.register(fastifyWebsocket)
+
+  await fastify.listen({ port: 0 })
+
+  await fastify.close()
+
+  t.equal(fastify.server.listeners('upgrade').length, 0)
+})
+
+test('clashing upgrade handler', async (t) => {
+  const fastify = Fastify()
+  t.teardown(() => fastify.close())
+
+  fastify.server.on('upgrade', (req, socket, head) => {
+    const res = new http.ServerResponse(req)
+    res.assignSocket(socket)
+    res.end()
+    socket.destroy()
+  })
+
+  await fastify.register(fastifyWebsocket)
+
+  fastify.get('/', { websocket: true }, (connection) => {
+    t.fail('this should never be invoked')
+  })
+
+  await fastify.listen({ port: 0 })
+
+  const ws = new WebSocket('ws://localhost:' + fastify.server.address().port)
+  await once(ws, 'error')
+})