feat: add message validation

index.js

@@ -6,6 +6,7 @@ const WebSocket = require('ws')
 
 const kWs = Symbol('ws-socket')
 const kWsHead = Symbol('ws-head')
+const kWebSocketSchema = Symbol('websocket-body-schema')
 
 function fastifyWebsocket (fastify, opts, next) {
   fastify.decorateRequest('ws', null)
@@ -44,7 +45,7 @@ function fastifyWebsocket (fastify, opts, next) {
     rawRequest[kWsHead] = head
 
     if (closing) {
-      handleUpgrade(rawRequest, (connection) => {
+      handleUpgrade(rawRequest, null, (connection) => {
         connection.socket.close(1001)
       })
     } else {
@@ -54,11 +55,13 @@ function fastifyWebsocket (fastify, opts, next) {
     }
   })
 
-  const handleUpgrade = (rawRequest, callback) => {
+  const handleUpgrade = (rawRequest, request, callback) => {
     wss.handleUpgrade(rawRequest, rawRequest[kWs], rawRequest[kWsHead], (socket) => {
       wss.emit('connection', socket, rawRequest)
-
       const connection = WebSocket.createWebSocketStream(socket, opts.connectionOptions)
+      socket.afterDuplex = true
+      socket.validator = request.context[kWebSocketSchema]
+      socket.strict = opts.strictMode ? opts.strictMode : false
       connection.socket = socket
 
       connection.socket.on('newListener', event => {
@@ -84,7 +87,7 @@ function fastifyWebsocket (fastify, opts, 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 => {
+      handleUpgrade(request.raw, request, connection => {
         // Handle the error
         errorHandler.call(this, error, connection, request, reply)
       })
@@ -119,10 +122,12 @@ function fastifyWebsocket (fastify, opts, next) {
     // we always override the route handler so we can close websocket connections to routes to handlers that don't support websocket connections
     // This is not an arrow function to fetch the encapsulated this
     routeOptions.handler = function (request, reply) {
+      request.context[kWebSocketSchema] = request.context.schema ? fastify.validatorCompiler({ schema: request.context.schema.body }) : null
+
       // within the route handler, we check if there has been a connection upgrade by looking at request.raw[kWs]. we need to dispatch the normal HTTP handler if not, and hijack to dispatch the websocket handler if so
       if (request.raw[kWs]) {
         reply.hijack()
-        handleUpgrade(request.raw, connection => {
+        handleUpgrade(request.raw, request, connection => {
           let result
           try {
             if (isWebsocketRoute) {
@@ -200,7 +205,48 @@ function close (fastify, done) {
   server.close(done)
 }
 
+class ValidateWebSocket extends WebSocket {
+  constructor (...args) {
+    super(...args)
+    this.afterDuplex = false
+    this.validator = null
+    this.strict = false
+  }
+
+  wrapValidation (handler) {
+    return (message, isBinary) => {
+      if (isBinary) return handler(message, isBinary)
+      try {
+        if (!this.validator(JSON.parse(message.toString()))) {
+          if (this.strict) {
+            return this.close(1003, 'Unsupported payload')
+          } else {
+            return this.send(JSON.stringify(this.validator.errors))
+          }
+        }
+      } catch (e) {
+        if (this.strict) {
+          return this.close(1003, 'Unsupported payload')
+        } else {
+          return this.send('Unsupported payload')
+        }
+      }
+      return handler(message, isBinary)
+    }
+  }
+
+  on (event, handler, options) {
+    if (!this.afterDuplex) return super.on(event, handler)
+    if (event === 'message') {
+      super.on(event, this.wrapValidation(handler), options)
+    } else {
+      super.on(event, handler, options)
+    }
+  }
+}
+
 module.exports = fp(fastifyWebsocket, {
   fastify: '>= 3.11.0',
   name: 'fastify-websocket'
 })
+module.exports.ValidateWebSocket = ValidateWebSocket