Merge branch 'refs/heads/refactor-server'

backend/server/src/ConnectionManager.ts

@@ -0,0 +1,58 @@
+import { Socket } from "socket.io"
+
+class Counter {
+    private count: number = 0
+
+    increment() {
+        this.count++
+    }
+
+    decrement() {
+        this.count = Math.max(0, this.count - 1)
+    }
+
+    getValue(): number {
+        return this.count
+    }
+}
+
+// Owner Connection Management
+export class ConnectionManager {
+    // Counts how many times the owner is connected to a sandbox
+    private ownerConnections: Record<string, Counter> = {}
+    // Stores all sockets connected to a given sandbox
+    private sockets: Record<string, Set<Socket>> = {}
+
+    // Checks if the owner of a sandbox is connected
+    ownerIsConnected(sandboxId: string): boolean {
+        return this.ownerConnections[sandboxId]?.getValue() > 0
+    }
+
+    // Adds a connection for a sandbox
+    addConnectionForSandbox(socket: Socket, sandboxId: string, isOwner: boolean) {
+        this.sockets[sandboxId] ??= new Set()
+        this.sockets[sandboxId].add(socket)
+
+        // If the connection is for the owner, increments the owner connection counter
+        if (isOwner) {
+            this.ownerConnections[sandboxId] ??= new Counter()
+            this.ownerConnections[sandboxId].increment()
+        }
+    }
+
+    // Removes a connection for a sandbox
+    removeConnectionForSandbox(socket: Socket, sandboxId: string, isOwner: boolean) {
+        this.sockets[sandboxId]?.delete(socket)
+
+        // If the connection being removed is for the owner, decrements the owner connection counter
+        if (isOwner) {
+            this.ownerConnections[sandboxId]?.decrement()
+        }
+    }
+
+    // Returns the set of sockets connected to a given sandbox
+    connectionsForSandbox(sandboxId: string): Set<Socket> {
+        return this.sockets[sandboxId] ?? new Set();
+    }
+
+}

backend/server/src/FileManager.ts

@@ -4,12 +4,6 @@ import RemoteFileStorage from "./RemoteFileStorage"
 import { MAX_BODY_SIZE } from "./ratelimit"
 import { TFile, TFileData, TFolder } from "./types"
 
