Ensure that setErrorHandler functions run for errors thrown before th…

…e websocket is upgraded (#216)

README.md

@@ -192,11 +192,11 @@ fastify.listen(3000, err => {
 
 ### Custom error handler:
 
-You can optionally provide a custom errorHandler that will be used to handle any cleaning up:
+You can optionally provide a custom `errorHandler` that will be used to handle any cleaning up of established websocket connections. The `errorHandler` will be called if any errors are thrown by your websocket route handler after the connection has been established. Note that neither Fastify's `onError` hook or functions registered with `fastify.setErrorHandler` will be called for errors thrown during a websocket request handler.
 
-```js
-'use strict'
+Neither the `errorHandler` passed to this plugin or fastify's `onError` hook will be called for errors encountered during message processing for your connection. If you want to handle unexpected errors within your `message` event handlers, you'll need to use your own `try { } catch {}` statements and decide what to send back over the websocket.
 
+```js
 const fastify = require('fastify')()
 
 fastify.register(require('@fastify/websocket'), {
@@ -231,6 +231,7 @@ fastify.listen(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.
 ## Options
 
 `@fastify/websocket` accept these options for [`ws`](https://github.com/websockets/ws/blob/master/doc/ws.md#new-websocketserveroptions-callback) :

index.js

@@ -80,19 +80,6 @@ function fastifyWebsocket (fastify, opts, next) {
     done()
   })
 
-  fastify.addHook('onError', (request, reply, error, done) => {
-    /* istanbul ignore next */
-    if (request.raw[kWs]) {
-      // Hijack reply to prevent fastify from sending the error after onError hooks are done running
-      reply.hijack()
-      handleUpgrade(request.raw, connection => {
-        // Handle the error
-        errorHandler.call(this, error, connection, request, reply)
-      })
-    }
-    done()
-  })
-
   fastify.addHook('onResponse', (request, reply, done) => {
     if (request.ws) {
       request.raw[kWs].destroy()

test/hooks.js

@@ -93,16 +93,14 @@ test('Should run onError hook before handler is executed (error thrown in onRequ
     fastify.addHook('onError', async (request, reply) => t.ok('called', 'onError'))
 
     fastify.get('/echo', { websocket: true }, (conn, request) => {
-      t.teardown(conn.destroy.bind(conn))
+      t.fail()
     })
   })
 
   fastify.listen({ port: 0 }, function (err) {
     t.error(err)
     const ws = new WebSocket('ws://localhost:' + (fastify.server.address()).port + '/echo')
-    const client = WebSocket.createWebSocketStream(ws, { encoding: 'utf8' })
-    t.teardown(client.destroy.bind(client))
-    ws.on('close', code => t.equal(code, 1006))
+    ws.on('unexpected-response', (_request, response) => t.equal(response.statusCode, 500))
   })
 })
 
@@ -130,9 +128,7 @@ test('Should run onError hook before handler is executed (error thrown in preVal
   fastify.listen({ port: 0 }, function (err) {
     t.error(err)
     const ws = new WebSocket('ws://localhost:' + (fastify.server.address()).port + '/echo')
-    const client = WebSocket.createWebSocketStream(ws, { encoding: 'utf8' })
-    t.teardown(client.destroy.bind(client))
-    ws.on('close', code => t.equal(code, 1006))
+    ws.on('unexpected-response', (_request, response) => t.equal(response.statusCode, 500))
   })
 })
 
@@ -152,7 +148,7 @@ test('onError hooks can send a reply and prevent hijacking', t => {
 
     fastify.addHook('onError', async (request, reply) => {
       t.ok('called', 'onError')
-      await reply.code(404).send('there was an error')
+      await reply.code(501).send('there was an error')
     })
 
     fastify.get('/echo', { websocket: true }, (conn, request) => {
@@ -163,9 +159,39 @@ test('onError hooks can send a reply and prevent hijacking', t => {
   fastify.listen({ port: 0 }, function (err) {
     t.error(err)
     const ws = new WebSocket('ws://localhost:' + (fastify.server.address()).port + '/echo')
-    const client = WebSocket.createWebSocketStream(ws, { encoding: 'utf8' })
-    t.teardown(client.destroy.bind(client))
-    ws.on('close', code => t.equal(code, 1006))
+    ws.on('unexpected-response', (_request, response) => t.equal(response.statusCode, 501))
+  })
+})
+
+test('setErrorHandler functions can send a reply and prevent hijacking', t => {
+  t.plan(4)
+  const fastify = Fastify()
+
+  t.teardown(() => fastify.close())
+
+  fastify.register(fastifyWebsocket)
+
+  fastify.register(async function (fastify) {
+    fastify.addHook('preValidation', async (request, reply) => {
+      await Promise.resolve()
+      throw new Error('Fail')
+    })
+
+    fastify.setErrorHandler(async (error, request, reply) => {
+      t.ok('called', 'onError')
+      t.ok(error)
+      await reply.code(501).send('there was an error')
+    })
+
+    fastify.get('/echo', { websocket: true }, (conn, request) => {
+      t.fail()
+    })
+  })
+
+  fastify.listen({ port: 0 }, function (err) {
+    t.error(err)
+    const ws = new WebSocket('ws://localhost:' + (fastify.server.address()).port + '/echo')
+    ws.on('unexpected-response', (_request, response) => t.equal(response.statusCode, 501))
   })
 })