From b8bd1b51c11e0d08465a45a5fc07586ba66a56ad Mon Sep 17 00:00:00 2001
From: Andrew Kingston <andrew@kingston.dev>
Date: Mon, 22 Aug 2022 18:24:34 +0100
Subject: [PATCH] Clean up websocket initialisation

---
 packages/client/src/websocket.js              |  9 ++++++---
 packages/server/src/api/controllers/plugin.ts |  4 ++--
 packages/server/src/app.ts                    | 12 +++--------
 packages/server/src/websocket.ts              | 20 +++++++++++++++++++
 4 files changed, 31 insertions(+), 14 deletions(-)
 create mode 100644 packages/server/src/websocket.ts

diff --git a/packages/client/src/websocket.js b/packages/client/src/websocket.js
index 6acb442e7a..8c151e6ab9 100644
--- a/packages/client/src/websocket.js
+++ b/packages/client/src/websocket.js
@@ -4,20 +4,23 @@ import { io } from "socket.io-client"
 
 export const initWebsocket = () => {
   const { inBuilder, location } = get(builderStore)
-  console.log(location)
+
+  // Only connect when we're inside the builder preview, for now
   if (!inBuilder || !location) {
     return
   }
 
-  // Websocket
+  // Initialise connection
   const tls = location.protocol === "https:"
   const proto = tls ? "wss:" : "ws:"
   const host = location.hostname
   const port = location.port || (tls ? 433 : 80)
   console.log(`${proto}//${host}:${port}`)
   const socket = io(`${proto}//${host}:${port}`, {
-    path: "/socket/",
+    path: "/socket/client",
   })
+
+  // Event handlers
   socket.on("plugin-update", data => {
     builderStore.actions.updateUsedPlugin(data.name, data.hash)
   })
diff --git a/packages/server/src/api/controllers/plugin.ts b/packages/server/src/api/controllers/plugin.ts
index 1c9241a53c..b237a70421 100644
--- a/packages/server/src/api/controllers/plugin.ts
+++ b/packages/server/src/api/controllers/plugin.ts
@@ -4,7 +4,7 @@ import { getGlobalDB } from "@budibase/backend-core/tenancy"
 import { generatePluginID, getPluginParams } from "../../db/utils"
 import { uploadDirectory } from "@budibase/backend-core/objectStore"
 import { PluginType, FileType } from "@budibase/types"
-import { io } from "../../app"
+import { ClientAppSocket } from "../../app"
 
 export async function getPlugins(type?: PluginType) {
   const db = getGlobalDB()
@@ -90,7 +90,7 @@ export async function processPlugin(plugin: FileType) {
     jsUrl: `${bucketPath}${jsFileName}`,
   }
   const response = await db.put(doc)
-  io.sockets.emit("plugin-update", { name, hash })
+  ClientAppSocket.emit("plugin-update", { name, hash })
   return {
     ...doc,
     _rev: response.rev,
diff --git a/packages/server/src/app.ts b/packages/server/src/app.ts
index 8d2e1cdea8..3238fe4d26 100644
--- a/packages/server/src/app.ts
+++ b/packages/server/src/app.ts
@@ -18,12 +18,12 @@ const { logAlert } = require("@budibase/backend-core/logging")
 const { pinoSettings } = require("@budibase/backend-core")
 const { Thread } = require("./threads")
 const fs = require("fs")
-const SocketIO = require("socket.io")
 import redis from "./utilities/redis"
 import * as migrations from "./migrations"
 import { events, installation, tenancy } from "@budibase/backend-core"
 import { createAdminUser, getChecklist } from "./utilities/workerRequests"
 import { watch } from "./watch"
+import { Websocket } from "./websocket"
 
 const app = new Koa()
 
@@ -69,14 +69,8 @@ if (env.isProd()) {
 const server = http.createServer(app.callback())
 destroyable(server)
 
-// Websocket
-export const io = SocketIO(server, {
-  path: "/socket/",
-  cors: {
-    origin: ["https://hmr.lan.kingston.dev"],
-    methods: ["GET", "POST"],
-  },
-})
+// initialise websockets
+export const ClientAppSocket = new Websocket(server, "/socket/client")
 
 let shuttingDown = false,
   errCode = 0
diff --git a/packages/server/src/websocket.ts b/packages/server/src/websocket.ts
new file mode 100644
index 0000000000..db61c34358
--- /dev/null
+++ b/packages/server/src/websocket.ts
@@ -0,0 +1,20 @@
+import SocketIO from "socket.io"
+
+export class Websocket {
+  socketIO: any
+
+  constructor(server: any, path: string) {
+    // @ts-ignore
+    this.socketIO = SocketIO(server, {
+      path,
+      cors: {
+        origin: "*",
+      },
+    })
+  }
+
+  // Emit an event to all sockets
+  emit(event: string, payload: any) {
+    this.socketIO.sockets.emit(event, payload)
+  }
+}