diff --git a/index.js b/index.js index 7286240..6b9e615 100644 --- a/index.js +++ b/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 diff --git a/package.json b/package.json index a7feb51..66ef220 100644 --- a/package.json +++ b/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", diff --git a/test/base.test.js b/test/base.test.js index 90d39b9..1d5244f 100644 --- a/test/base.test.js +++ b/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') +})