Added basic websocket support (#89)

* Added basic websocket support

* Websocket support

* Update index.js

Co-authored-by: James Sumners <[email protected]>

* Fix old nodes

Co-authored-by: James Sumners <[email protected]>

Author: mcollina

README.md

@@ -118,6 +118,20 @@ configuration passed to the route.
 
 Object with [reply options](https://github.com/fastify/fastify-reply-from#replyfromsource-opts) for `fastify-reply-from`.
 
+## websocket
+
+This module has _partial_ support for forwarding websockets by passing a
+`websocket` option. All those options are going to be forwarded to
+[`fastify-websocket`](https://github.com/fastify/fastify-websocket).
+A few things are missing:
+
+1. forwarding headers as well as `rewriteHeaders`
+2. support for paths, `prefix` and `rewritePrefix`
+3. request id logging
+
+Pull requests are welcome to finish this feature.
+
+
 ## Benchmarks
 
 The following benchmarks where generated on a Macbook 2018 with i5 and
@@ -137,6 +151,7 @@ URL`.
 
 * [ ] Generate unique request ids and implement request tracking
 * [ ] Perform validations for incoming data
+* [ ] Finish implementing websocket (follow TODO)
 
 ## License
 

index.js

@@ -1,6 +1,10 @@
 'use strict'
 
 const From = require('fastify-reply-from')
+const WebSocketPlugin = require('fastify-websocket')
+const WebSocket = require('ws')
+const { pipeline } = require('stream')
+const nonWsMethods = ['DELETE', 'HEAD', 'PATCH', 'POST', 'PUT', 'OPTIONS']
 
 module.exports = async function (fastify, opts) {
   if (!opts.upstream) {
@@ -42,12 +46,59 @@ module.exports = async function (fastify, opts) {
     done(null, req)
   }
 
-  fastify.all('/', { preHandler, config: opts.config || {} }, reply)
-  fastify.all('/*', { preHandler, config: opts.config || {} }, reply)
+  if (opts.websocket) {
+    fastify.register(WebSocketPlugin, opts.websocket)
+  }
+
+  fastify.get('/', {
+    preHandler,
+    config: opts.config || {},
+    handler,
+    wsHandler
+  })
+  fastify.get('/*', {
+    preHandler,
+    config: opts.config || {},
+    handler,
+    wsHandler
+  })
+
+  fastify.route({
+    url: '/',
+    method: nonWsMethods,
+    preHandler,
+    config: opts.config || {},
+    handler
+  })
+  fastify.route({
+    url: '/*',
+    method: nonWsMethods,
+    preHandler,
+    config: opts.config || {},
+    handler
+  })
 
-  function reply (request, reply) {
+  function handler (request, reply) {
     var dest = request.req.url
     dest = dest.replace(this.prefix, rewritePrefix)
     reply.from(dest || '/', replyOpts)
   }
+
+  function wsHandler (conn, req) {
+    // TODO support paths and querystrings
+    // TODO support rewriteHeader
+    // TODO support rewritePrefix
+    const ws = new WebSocket(opts.upstream)
+    const stream = WebSocket.createWebSocketStream(ws)
+
+    // TODO fastify-websocket should create a logger for each connection
+    fastify.log.info('starting websocket tunnel')
+    pipeline(conn, stream, conn, function (err) {
+      if (err) {
+        fastify.log.info({ err }, 'websocket tunnel terminated with error')
+        return
+      }
+      fastify.log.info('websocket tunnel terminated')
+    })
+  }
 }

package.json

@@ -4,7 +4,7 @@
   "description": "proxy http requests, for Fastify",
   "main": "index.js",
   "scripts": {
-    "test": "standard | snazzy && tap test/test.js",
+    "test": "standard | snazzy && tap test/*.js",
     "lint:typescript": "standard --fix --parser @typescript-eslint/parser --plugin typescript test/types/*.ts",
     "typescript": "tsc --project ./test/types/tsconfig.json"
   },
@@ -42,6 +42,8 @@
     "typescript": "^3.4.5"
   },
   "dependencies": {
-    "fastify-reply-from": "^2.0.0"
+    "fastify-reply-from": "^2.0.0",
+    "fastify-websocket": "^1.1.2",
+    "ws": "^7.2.5"
   }
 }

test/websocket.js

@@ -0,0 +1,49 @@
+'use strict'
+
+const { test } = require('tap')
+const Fastify = require('fastify')
+const proxy = require('../')
+const WebSocket = require('ws')
+const { createServer } = require('http')
+const { promisify } = require('util')
+const { once } = require('events')
+
+test('basic websocket proxy', async (t) => {
+  t.plan(2)
+
+  const origin = createServer()
+  const wss = new WebSocket.Server({ server: origin })
+  t.tearDown(wss.close.bind(wss))
+  t.tearDown(origin.close.bind(origin))
+
+  wss.on('connection', (ws) => {
+    ws.on('message', (message) => {
+      t.equal(message.toString(), 'hello')
+      // echo
+      ws.send(message)
+    })
+  })
+
+  await promisify(origin.listen.bind(origin))(0)
+
+  const server = Fastify()
+  server.register(proxy, {
+    upstream: `http://localhost:${origin.address().port}`,
+    websocket: true
+  })
+
+  await server.listen(0)
+  t.tearDown(server.close.bind(server))
+
+  const ws = new WebSocket(`http://localhost:${server.server.address().port}`)
+
+  await once(ws, 'open')
+
+  const stream = WebSocket.createWebSocketStream(ws)
+
+  stream.write('hello')
+
+  const [buf] = await once(stream, 'data')
+
+  t.is(buf.toString(), 'hello')
+})