adding websocket tests

Author: jkyberneees

lib/default-hooks.js

@@ -5,6 +5,9 @@ const toArray = require('stream-to-array')
 const TRANSFER_ENCODING_HEADER_NAME = 'transfer-encoding'
 
 module.exports = {
+  websocket: {
+    onOpenNoOp (ws, searchParams) { }
+  },
   lambda: {
     onRequestNoOp (req, res) { },
     onResponse (req, res, response) {

lib/ws-proxy.js

@@ -1,23 +1,23 @@
 'use strict'
 
 const WebSocket = require('faye-websocket')
-
-function onOpenNoOp () {}
+const { onOpenNoOp } = require('./default-hooks').websocket
 
 module.exports = (config) => {
   const { routes, server } = config
 
-  const prefix2route = {}
+  const prefix2route = new Map()
   for (const route of routes) {
-    prefix2route[route.prefix] = route
+    prefix2route.set(route.prefix, route)
   }
 
   server.on('upgrade', async (req, socket, body) => {
     if (WebSocket.isWebSocket(req)) {
       const url = new URL('http://fw' + req.url)
-      const route = prefix2route[url.pathname]
+      const prefix = url.pathname || '/'
 
-      if (route) {
+      if (prefix2route.has(prefix)) {
+        const route = prefix2route.get(prefix)
         const subProtocols = route.subProtocols || []
         route.hooks = route.hooks || {}
         const onOpen = route.hooks.onOpen || onOpenNoOp

test/ws-proxy.test.js

@@ -0,0 +1,138 @@
+'use strict'
+
+const gateway = require('../index')
+const WebSocket = require('faye-websocket')
+const http = require('http')
+
+/* global describe, it */
+const expect = require('chai').expect
+
+describe('ws-proxy', () => {
+  let gw, echoServer
+
+  it('initialize', async () => {
+    echoServer = http.createServer()
+    echoServer.on('upgrade', (request, socket, body) => {
+      if (WebSocket.isWebSocket(request)) {
+        const ws = new WebSocket(request, socket, body)
+
+        ws.on('message', (event) => {
+          ws.send(JSON.stringify({
+            data: event.data,
+            url: request.url
+          }))
+        })
+      }
+    })
+    echoServer.listen(3000)
+
+    gw = gateway({
+      routes: [{
+        proxyType: 'websocket',
+        prefix: '/',
+        target: 'ws://127.0.0.1:3000'
+      }, {
+        proxyType: 'websocket',
+        prefix: '/echo',
+        target: 'ws://127.0.0.1:3000'
+      }, {
+        proxyType: 'websocket',
+        prefix: '/echo-auth',
+        target: 'ws://127.0.0.1:3000',
+        hooks: {
+          onOpen (ws, searchParams) {
+            if (searchParams.get('accessToken') !== '12345') {
+              const err = new Error('Unauthorized')
+              err.closeEventCode = 4401
+
+              throw err
+            }
+          }
+        }
+      }, {
+        proxyType: 'websocket',
+        prefix: '/echo-params',
+        target: 'ws://127.0.0.1:3000',
+        hooks: {
+          onOpen (ws, searchParams) {
+            searchParams.set('x-token', 'abc')
+          }
+        }
+      }]
+    })
+
+    await gw.start(8080)
+  })
+
+  it('should echo using default prefix', (done) => {
+    const ws = new WebSocket.Client('ws://127.0.0.1:8080')
+    const msg = 'hello'
+
+    ws.on('message', (event) => {
+      const { data } = JSON.parse(event.data)
+      expect(data).equals('hello')
+
+      ws.close()
+      done()
+    })
+
+    ws.send(msg)
+  })
+
+  it('should echo', (done) => {
+    const ws = new WebSocket.Client('ws://127.0.0.1:8080/echo')
+    const msg = 'hello'
+
+    ws.on('message', (event) => {
+      const { data } = JSON.parse(event.data)
+      expect(data).equals('hello')
+
+      ws.close()
+      done()
+    })
+
+    ws.send(msg)
+  })
+
+  it('should fail auth', (done) => {
+    const ws = new WebSocket.Client('ws://127.0.0.1:8080/echo-auth?accessToken=2')
+    ws.on('close', (event) => {
+      done()
+    })
+  })
+
+  it('should pass auth', (done) => {
+    const ws = new WebSocket.Client('ws://127.0.0.1:8080/echo-auth?accessToken=12345')
+    const msg = 'hello'
+
+    ws.on('message', (event) => {
+      const { data } = JSON.parse(event.data)
+      expect(data).equals('hello')
+
+      ws.close()
+      done()
+    })
+
+    ws.send(msg)
+  })
+
+  it('should rewrite search params', (done) => {
+    const ws = new WebSocket.Client('ws://127.0.0.1:8080/echo-params')
+    const msg = 'hello'
+
+    ws.on('message', (event) => {
+      const { url } = JSON.parse(event.data)
+      expect(url).contains('?x-token=abc')
+
+      ws.close()
+      done()
+    })
+
+    ws.send(msg)
+  })
+
+  it('shutdown', async () => {
+    await gw.close()
+    echoServer.close()
+  })
+})