feat: The target websocket watch the specific url instead of upstream… (#304)

* feat: The target websocket watch the specific url instead of upstream(host) with path

* feat: The target websocket watch the specific url instead of upstream(host) with path, add test for type

* feat: The target websocket watch the specific url instead of upstream(host) with path, more strict types

---------

Co-authored-by: Gulin Anton <[email protected]>

Author: Dokki

README.md

@@ -187,6 +187,14 @@ A few things are missing:
 
 Pull requests are welcome to finish this feature.
 
+### `wsUpstream`
+Working only if property `websocket` is `true`.
+
+An URL (including protocol) that represents the target websockets to use for proxying websockets.
+Accepted both `https://` and `wss://`.
+
+Note that if property `wsUpstream` not specified then proxy will try to connect with the `upstream` property.
+
 ### `wsServerOptions`
 
 The options passed to [`new ws.Server()`](https://github.com/websockets/ws/blob/HEAD/doc/ws.md#class-websocketserver).

index.js

@@ -147,9 +147,19 @@ class WebSocketProxy {
 
   findUpstream (request) {
     const source = new URL(request.url, 'ws://127.0.0.1')
+
     for (const { prefix, rewritePrefix, upstream, wsClientOptions } of this.prefixList) {
+      // If the "upstream" have path then need get the Base of url, otherwise the "target" path will be broken.
+      // Example: upstream is "ws://localhost:22/some/path" and after this code
+      // "target = new URL(source.pathname.replace(prefix, rewritePrefix), upstream)"
+      // The target.pathname will be "some/some/path"
+      const upstreamUrl = new URL(upstream)
+      const upstreamBase = upstreamUrl.pathname && upstreamUrl.pathname !== '/'
+        ? upstreamUrl.href.replace(upstreamUrl.pathname, '')
+        : upstream
+
       if (source.pathname.startsWith(prefix)) {
-        const target = new URL(source.pathname.replace(prefix, rewritePrefix), upstream)
+        const target = new URL(source.pathname.replace(prefix, rewritePrefix), upstreamBase)
         target.search = source.search
         return { target, wsClientOptions }
       }
@@ -195,8 +205,16 @@ function setupWebSocketProxy (fastify, options, rewritePrefix) {
     httpWss.set(fastify.server, wsProxy)
   }
 
-  if (options.upstream !== '') {
-    wsProxy.addUpstream(fastify.prefix, rewritePrefix, options.upstream, options.wsClientOptions)
+  if (
+    (typeof options.wsUpstream === 'string' && options.wsUpstream !== '') ||
+    (typeof options.upstream === 'string' && options.upstream !== '')
+  ) {
+    wsProxy.addUpstream(
+      fastify.prefix,
+      rewritePrefix,
+      options.wsUpstream ? options.wsUpstream : options.upstream,
+      options.wsClientOptions
+    )
     // The else block is validate earlier in the code
   } else {
     wsProxy.findUpstream = function (request) {

test/websocket.js

@@ -407,3 +407,68 @@ test('Should gracefully close when clients attempt to connect after calling clos
   await server.close()
   await p
 })
+
+test('Proxy websocket with custom upstream url', async (t) => {
+  t.plan(5)
+
+  const origin = createServer()
+  const wss = new WebSocket.Server({ server: origin })
+
+  t.teardown(wss.close.bind(wss))
+  t.teardown(origin.close.bind(origin))
+
+  const serverMessages = []
+  wss.on('connection', (ws, request) => {
+    ws.on('message', (message, binary) => {
+      // Also need save request.url for check from what url the message is coming.
+      serverMessages.push([message.toString(), binary, request.url])
+      ws.send(message, { binary })
+    })
+  })
+
+  await promisify(origin.listen.bind(origin))({ port: 0 })
+  // Path for wsUpstream and for later check.
+  const path = '/some/path'
+  const server = Fastify()
+  server.register(proxy, {
+    upstream: `ws://localhost:${origin.address().port}`,
+    // Start proxy with different upstream, added path.
+    wsUpstream: `ws://localhost:${origin.address().port}${path}`,
+    websocket: true
+  })
+
+  await server.listen({ port: 0 })
+  t.teardown(server.close.bind(server))
+
+  // Start websocket with different upstream for connect, added path.
+  const ws = new WebSocket(`ws://localhost:${server.server.address().port}${path}`)
+  await once(ws, 'open')
+
+  const data = [{ message: 'hello', binary: false }, { message: 'fastify', binary: true, isBuffer: true }]
+  const dataLength = data.length
+  let dataIndex = 0
+
+  for (; dataIndex < dataLength; dataIndex++) {
+    const { message: msg, binary, isBuffer } = data[dataIndex]
+    const message = isBuffer
+      ? Buffer.from(msg)
+      : msg
+
+    ws.send(message, { binary })
+
+    const [reply, binaryAnswer] = await once(ws, 'message')
+
+    t.equal(reply.toString(), msg)
+    t.equal(binaryAnswer, binary)
+  }
+  // Also check "path", must be the same.
+  t.strictSame(serverMessages, [
+    ['hello', false, path],
+    ['fastify', true, path]
+  ])
+
+  await Promise.all([
+    once(ws, 'close'),
+    server.close()
+  ])
+})

types/index.d.ts

@@ -9,7 +9,19 @@ import {
 
 import { ClientOptions, ServerOptions } from 'ws';
 
-type FastifyHttpProxy = FastifyPluginCallback<fastifyHttpProxy.FastifyHttpProxyOptions>;
+interface FastifyHttpProxyWebsocketOptionsEnabled {
+  websocket: true;
+  wsUpstream?: string;
+}
+interface FastifyHttpProxyWebsocketOptionsDisabled {
+  websocket?: false | never;
+  wsUpstream?: never;
+}
+
+type FastifyHttpProxy = FastifyPluginCallback<
+  fastifyHttpProxy.FastifyHttpProxyOptions
+  & (FastifyHttpProxyWebsocketOptionsEnabled | FastifyHttpProxyWebsocketOptionsDisabled)
+>;
 
 declare namespace fastifyHttpProxy {
   export interface FastifyHttpProxyOptions extends FastifyReplyFromOptions {
@@ -21,7 +33,6 @@ declare namespace fastifyHttpProxy {
     beforeHandler?: preHandlerHookHandler;
     config?: Object;
     replyOptions?: FastifyReplyFromHooks;
-    websocket?: boolean;
     wsClientOptions?: ClientOptions;
     wsServerOptions?: ServerOptions;
     httpMethods?: string[];

types/index.test-d.ts

@@ -42,11 +42,28 @@ app.register(fastifyHttpProxy, {
       timeout: 20000
     }
   },
-  constraints: { version: '1.0.2' }
+  constraints: { version: '1.0.2' },
+  websocket: true,
+  wsUpstream: 'ws://origin.asd/connection'
 });
 
 expectError(
   app.register(fastifyHttpProxy, {
     thisOptionDoesNotExist: 'triggers a typescript error'
   })
 );
+
+expectError(
+  app.register(fastifyHttpProxy, {
+    upstream: 'http://origin.asd',
+    wsUpstream: 'ws://origin.asd'
+  })
+);
+
+expectError(
+  app.register(fastifyHttpProxy, {
+    upstream: 'http://origin.asd',
+    websocket: false,
+    wsUpstream: 'asdf'
+  })
+);