-// Define the structure for sandbox files
-export type SandboxFiles = {
-  files: (TFolder | TFile)[]
-  fileData: TFileData[]
-}
-
 // Convert list of paths to the hierchical file structure used by the editor
 function generateFileStructure(paths: string[]): (TFolder | TFile)[] {
   const root: TFolder = { id: "/", type: "folder", name: "/", children: [] }
@@ -52,20 +46,22 @@ function generateFileStructure(paths: string[]): (TFolder | TFile)[] {
 export class FileManager {
   private sandboxId: string
   private sandbox: Sandbox
-  public sandboxFiles: SandboxFiles
+  public files: (TFolder | TFile)[]
+  public fileData: TFileData[]
   private fileWatchers: WatchHandle[] = []
   private dirName = "/home/user/project"
-  private refreshFileList: (files: SandboxFiles) => void
+  private refreshFileList: ((files: (TFolder | TFile)[]) => void) | null
 
   // Constructor to initialize the FileManager
   constructor(
     sandboxId: string,
     sandbox: Sandbox,
-    refreshFileList: (files: SandboxFiles) => void
+    refreshFileList: ((files: (TFolder | TFile)[]) => void) | null
   ) {
     this.sandboxId = sandboxId
     this.sandbox = sandbox
-    this.sandboxFiles = { files: [], fileData: [] }
+    this.files = []
+    this.fileData = []
     this.refreshFileList = refreshFileList
   }
 
@@ -110,16 +106,16 @@ export class FileManager {
   private async updateFileData(): Promise<TFileData[]> {
     const remotePaths = await RemoteFileStorage.getSandboxPaths(this.sandboxId)
     const localPaths = this.getLocalFileIds(remotePaths)
-    this.sandboxFiles.fileData = await this.generateFileData(localPaths)
+    this.fileData = await this.generateFileData(localPaths)
-    return this.sandboxFiles.fileData
+    return this.fileData
   }
 
   // Update file structure
   private async updateFileStructure(): Promise<(TFolder | TFile)[]> {
     const remotePaths = await RemoteFileStorage.getSandboxPaths(this.sandboxId)
     const localPaths = this.getLocalFileIds(remotePaths)
-    this.sandboxFiles.files = generateFileStructure(localPaths)
+    this.files = generateFileStructure(localPaths)
-    return this.sandboxFiles.files
+    return this.files
   }
 
   // Initialize the FileManager
@@ -130,7 +126,7 @@ export class FileManager {
     await this.updateFileData()
 
     // Copy all files from the project to the container
-    const promises = this.sandboxFiles.fileData.map(async (file) => {
+    const promises = this.fileData.map(async (file) => {
       try {
         const filePath = path.join(this.dirName, file.id)
         const parentDirectory = path.dirname(filePath)
@@ -209,7 +205,7 @@ export class FileManager {
             // Handle file/directory creation event
             if (event.type === "create") {
               const folder = findFolderById(
-                this.sandboxFiles.files,
+                this.files,
                 sandboxDirectory
               ) as TFolder
               const isDir = await this.isDirectory(containerFilePath)
@@ -232,7 +228,7 @@ export class FileManager {
                 folder.children.push(newItem)
               } else {
                 // If folder doesn't exist, add the new item to the root
-                this.sandboxFiles.files.push(newItem)
+                this.files.push(newItem)
               }
 
               if (!isDir) {
@@ -241,7 +237,7 @@ export class FileManager {
                 )
                 const fileContents =
                   typeof fileData === "string" ? fileData : ""
-                this.sandboxFiles.fileData.push({
+                this.fileData.push({
                   id: sandboxFilePath,
                   data: fileContents,
                 })
@@ -253,7 +249,7 @@ export class FileManager {
             // Handle file/directory removal or rename event
             else if (event.type === "remove" || event.type == "rename") {
               const folder = findFolderById(
-                this.sandboxFiles.files,
+                this.files,
                 sandboxDirectory
               ) as TFolder
               const isDir = await this.isDirectory(containerFilePath)
@@ -269,13 +265,13 @@ export class FileManager {
                 )
               } else {
                 // Remove from the root if it's not inside a folder
-                this.sandboxFiles.files = this.sandboxFiles.files.filter(
+                this.files = this.files.filter(
                   (file: TFolder | TFile) => !isFileMatch(file)
                 )
               }
 
               // Also remove any corresponding file data
-              this.sandboxFiles.fileData = this.sandboxFiles.fileData.filter(
+              this.fileData = this.fileData.filter(
                 (file: TFileData) => !isFileMatch(file)
               )
 
@@ -285,10 +281,10 @@ export class FileManager {
             // Handle file write event
             else if (event.type === "write") {
               const folder = findFolderById(
-                this.sandboxFiles.files,
+                this.files,
                 sandboxDirectory
               ) as TFolder
-              const fileToWrite = this.sandboxFiles.fileData.find(
+              const fileToWrite = this.fileData.find(
                 (file) => file.id === sandboxFilePath
               )
 
@@ -308,7 +304,7 @@ export class FileManager {
                   )
                   const fileContents =
                     typeof fileData === "string" ? fileData : ""
-                  this.sandboxFiles.fileData.push({
+                  this.fileData.push({
                     id: sandboxFilePath,
                     data: fileContents,
                   })
@@ -318,7 +314,9 @@ export class FileManager {
             }
 
             // Tell the client to reload the file list
-            this.refreshFileList(this.sandboxFiles)
+            if (event.type !== "chmod") {
+              this.refreshFileList?.(this.files)
+            }
           } catch (error) {
             console.error(
               `Error handling ${event.type} event for ${event.name}:`,
@@ -350,7 +348,7 @@ export class FileManager {
 
   // Get file content
   async getFile(fileId: string): Promise<string | undefined> {
-    const file = this.sandboxFiles.fileData.find((f) => f.id === fileId)
+    const file = this.fileData.find((f) => f.id === fileId)
     return file?.data
   }
 
@@ -368,7 +366,7 @@ export class FileManager {
       throw new Error("File size too large. Please reduce the file size.")
     }
     await RemoteFileStorage.saveFile(this.getRemoteFileId(fileId), body)
-    const file = this.sandboxFiles.fileData.find((f) => f.id === fileId)
+    const file = this.fileData.find((f) => f.id === fileId)
     if (!file) return
     file.data = body
 
@@ -381,9 +379,9 @@ export class FileManager {
     fileId: string,
     folderId: string
   ): Promise<(TFolder | TFile)[]> {
-    const fileData = this.sandboxFiles.fileData.find((f) => f.id === fileId)
+    const fileData = this.fileData.find((f) => f.id === fileId)
-    const file = this.sandboxFiles.files.find((f) => f.id === fileId)
+    const file = this.files.find((f) => f.id === fileId)
-    if (!fileData || !file) return this.sandboxFiles.files
+    if (!fileData || !file) return this.files
 
     const parts = fileId.split("/")
     const newFileId = folderId + "/" + parts.pop()
@@ -427,13 +425,13 @@ export class FileManager {
     await this.sandbox.files.write(path.posix.join(this.dirName, id), "")
     await this.fixPermissions()
 
-    this.sandboxFiles.files.push({
+    this.files.push({
       id,
       name,
       type: "file",
     })
 
-    this.sandboxFiles.fileData.push({
+    this.fileData.push({
       id,
       data: "",
     })
@@ -451,8 +449,8 @@ export class FileManager {
 
   // Rename a file
   async renameFile(fileId: string, newName: string): Promise<void> {
-    const fileData = this.sandboxFiles.fileData.find((f) => f.id === fileId)
+    const fileData = this.fileData.find((f) => f.id === fileId)
-    const file = this.sandboxFiles.files.find((f) => f.id === fileId)
+    const file = this.files.find((f) => f.id === fileId)
     if (!fileData || !file) return
 
     const parts = fileId.split("/")
@@ -468,11 +466,11 @@ export class FileManager {
 
   // Delete a file
   async deleteFile(fileId: string): Promise<(TFolder | TFile)[]> {
-    const file = this.sandboxFiles.fileData.find((f) => f.id === fileId)
+    const file = this.fileData.find((f) => f.id === fileId)
-    if (!file) return this.sandboxFiles.files
+    if (!file) return this.files
 
     await this.sandbox.files.remove(path.posix.join(this.dirName, fileId))
-    this.sandboxFiles.fileData = this.sandboxFiles.fileData.filter(
+    this.fileData = this.fileData.filter(
       (f) => f.id !== fileId
     )
 
@@ -487,7 +485,7 @@ export class FileManager {
     await Promise.all(
       files.map(async (file) => {
         await this.sandbox.files.remove(path.posix.join(this.dirName, file))
-        this.sandboxFiles.fileData = this.sandboxFiles.fileData.filter(
+        this.fileData = this.fileData.filter(
           (f) => f.id !== file
         )
         await RemoteFileStorage.deleteFile(this.getRemoteFileId(file))

backend/server/src/Sandbox.ts

@@ -0,0 +1,243 @@
+import { Sandbox as E2BSandbox } from "e2b"
+import { Socket } from "socket.io"
+import { AIWorker } from "./AIWorker"
+import { CONTAINER_TIMEOUT } from "./constants"
+import { DokkuClient } from "./DokkuClient"
+import { FileManager } from "./FileManager"
+import {
+    createFileRL,
+    createFolderRL,
+    deleteFileRL,
+    renameFileRL,
+    saveFileRL,
+} from "./ratelimit"
+import { SecureGitClient } from "./SecureGitClient"
+import { TerminalManager } from "./TerminalManager"
+import { TFile, TFolder } from "./types"
+import { LockManager } from "./utils"
+
+const lockManager = new LockManager()
+
+// Define a type for SocketHandler functions
+type SocketHandler<T = Record<string, any>> = (args: T) => any;
+
+// Extract port number from a string
+function extractPortNumber(inputString: string): number | null {
+    const cleanedString = inputString.replace(/\x1B\[[0-9;]*m/g, "")
+    const regex = /http:\/\/localhost:(\d+)/
+    const match = cleanedString.match(regex)
+    return match ? parseInt(match[1]) : null
+}
+
+type ServerContext = {
+    aiWorker: AIWorker;
+    dokkuClient: DokkuClient | null;
+    gitClient: SecureGitClient | null;
+};
+
+export class Sandbox {
+    // Sandbox properties:
+    sandboxId: string;
+    fileManager: FileManager | null;
+    terminalManager: TerminalManager | null;
+    container: E2BSandbox | null;
+    // Server context:
+    dokkuClient: DokkuClient | null;
+    gitClient: SecureGitClient | null;
+    aiWorker: AIWorker;
+
+    constructor(sandboxId: string, { aiWorker, dokkuClient, gitClient }: ServerContext) {
+        // Sandbox properties:
+        this.sandboxId = sandboxId;
+        this.fileManager = null;
+        this.terminalManager = null;
+        this.container = null;
+        // Server context:
+        this.aiWorker = aiWorker;
+        this.dokkuClient = dokkuClient;
+        this.gitClient = gitClient;
+    }
+
+    // Initializes the container for the sandbox environment
+    async initialize(
+        fileWatchCallback: ((files: (TFolder | TFile)[]) => void) | undefined
+    ) {
+        // Acquire a lock to ensure exclusive access to the sandbox environment
+        await lockManager.acquireLock(this.sandboxId, async () => {
+            // Check if a container already exists and is running
+            if (this.container && await this.container.isRunning()) {
+                console.log(`Found existing container ${this.sandboxId}`)
+            } else {
+                console.log("Creating container", this.sandboxId)
+                // Create a new container with a specified timeout
+                this.container = await E2BSandbox.create({
+                    timeoutMs: CONTAINER_TIMEOUT,
+                })
+            }
+        })
+        // Ensure a container was successfully created
+        if (!this.container) throw new Error("Failed to create container")
+
+        // Initialize the terminal manager if it hasn't been set up yet
+        if (!this.terminalManager) {
+            this.terminalManager = new TerminalManager(this.container)
+            console.log(`Terminal manager set up for ${this.sandboxId}`)
+        }
+
+        // Initialize the file manager if it hasn't been set up yet
+        if (!this.fileManager) {
+            this.fileManager = new FileManager(
+                this.sandboxId,
+                this.container,
+                fileWatchCallback ?? null
+            )
+            // Initialize the file manager and emit the initial files
+            await this.fileManager.initialize()
+        }
+    }
+
+    // Called when the client disconnects from the Sandbox
+    async disconnect() {
+        // Close all terminals managed by the terminal manager
+        await this.terminalManager?.closeAllTerminals()
+        // This way the terminal manager will be set up again if we reconnect
+        this.terminalManager = null;
+        // Close all file watchers managed by the file manager
+        await this.fileManager?.closeWatchers()
+        // This way the file manager will be set up again if we reconnect
+        this.fileManager = null;
+    }
+
+    handlers(connection: { userId: string, isOwner: boolean, socket: Socket }) {
+
+        // Handle heartbeat from a socket connection
+        const handleHeartbeat: SocketHandler = (_: any) => {
+            // Only keep the sandbox alive if the owner is still connected
+            if (connection.isOwner) {
+                this.container?.setTimeout(CONTAINER_TIMEOUT)
+            }
+        }
+
+        // Handle getting a file
+        const handleGetFile: SocketHandler = ({ fileId }: any) => {
+            return this.fileManager?.getFile(fileId)
+        }
+
+        // Handle getting a folder
+        const handleGetFolder: SocketHandler = ({ folderId }: any) => {
+            return this.fileManager?.getFolder(folderId)
+        }
+
+        // Handle saving a file
+        const handleSaveFile: SocketHandler = async ({ fileId, body }: any) => {
+            await saveFileRL.consume(connection.userId, 1);
+            return this.fileManager?.saveFile(fileId, body)
+        }
+
+        // Handle moving a file
+        const handleMoveFile: SocketHandler = ({ fileId, folderId }: any) => {
+            return this.fileManager?.moveFile(fileId, folderId)
+        }
+
+        // Handle listing apps
+        const handleListApps: SocketHandler = async (_: any) => {
+            if (!this.dokkuClient) throw Error("Failed to retrieve apps list: No Dokku client")
+            return { success: true, apps: await this.dokkuClient.listApps() }
+        }
+
+        // Handle deploying code
+        const handleDeploy: SocketHandler = async (_: any) => {
+            if (!this.gitClient) throw Error("No git client")
+            if (!this.fileManager) throw Error("No file manager")
+            await this.gitClient.pushFiles(this.fileManager?.fileData, this.sandboxId)
+            return { success: true }
+        }
+
+        // Handle creating a file
+        const handleCreateFile: SocketHandler = async ({ name }: any) => {
+            await createFileRL.consume(connection.userId, 1);
+            return { "success": await this.fileManager?.createFile(name) }
+        }
+
+        // Handle creating a folder
+        const handleCreateFolder: SocketHandler = async ({ name }: any) => {
+            await createFolderRL.consume(connection.userId, 1);
+            return { "success": await this.fileManager?.createFolder(name) }
+        }
+
+        // Handle renaming a file
+        const handleRenameFile: SocketHandler = async ({ fileId, newName }: any) => {
+            await renameFileRL.consume(connection.userId, 1)
+            return this.fileManager?.renameFile(fileId, newName)
+        }
+
+        // Handle deleting a file
+        const handleDeleteFile: SocketHandler = async ({ fileId }: any) => {
+            await deleteFileRL.consume(connection.userId, 1)
+            return this.fileManager?.deleteFile(fileId)
+        }
+
+        // Handle deleting a folder
+        const handleDeleteFolder: SocketHandler = ({ folderId }: any) => {
+            return this.fileManager?.deleteFolder(folderId)
+        }
+
+        // Handle creating a terminal session
+        const handleCreateTerminal: SocketHandler = async ({ id }: any) => {
+            await lockManager.acquireLock(this.sandboxId, async () => {
+                await this.terminalManager?.createTerminal(id, (responseString: string) => {
+                    connection.socket.emit("terminalResponse", { id, data: responseString })
+                    const port = extractPortNumber(responseString)
+                    if (port) {
+                        connection.socket.emit(
+                            "previewURL",
+                            "https://" + this.container?.getHost(port)
+                        )
+                    }
+                })
+            })
+        }
+
+        // Handle resizing a terminal
+        const handleResizeTerminal: SocketHandler = ({ dimensions }: any) => {
+            this.terminalManager?.resizeTerminal(dimensions)
+        }
+
+        // Handle sending data to a terminal
+        const handleTerminalData: SocketHandler = ({ id, data }: any) => {
+            return this.terminalManager?.sendTerminalData(id, data)
+        }
+
+        // Handle closing a terminal
+        const handleCloseTerminal: SocketHandler = ({ id }: any) => {
+            return this.terminalManager?.closeTerminal(id)
+        }
+
+        // Handle generating code
+        const handleGenerateCode: SocketHandler = ({ fileName, code, line, instructions }: any) => {
+            return this.aiWorker.generateCode(connection.userId, fileName, code, line, instructions)
+        }
+
+        return {
+            "heartbeat": handleHeartbeat,
+            "getFile": handleGetFile,
+            "getFolder": handleGetFolder,
+            "saveFile": handleSaveFile,
+            "moveFile": handleMoveFile,
+            "list": handleListApps,
+            "deploy": handleDeploy,
+            "createFile": handleCreateFile,
+            "createFolder": handleCreateFolder,
+            "renameFile": handleRenameFile,
+            "deleteFile": handleDeleteFile,
+            "deleteFolder": handleDeleteFolder,
+            "createTerminal": handleCreateTerminal,
+            "resizeTerminal": handleResizeTerminal,
+            "terminalData": handleTerminalData,
+            "closeTerminal": handleCloseTerminal,
+            "generateCode": handleGenerateCode,
+        };
+
+    }
+
+}

backend/server/src/constants.ts

@@ -0,0 +1,2 @@
+// The amount of time in ms that a container will stay alive without a hearbeat.
+export const CONTAINER_TIMEOUT = 120_000

backend/server/src/index.ts

@@ -1,42 +1,39 @@
 import cors from "cors"
 import dotenv from "dotenv"
-import { Sandbox } from "e2b"
 import express, { Express } from "express"
 import fs from "fs"
 import { createServer } from "http"
-import { Server } from "socket.io"
+import { Server, Socket } from "socket.io"
-import { z } from "zod"
 import { AIWorker } from "./AIWorker"
+
+import { ConnectionManager } from "./ConnectionManager"
 import { DokkuClient } from "./DokkuClient"
-import { FileManager, SandboxFiles } from "./FileManager"
+import { Sandbox } from "./Sandbox"
-import {
-  createFileRL,
-  createFolderRL,
-  deleteFileRL,
-  renameFileRL,
-  saveFileRL,
-} from "./ratelimit"
 import { SecureGitClient } from "./SecureGitClient"
-import { TerminalManager } from "./TerminalManager"
+import { socketAuth } from "./socketAuth"; // Import the new socketAuth middleware
-import { User } from "./types"
+import { TFile, TFolder } from "./types"
-import { LockManager } from "./utils"
+
+// Log errors and send a notification to the client
+export const handleErrors = (message: string, error: any, socket: Socket) => {
+  console.error(message, error);
+  socket.emit("error", `${message} ${error.message ?? error}`);
+};
 
 // Handle uncaught exceptions
 process.on("uncaughtException", (error) => {
   console.error("Uncaught Exception:", error)
   // Do not exit the process
-  // You can add additional logging or recovery logic here
 })
 
 // Handle unhandled promise rejections
 process.on("unhandledRejection", (reason, promise) => {
   console.error("Unhandled Rejection at:", promise, "reason:", reason)
   // Do not exit the process
-  // You can also handle the rejected promise here if needed
 })
 
-// The amount of time in ms that a container will stay alive without a hearbeat.
+// Initialize containers and managers
-const CONTAINER_TIMEOUT = 120_000
+const connections = new ConnectionManager()
+const sandboxes: Record<string, Sandbox> = {}
 
 // Load environment variables
 dotenv.config()
@@ -48,102 +45,23 @@ app.use(cors())
 const httpServer = createServer(app)
 const io = new Server(httpServer, {
   cors: {
-    origin: "*",
+    origin: "*", // Allow connections from any origin
   },
 })
 
-// Check if the sandbox owner is connected
-function isOwnerConnected(sandboxId: string): boolean {
-  return (connections[sandboxId] ?? 0) > 0
-}
-
-// Extract port number from a string
-function extractPortNumber(inputString: string): number | null {
-  const cleanedString = inputString.replace(/\x1B\[[0-9;]*m/g, "")
-  const regex = /http:\/\/localhost:(\d+)/
-  const match = cleanedString.match(regex)
-  return match ? parseInt(match[1]) : null
-}
-
-// Initialize containers and managers
-const containers: Record<string, Sandbox> = {}
-const connections: Record<string, number> = {}
-const fileManagers: Record<string, FileManager> = {}
-const terminalManagers: Record<string, TerminalManager> = {}
-
 // Middleware for socket authentication
-io.use(async (socket, next) => {
+io.use(socketAuth) // Use the new socketAuth middleware
-  // Define the schema for handshake query validation
-  const handshakeSchema = z.object({
-    userId: z.string(),
-    sandboxId: z.string(),
-    EIO: z.string(),
-    transport: z.string(),
-  })
-
-  const q = socket.handshake.query
-  const parseQuery = handshakeSchema.safeParse(q)
-
-  // Check if the query is valid according to the schema
-  if (!parseQuery.success) {
-    next(new Error("Invalid request."))
-    return
-  }
-
-  const { sandboxId, userId } = parseQuery.data
-  // Fetch user data from the database
-  const dbUser = await fetch(
-    `${process.env.DATABASE_WORKER_URL}/api/user?id=${userId}`,
-    {
-      headers: {
-        Authorization: `${process.env.WORKERS_KEY}`,
-      },
-    }
-  )
-  const dbUserJSON = (await dbUser.json()) as User
-
-  // Check if user data was retrieved successfully
-  if (!dbUserJSON) {
-    next(new Error("DB error."))
-    return
-  }
-
-  // Check if the user owns the sandbox or has shared access
-  const sandbox = dbUserJSON.sandbox.find((s) => s.id === sandboxId)
-  const sharedSandboxes = dbUserJSON.usersToSandboxes.find(
-    (uts) => uts.sandboxId === sandboxId
-  )
-
-  // If user doesn't own or have shared access to the sandbox, deny access
-  if (!sandbox && !sharedSandboxes) {
-    next(new Error("Invalid credentials."))
-    return
-  }
-
-  // Set socket data with user information
-  socket.data = {
-    userId,
-    sandboxId: sandboxId,
-    isOwner: sandbox !== undefined,
-  }
-
-  // Allow the connection
-  next()
-})
-
-// Initialize lock manager
-const lockManager = new LockManager()
 
 // Check for required environment variables
 if (!process.env.DOKKU_HOST)
-  console.error("Environment variable DOKKU_HOST is not defined")
+  console.warn("Environment variable DOKKU_HOST is not defined")
 if (!process.env.DOKKU_USERNAME)
-  console.error("Environment variable DOKKU_USERNAME is not defined")
+  console.warn("Environment variable DOKKU_USERNAME is not defined")
 if (!process.env.DOKKU_KEY)
-  console.error("Environment variable DOKKU_KEY is not defined")
+  console.warn("Environment variable DOKKU_KEY is not defined")
 
 // Initialize Dokku client
-const client =
+const dokkuClient =
   process.env.DOKKU_HOST && process.env.DOKKU_KEY && process.env.DOKKU_USERNAME
     ? new DokkuClient({
       host: process.env.DOKKU_HOST,
@@ -151,10 +69,10 @@ const client =
       privateKey: fs.readFileSync(process.env.DOKKU_KEY),
     })
     : null
-client?.connect()
+dokkuClient?.connect()
 
 // Initialize Git client used to deploy Dokku apps
-const git =
+const gitClient =
   process.env.DOKKU_HOST && process.env.DOKKU_KEY
     ? new SecureGitClient(
       `dokku@${process.env.DOKKU_HOST}`,
@@ -170,360 +88,91 @@ const aiWorker = new AIWorker(
   process.env.WORKERS_KEY!
 )
 
-// Handle socket connections
+// Handle a client connecting to the server
 io.on("connection", async (socket) => {
   try {
+    // This data comes is added by our authentication middleware
     const data = socket.data as {
       userId: string
       sandboxId: string
       isOwner: boolean
     }
 
-    // Handle connection based on user type (owner or not)
+    // Register the connection
-    if (data.isOwner) {
+    connections.addConnectionForSandbox(socket, data.sandboxId, data.isOwner)
-      connections[data.sandboxId] = (connections[data.sandboxId] ?? 0) + 1
+
-    } else {
+    // Disable access unless the sandbox owner is connected
-      if (!isOwnerConnected(data.sandboxId)) {
+    if (!data.isOwner && !connections.ownerIsConnected(data.sandboxId)) {
       socket.emit("disableAccess", "The sandbox owner is not connected.")
       return
     }
-    }
 
-    // Create or retrieve container
+    try {
-    const createdContainer = await lockManager.acquireLock(
+      // Create or retrieve the sandbox manager for the given sandbox ID
+      const sandbox = sandboxes[data.sandboxId] ?? new Sandbox(
         data.sandboxId,
-      async () => {
+        {
-        try {
+          aiWorker, dokkuClient, gitClient,
-          // Start a new container if the container doesn't exist or it timed out.
-          if (
-            !containers[data.sandboxId] ||
-            !(await containers[data.sandboxId].isRunning())
-          ) {
-            containers[data.sandboxId] = await Sandbox.create({
-              timeoutMs: CONTAINER_TIMEOUT,
-            })
-            console.log("Created container ", data.sandboxId)
-            return true
-          }
-        } catch (e: any) {
-          console.error(`Error creating container ${data.sandboxId}:`, e)
-          socket.emit("error", `Error: container creation. ${e.message ?? e}`)
-        }
         }
       )
+      sandboxes[data.sandboxId] = sandbox
 
-    // Function to send loaded event
+      // This callback recieves an update when the file list changes, and notifies all relevant connections.
-    const sendLoadedEvent = (files: SandboxFiles) => {
+      const sendFileNotifications = (files: (TFolder | TFile)[]) => {
-      socket.emit("loaded", files.files)
+        connections.connectionsForSandbox(data.sandboxId).forEach((socket: Socket) => {
-    }
+          socket.emit("loaded", files);
+        });
+      };
 
-    // Initialize file and terminal managers if container was created
+      // Initialize the sandbox container
-    if (createdContainer) {
+      // The file manager and terminal managers will be set up if they have been closed
-      fileManagers[data.sandboxId] = new FileManager(
+      await sandbox.initialize(sendFileNotifications)
-        data.sandboxId,
+      socket.emit("loaded", sandbox.fileManager?.files)
-        containers[data.sandboxId],
-        sendLoadedEvent
-      )
-      terminalManagers[data.sandboxId] = new TerminalManager(
-        containers[data.sandboxId]
-      )
-      console.log(`terminal manager set up for ${data.sandboxId}`)
-      await fileManagers[data.sandboxId].initialize()
-    }
 
-    const fileManager = fileManagers[data.sandboxId]
+      // Register event handlers for the sandbox
-    const terminalManager = terminalManagers[data.sandboxId]
+      // For each event handler, listen on the socket for that event
-
+      // Pass connection-specific information to the handlers
-    // Load file list from the file manager into the editor
+      Object.entries(sandbox.handlers({
-    sendLoadedEvent(fileManager.sandboxFiles)
+        userId: data.userId,
-
+        isOwner: data.isOwner,
-    // Handle various socket events (heartbeat, file operations, terminal operations, etc.)
+        socket
-    socket.on("heartbeat", async () => {
+      })).forEach(([event, handler]) => {
+        socket.on(event, async (options: any, callback?: (response: any) => void) => {
           try {
-        // This keeps the container alive for another CONTAINER_TIMEOUT seconds.
+            const result = await handler(options)
-        // The E2B docs are unclear, but the timeout is relative to the time of this method call.
+            callback?.(result);
-        await containers[data.sandboxId].setTimeout(CONTAINER_TIMEOUT)
           } catch (e: any) {
-        console.error("Error setting timeout:", e)
+            handleErrors(`Error processing event "${event}":`, e, socket);
-        socket.emit("error", `Error: set timeout. ${e.message ?? e}`)
           }
-    })
+        });
+      });
 
-    // Handle request to get file content
+      // Handle disconnection event
-    socket.on("getFile", async (fileId: string, callback) => {
-      try {
-        const fileContent = await fileManager.getFile(fileId)
-        callback(fileContent)
-      } catch (e: any) {
-        console.error("Error getting file:", e)
-        socket.emit("error", `Error: get file. ${e.message ?? e}`)
-      }
-    })
-
-    // Handle request to get folder contents
-    socket.on("getFolder", async (folderId: string, callback) => {
-      try {
-        const files = await fileManager.getFolder(folderId)
-        callback(files)
-      } catch (e: any) {
-        console.error("Error getting folder:", e)
-        socket.emit("error", `Error: get folder. ${e.message ?? e}`)
-      }
-    })
-
-    // Handle request to save file
-    socket.on("saveFile", async (fileId: string, body: string) => {
-      try {
-        await saveFileRL.consume(data.userId, 1)
-        await fileManager.saveFile(fileId, body)
-      } catch (e: any) {
-        console.error("Error saving file:", e)
-        socket.emit("error", `Error: file saving. ${e.message ?? e}`)
-      }
-    })
-
-    // Handle request to move file
-    socket.on(
-      "moveFile",
-      async (fileId: string, folderId: string, callback) => {
-        try {
-          const newFiles = await fileManager.moveFile(fileId, folderId)
-          callback(newFiles)
-        } catch (e: any) {
-          console.error("Error moving file:", e)
-          socket.emit("error", `Error: file moving. ${e.message ?? e}`)
-        }
-      }
-    )
-
-    interface CallbackResponse {
-      success: boolean
-      apps?: string[]
-      message?: string
-    }
-
-    // Handle request to list apps
-    socket.on(
-      "list",
-      async (callback: (response: CallbackResponse) => void) => {
-        console.log("Retrieving apps list...")
-        try {
-          if (!client)
-            throw Error("Failed to retrieve apps list: No Dokku client")
-          callback({
-            success: true,
-            apps: await client.listApps(),
-          })
-        } catch (error) {
-          callback({
-            success: false,
-            message: "Failed to retrieve apps list",
-          })
-        }
-      }
-    )
-
-    // Handle request to deploy project
-    socket.on(
-      "deploy",
-      async (callback: (response: CallbackResponse) => void) => {
-        try {
-          // Push the project files to the Dokku server
-          console.log("Deploying project ${data.sandboxId}...")
-          if (!git) throw Error("Failed to retrieve apps list: No git client")
-          // Remove the /project/[id]/ component of each file path:
-          const fixedFilePaths = fileManager.sandboxFiles.fileData.map(
-            (file) => {
-              return {
-                ...file,
-                id: file.id.split("/").slice(2).join("/"),
-              }
-            }
-          )
-          // Push all files to Dokku.
-          await git.pushFiles(fixedFilePaths, data.sandboxId)
-          callback({
-            success: true,
-          })
-        } catch (error) {
-          callback({
-            success: false,
-            message: "Failed to deploy project: " + error,
-          })
-        }
-      }
-    )
-
-    // Handle request to create a new file
-    socket.on("createFile", async (name: string, callback) => {
-      try {
-        await createFileRL.consume(data.userId, 1)
-        const success = await fileManager.createFile(name)
-        callback({ success })
-      } catch (e: any) {
-        console.error("Error creating file:", e)
-        socket.emit("error", `Error: file creation. ${e.message ?? e}`)
-      }
-    })
-
-    // Handle request to create a new folder
-    socket.on("createFolder", async (name: string, callback) => {
-      try {
-        await createFolderRL.consume(data.userId, 1)
-        await fileManager.createFolder(name)
-        callback()
-      } catch (e: any) {
-        console.error("Error creating folder:", e)
-        socket.emit("error", `Error: folder creation. ${e.message ?? e}`)
-      }
-    })
-
-    // Handle request to rename a file
-    socket.on("renameFile", async (fileId: string, newName: string) => {
-      try {
-        await renameFileRL.consume(data.userId, 1)
-        await fileManager.renameFile(fileId, newName)
-      } catch (e: any) {
-        console.error("Error renaming file:", e)
-        socket.emit("error", `Error: file renaming. ${e.message ?? e}`)
-      }
-    })
-
-    // Handle request to delete a file
-    socket.on("deleteFile", async (fileId: string, callback) => {
-      try {
-        await deleteFileRL.consume(data.userId, 1)
-        const newFiles = await fileManager.deleteFile(fileId)
-        callback(newFiles)
-      } catch (e: any) {
-        console.error("Error deleting file:", e)
-        socket.emit("error", `Error: file deletion. ${e.message ?? e}`)
-      }
-    })
-
-    // Handle request to delete a folder
-    socket.on("deleteFolder", async (folderId: string, callback) => {
-      try {
-        const newFiles = await fileManager.deleteFolder(folderId)
-        callback(newFiles)
-      } catch (e: any) {
-        console.error("Error deleting folder:", e)
-        socket.emit("error", `Error: folder deletion. ${e.message ?? e}`)
-      }
-    })
-
-    // Handle request to create a new terminal
-    socket.on("createTerminal", async (id: string, callback) => {
-      try {
-        await lockManager.acquireLock(data.sandboxId, async () => {
-          let terminalManager = terminalManagers[data.sandboxId]
-          if (!terminalManager) {
-            terminalManager = terminalManagers[data.sandboxId] =
-              new TerminalManager(containers[data.sandboxId])
-          }
-
-          await terminalManager.createTerminal(id, (responseString: string) => {
-            socket.emit("terminalResponse", { id, data: responseString })
-            const port = extractPortNumber(responseString)
-            if (port) {
-              socket.emit(
-                "previewURL",
-                "https://" + containers[data.sandboxId].getHost(port)
-              )
-            }
-          })
-        })
-        callback()
-      } catch (e: any) {
-        console.error(`Error creating terminal ${id}:`, e)
-        socket.emit("error", `Error: terminal creation. ${e.message ?? e}`)
-      }
-    })
-
-    // Handle request to resize terminal
-    socket.on(
-      "resizeTerminal",
-      (dimensions: { cols: number; rows: number }) => {
-        try {
-          terminalManager.resizeTerminal(dimensions)
-        } catch (e: any) {
-          console.error("Error resizing terminal:", e)
-          socket.emit("error", `Error: terminal resizing. ${e.message ?? e}`)
-        }
-      }
-    )
-
-    // Handle terminal input data
-    socket.on("terminalData", async (id: string, data: string) => {
-      try {
-        await terminalManager.sendTerminalData(id, data)
-      } catch (e: any) {
-        console.error("Error writing to terminal:", e)
-        socket.emit("error", `Error: writing to terminal. ${e.message ?? e}`)
-      }
-    })
-
-    // Handle request to close terminal
-    socket.on("closeTerminal", async (id: string, callback) => {
-      try {
-        await terminalManager.closeTerminal(id)
-        callback()
-      } catch (e: any) {
-        console.error("Error closing terminal:", e)
-        socket.emit("error", `Error: closing terminal. ${e.message ?? e}`)
-      }
-    })
-
-    // Handle request to generate code
-    socket.on(
-      "generateCode",
-      async (
-        fileName: string,
-        code: string,
-        line: number,
-        instructions: string,
-        callback
-      ) => {
-        try {
-          const result = await aiWorker.generateCode(
-            data.userId,
-            fileName,
-            code,
-            line,
-            instructions
-          )
-          callback(result)
-        } catch (e: any) {
-          console.error("Error generating code:", e)
-          socket.emit("error", `Error: code generation. ${e.message ?? e}`)
-        }
-      }
-    )
-
-    // Handle socket disconnection
       socket.on("disconnect", async () => {
         try {
-        if (data.isOwner) {
+          // Deregister the connection
-          connections[data.sandboxId]--
+          connections.removeConnectionForSandbox(socket, data.sandboxId, data.isOwner)
-        }
 
-        await terminalManager.closeAllTerminals()
+          // If the owner has disconnected from all sockets, close open terminals and file watchers.o
-        await fileManager.closeWatchers()
+          // The sandbox itself will timeout after the heartbeat stops.
-
+          if (data.isOwner && !connections.ownerIsConnected(data.sandboxId)) {
-        if (data.isOwner && connections[data.sandboxId] <= 0) {
+            await sandbox.disconnect()
             socket.broadcast.emit(
               "disableAccess",
               "The sandbox owner has disconnected."
             )
           }
         } catch (e: any) {
-        console.log("Error disconnecting:", e)
+          handleErrors("Error disconnecting:", e, socket);
-        socket.emit("error", `Error: disconnecting. ${e.message ?? e}`)
         }
       })
+
     } catch (e: any) {
-    console.error("Error connecting:", e)
+      handleErrors(`Error initializing sandbox ${data.sandboxId}:`, e, socket);
-    socket.emit("error", `Error: connection. ${e.message ?? e}`)
+    }
+
+  } catch (e: any) {
+    handleErrors("Error connecting:", e, socket);
   }
 })
 

backend/server/src/socketAuth.ts

@@ -0,0 +1,63 @@
+import { Socket } from "socket.io"
+import { z } from "zod"
+import { User } from "./types"
+
+// Middleware for socket authentication
+export const socketAuth = async (socket: Socket, next: Function) => {
+    // Define the schema for handshake query validation
+    const handshakeSchema = z.object({
+        userId: z.string(),
+        sandboxId: z.string(),
+        EIO: z.string(),
+        transport: z.string(),
+    })
+
+    const q = socket.handshake.query
+    const parseQuery = handshakeSchema.safeParse(q)
+
+    // Check if the query is valid according to the schema
+    if (!parseQuery.success) {
+        next(new Error("Invalid request."))
+        return
+    }
+
+    const { sandboxId, userId } = parseQuery.data
+    // Fetch user data from the database
+    const dbUser = await fetch(
+        `${process.env.DATABASE_WORKER_URL}/api/user?id=${userId}`,
+        {
+            headers: {
+                Authorization: `${process.env.WORKERS_KEY}`,
+            },
+        }
+    )
+    const dbUserJSON = (await dbUser.json()) as User
+
+    // Check if user data was retrieved successfully
+    if (!dbUserJSON) {
+        next(new Error("DB error."))
+        return
+    }
+
+    // Check if the user owns the sandbox or has shared access
+    const sandbox = dbUserJSON.sandbox.find((s) => s.id === sandboxId)
+    const sharedSandboxes = dbUserJSON.usersToSandboxes.find(
+        (uts) => uts.sandboxId === sandboxId
+    )
+
+    // If user doesn't own or have shared access to the sandbox, deny access
+    if (!sandbox && !sharedSandboxes) {
+        next(new Error("Invalid credentials."))
+        return
+    }
+
+    // Set socket data with user information
+    socket.data = {
+        userId,
+        sandboxId: sandboxId,
+        isOwner: sandbox !== undefined,
+    }
+
+    // Allow the connection
+    next()
+}

backend/server/src/types.ts

@@ -68,3 +68,8 @@ export type R2FileBody = R2FileData & {
   json: Promise<any>
   blob: Promise<Blob>
 }
+export interface DokkuResponse {
+  success: boolean
+  apps?: string[]
+  message?: string
+}

frontend/components/editor/generate.tsx

@@ -68,10 +68,12 @@ export default function GenerateInput({
     setCurrentPrompt(input)
     socket.emit(
       "generateCode",
-      data.fileName,
+      {
-      data.code,
+        fileName: data.fileName,
-      data.line,
+        code: data.code,
-      regenerate ? currentPrompt : input,
+        line: data.line,
+        instructions: regenerate ? currentPrompt : input
+      },
       (res: { response: string; success: boolean }) => {
         console.log("Generated code", res.response, res.success)
         // if (!res.success) {

frontend/components/editor/index.tsx

@@ -107,7 +107,7 @@ export default function CodeEditor({
 
   // Editor state
   const [editorLanguage, setEditorLanguage] = useState("plaintext")
-  console.log("editor language: ",editorLanguage)
+  console.log("editor language: ", editorLanguage)
   const [cursorLine, setCursorLine] = useState(0)
   const [editorRef, setEditorRef] =
     useState<monaco.editor.IStandaloneCodeEditor>()
@@ -207,7 +207,7 @@ export default function CodeEditor({
     )
     const fetchFileContent = (fileId: string): Promise<string> => {
       return new Promise((resolve) => {
-        socket?.emit("getFile", fileId, (content: string) => {
+        socket?.emit("getFile", { fileId }, (content: string) => {
           resolve(content)
         })
       })
@@ -532,7 +532,7 @@ export default function CodeEditor({
         )
         console.log(`Saving file...${activeFileId}`)
         console.log(`Saving file...${content}`)
-        socket?.emit("saveFile", activeFileId, content)
+        socket?.emit("saveFile", { fileId: activeFileId, body: content })
       }
     }, Number(process.env.FILE_SAVE_DEBOUNCE_DELAY) || 1000),
     [socket, fileContents]
@@ -649,7 +649,7 @@ export default function CodeEditor({
 
   // Socket event listener effect
   useEffect(() => {
-    const onConnect = () => {}
+    const onConnect = () => { }
 
     const onDisconnect = () => {
       setTerminals([])
@@ -715,7 +715,7 @@ export default function CodeEditor({
 
   // Debounced function to get file content
   const debouncedGetFile = (tabId: any, callback: any) => {
-    socket?.emit("getFile", tabId, callback)
+    socket?.emit("getFile", { fileId: tabId }, callback)
   } // 300ms debounce delay, adjust as needed
 
   const selectFile = (tab: TTab) => {
@@ -835,7 +835,7 @@ export default function CodeEditor({
       return false
     }
 
-    socket?.emit("renameFile", id, newName)
+    socket?.emit("renameFile", { fileId: id, newName })
     setTabs((prev) =>
       prev.map((tab) => (tab.id === id ? { ...tab, name: newName } : tab))
     )
@@ -844,7 +844,7 @@ export default function CodeEditor({
   }
 
   const handleDeleteFile = (file: TFile) => {
-    socket?.emit("deleteFile", file.id, (response: (TFolder | TFile)[]) => {
+    socket?.emit("deleteFile", { fileId: file.id }, (response: (TFolder | TFile)[]) => {
       setFiles(response)
     })
     closeTab(file.id)
@@ -854,11 +854,11 @@ export default function CodeEditor({
     setDeletingFolderId(folder.id)
     console.log("deleting folder", folder.id)
 
-    socket?.emit("getFolder", folder.id, (response: string[]) =>
+    socket?.emit("getFolder", { folderId: folder.id }, (response: string[]) =>
       closeTabs(response)
     )
 
-    socket?.emit("deleteFolder", folder.id, (response: (TFolder | TFile)[]) => {
+    socket?.emit("deleteFolder", { folderId: folder.id }, (response: (TFolder | TFile)[]) => {
       setFiles(response)
       setDeletingFolderId("")
     })
@@ -902,7 +902,7 @@ export default function CodeEditor({
         <DisableAccessModal
           message={disableAccess.message}
           open={disableAccess.isDisabled}
-          setOpen={() => {}}
+          setOpen={() => { }}
         />
         <Loading />
       </>

frontend/components/editor/sidebar/index.tsx

@@ -87,8 +87,10 @@ export default function Sidebar({
         setMovingId(fileId)
         socket.emit(
           "moveFile",
+          {
             fileId,
-          folderId,
+            folderId
+          },
           (response: (TFolder | TFile)[]) => {
             setFiles(response)
             setMovingId("")

frontend/components/editor/sidebar/new.tsx

@@ -27,7 +27,7 @@ export default function New({
         if (type === "file") {
           socket.emit(
             "createFile",
-            name,
+            { name },
             ({ success }: { success: boolean }) => {
               if (success) {
                 addNew(name, type)
@@ -35,7 +35,7 @@ export default function New({
             }
           )
         } else {
-          socket.emit("createFolder", name, () => {
+          socket.emit("createFolder", { name }, () => {
             addNew(name, type)
           })
         }

frontend/components/editor/terminals/terminal.tsx

@@ -65,12 +65,12 @@ export default function EditorTerminal({
     }
 
     const disposableOnData = term.onData((data) => {
-      socket.emit("terminalData", id, data)
+      socket.emit("terminalData", { id, data })
     })
 
     const disposableOnResize = term.onResize((dimensions) => {
       fitAddonRef.current?.fit()
-      socket.emit("terminalResize", dimensions)
+      socket.emit("terminalResize", { dimensions })
     })
     const resizeObserver = new ResizeObserver(
       debounce((entries) => {

frontend/context/TerminalContext.tsx

@@ -63,7 +63,7 @@ export const TerminalProvider: React.FC<{ children: React.ReactNode }> = ({
         terminals,
         setTerminals,
         setActiveTerminalId,
-        setClosingTerminal: () => {},
+        setClosingTerminal: () => { },
         socket,
         activeTerminalId,
       })
@@ -73,7 +73,7 @@ export const TerminalProvider: React.FC<{ children: React.ReactNode }> = ({
   const deploy = (callback: () => void) => {
     if (!socket) console.error("Couldn't deploy: No socket")
     console.log("Deploying...")
-    socket?.emit("deploy", () => {
+    socket?.emit("deploy", {}, () => {
       callback()
     })
   }

frontend/lib/terminal.ts

@@ -32,9 +32,9 @@ export const createTerminal = ({
   setActiveTerminalId(id)
 
   setTimeout(() => {
-    socket.emit("createTerminal", id, () => {
+    socket.emit("createTerminal", { id }, () => {
       setCreatingTerminal(false)
-      if (command) socket.emit("terminalData", id, command + "\n")
+      if (command) socket.emit("terminalData", { id, data: command + "\n" })
     })
   }, 1000)
 }
@@ -75,7 +75,7 @@ export const closeTerminal = ({
 
   setClosingTerminal(term.id)
 
-  socket.emit("closeTerminal", term.id, () => {
+  socket.emit("closeTerminal", { id: term.id }, () => {
     setClosingTerminal("")
 
     const nextId =

tests/index.ts

@@ -1,6 +1,6 @@
 // Import necessary modules
-import { io, Socket } from "socket.io-client";
 import dotenv from "dotenv";
+import { io, Socket } from "socket.io-client";
 
 dotenv.config();
 
@@ -21,7 +21,7 @@ socketRef.on("connect", async () => {
   console.log("Connected to the server");
   await new Promise((resolve) => setTimeout(resolve, 1000));
 
-  socketRef.emit("list", (response: CallbackResponse) => {
+  socketRef.emit("list", {}, (response: CallbackResponse) => {
     if (response.success) {
       console.log("List of apps:", response.apps);
     } else {
@@ -29,7 +29,7 @@ socketRef.on("connect", async () => {
     }
   });
 
-  socketRef.emit("deploy", (response: CallbackResponse) => {
+  socketRef.emit("deploy", {}, (response: CallbackResponse) => {
     if (response.success) {
       console.log("It worked!");
     } else {