Rewrite how middlewares are ran inside sockets, and fix collaboration

packages/server/src/websockets/grid.ts

@@ -1,12 +1,15 @@
 import authorized from "../middleware/authorized"
+import currentApp from "../middleware/currentapp"
 import { BaseSocket } from "./websocket"
-import { context, permissions } from "@budibase/backend-core"
+import { auth, permissions } from "@budibase/backend-core"
 import http from "http"
 import Koa from "koa"
 import { getTableId } from "../api/controllers/row/utils"
 import { Row, Table } from "@budibase/types"
 import { Socket } from "socket.io"
 import { GridSocketEvent } from "@budibase/shared-core"
+import { userAgent } from "koa-useragent"
+import { createContext, runMiddlewares } from "./middleware"
 
 const { PermissionType, PermissionLevel } = permissions
 
@@ -26,28 +29,27 @@ export default class GridSocket extends BaseSocket {
           return
         }
 
-        // Check if the user has permission to read this resource
+        // Create context
-        const middleware = authorized(
+        const ctx = createContext(this.app, socket, {
-          PermissionType.TABLE,
-          PermissionLevel.READ
-        )
-        const ctx = {
-          appId,
           resourceId: tableId,
-          roleId: socket.data.roleId,
+          appId,
-          user: { _id: socket.data._id },
+        })
-          isAuthenticated: socket.data.isAuthenticated,
+
-          request: {
+        // Construct full middleware chain to assess permissions
-            url: "/fake",
+        const middlewares = [
-          },
+          userAgent,
-          get: () => null,
+          auth.buildAuthMiddleware([], {
-          throw: () => {
+            publicAllowed: true,
-            // If they don't have access, immediately disconnect them
+          }),
-            socket.disconnect(true)
+          currentApp,
-          },
+          authorized(PermissionType.TABLE, PermissionLevel.READ),
-        }
+        ]
-        await context.doInAppContext(appId, async () => {
+
-          await middleware(ctx, async () => {
+        // Run all koa middlewares
+        try {
+          await runMiddlewares(ctx, middlewares, async () => {
+            // Middlewares are finished and we have permission
+            // Join room for this resource
             const room = `${appId}-${tableId}`
             await this.joinRoom(socket, room)
 
@@ -55,7 +57,9 @@ export default class GridSocket extends BaseSocket {
             const sessions = await this.getRoomSessions(room)
             callback({ users: sessions })
           })
-        })
+        } catch (error) {
+          socket.disconnect(true)
+        }
       }
     )
 

packages/server/src/websockets/middleware.ts

@@ -0,0 +1,85 @@
+import { Socket } from "socket.io"
+import Cookies from "cookies"
+import http from "http"
+import Koa from "koa"
+import { Header } from "@budibase/backend-core"
+
+/**
+ * Constructs a fake Koa context to use for manually running middlewares in
+ * sockets
+ * @param app the Koa app
+ * @param socket the socket.io socket instance
+ * @param options additional metadata to populate the context with
+ */
+export const createContext = (
+  app: Koa,
+  socket: Socket,
+  options?: WebsocketContextOptions
+) => {
+  const res = new http.ServerResponse(socket.request)
+  const context: WebsocketContext = {
+    ...app.createContext(socket.request, res),
+
+    // Additional overrides needed to make our middlewares work with this
+    // fake koa context
+    resourceId: options?.resourceId,
+    path: "/fake",
+    request: {
+      url: "/fake",
+      headers: {
+        [Header.APP_ID]: options?.appId,
+      },
+    },
+    cookies: new Cookies(socket.request, res),
+    get: (field: string) => socket.request.headers?.[field] as string,
+    throw: (...params: any[]) => {
+      // Throw has a bunch of different signatures, so we'll just stringify
+      // whatever params we get given
+      throw new Error(
+        ...(params?.join(" ") || "Unknown error in socket middleware")
+      )
+    },
+
+    // Needed for koa-useragent middleware
+    headers: socket.request.headers,
+    header: socket.request.headers,
+  }
+  return context
+}
+
+/**
+ * Runs a list of middlewares, nesting each callback inside each other mimic
+ * how the real middlewares run and ensuring that app or tenant context work
+ * as expected
+ * @param ctx the Koa context
+ * @param middlewares the array of middlewares to run
+ * @param callback a final callback for when all middlewares are completed
+ */
+export const runMiddlewares = async (
+  ctx: any,
+  middlewares: any[],
+  callback: Function
+) => {
+  if (!middlewares[0]) {
+    await callback()
+  } else {
+    await middlewares[0](ctx, async () => {
+      await runMiddlewares(ctx, middlewares.slice(1), callback)
+    })
+  }
+}
+
+export interface WebsocketContext extends Omit<Koa.Context, "request"> {
+  request: {
+    url: string
+    headers: {
+      [key: string]: string | undefined
+    }
+  }
+  cookies: Cookies
+}
+
+export interface WebsocketContextOptions {
+  appId?: string
+  resourceId?: string
+}

packages/server/src/websockets/websocket.ts

@@ -1,7 +1,6 @@
 import { Server } from "socket.io"
 import http from "http"
 import Koa from "koa"
-import Cookies from "cookies"
 import { userAgent } from "koa-useragent"
 import { auth, Header, redis } from "@budibase/backend-core"
 import { createAdapter } from "@socket.io/redis-adapter"
@@ -10,14 +9,17 @@ import { getSocketPubSubClients } from "../utilities/redis"
 import { SocketEvent, SocketSessionTTL } from "@budibase/shared-core"
 import { SocketSession } from "@budibase/types"
 import { v4 as uuid } from "uuid"
+import { createContext, runMiddlewares } from "./middleware"
 
 const anonUser = () => ({
   _id: uuid(),
   email: "user@mail.com",
   firstName: "Anonymous",
+  tenantId: "default",
 })
 
 export class BaseSocket {
+  app: Koa
   io: Server
   path: string
   redisClient?: redis.Client
@@ -28,6 +30,7 @@ export class BaseSocket {
     path: string = "/",
     additionalMiddlewares?: any[]
   ) {
+    this.app = app
     this.path = path
     this.io = new Server(server, {
       path,
@@ -45,52 +48,25 @@ export class BaseSocket {
 
     // Apply middlewares
     this.io.use(async (socket, next) => {
-      // Build fake koa context
+      const ctx = createContext(this.app, socket)
-      const res = new http.ServerResponse(socket.request)
-      const ctx: any = {
-        ...app.createContext(socket.request, res),
 
-        // Additional overrides needed to make our middlewares work with this
-        // fake koa context
-        cookies: new Cookies(socket.request, res),
-        get: (field: string) => socket.request.headers[field],
-        throw: (code: number, message: string) => {
-          throw new Error(message)
-        },
-
-        // Needed for koa-useragent middleware
-        headers: socket.request.headers,
-        header: socket.request.headers,
-
-        // We don't really care about the path since it will never contain
-        // an app ID
-        path: "/socket",
-      }
-
-      // Run all koa middlewares
       try {
-        for (let [idx, middleware] of middlewares.entries()) {
+        await runMiddlewares(ctx, middlewares, () => {
-          await middleware(ctx, () => {
+          // Middlewares are finished
-            if (idx === middlewares.length - 1) {
+          // Extract some data from our enriched koa context to persist
-              // Middlewares are finished
+          // as metadata for the socket
-              // Extract some data from our enriched koa context to persist
+          const user = ctx.user?._id ? ctx.user : anonUser()
-              // as metadata for the socket
+          const { _id, email, firstName, lastName } = user
-              const user = ctx.user?._id ? ctx.user : anonUser()
+          socket.data = {
-              const { _id, email, firstName, lastName } = user
+            _id,
-              socket.data = {
+            email,
-                _id,
+            firstName,
-                email,
+            lastName,
-                firstName,
+            sessionId: socket.id,
-                lastName,
+            connectedAt: Date.now(),
-                sessionId: socket.id,
+          }
-                connectedAt: Date.now(),
+          next()
-                isAuthenticated: ctx.isAuthenticated,
+        })
-                roleId: ctx.roleId,
-              }
-              next()
-            }
-          })
-        }
       } catch (error: any) {
         next(error)
       }