diff --git a/backend/server/src/ConnectionManager.ts b/backend/server/src/ConnectionManager.ts new file mode 100644 index 0000000..45b5432 --- /dev/null +++ b/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 = {} + // Stores all sockets connected to a given sandbox + private sockets: Record> = {} + + // 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 { + return this.sockets[sandboxId] ?? new Set(); + } + +} \ No newline at end of file diff --git a/backend/server/src/FileManager.ts b/backend/server/src/FileManager.ts index 278d060..2f58fbf 100644 --- a/backend/server/src/FileManager.ts +++ b/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 { const remotePaths = await RemoteFileStorage.getSandboxPaths(this.sandboxId) const localPaths = this.getLocalFileIds(remotePaths) - this.sandboxFiles.fileData = await this.generateFileData(localPaths) - return this.sandboxFiles.fileData + this.fileData = await this.generateFileData(localPaths) + 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) - return this.sandboxFiles.files + this.files = generateFileStructure(localPaths) + 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 { - 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 file = this.sandboxFiles.files.find((f) => f.id === fileId) - if (!fileData || !file) return this.sandboxFiles.files + const fileData = this.fileData.find((f) => f.id === fileId) + const file = this.files.find((f) => f.id === fileId) + 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 { - const fileData = this.sandboxFiles.fileData.find((f) => f.id === fileId) - const file = this.sandboxFiles.files.find((f) => f.id === fileId) + const fileData = this.fileData.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) - if (!file) return this.sandboxFiles.files + const file = this.fileData.find((f) => f.id === fileId) + 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)) diff --git a/backend/server/src/Sandbox.ts b/backend/server/src/Sandbox.ts new file mode 100644 index 0000000..fda2237 --- /dev/null +++ b/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> = (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, + }; + + } + +} \ No newline at end of file diff --git a/backend/server/src/constants.ts b/backend/server/src/constants.ts new file mode 100644 index 0000000..dfd5ce3 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/backend/server/src/index.ts b/backend/server/src/index.ts index 788c329..cf95824 100644 --- a/backend/server/src/index.ts +++ b/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 { z } from "zod" +import { Server, Socket } from "socket.io" import { AIWorker } from "./AIWorker" + +import { ConnectionManager } from "./ConnectionManager" import { DokkuClient } from "./DokkuClient" -import { FileManager, SandboxFiles } from "./FileManager" -import { - createFileRL, - createFolderRL, - deleteFileRL, - renameFileRL, - saveFileRL, -} from "./ratelimit" +import { Sandbox } from "./Sandbox" import { SecureGitClient } from "./SecureGitClient" -import { TerminalManager } from "./TerminalManager" -import { User } from "./types" -import { LockManager } from "./utils" +import { socketAuth } from "./socketAuth"; // Import the new socketAuth middleware +import { TFile, TFolder } from "./types" + +// 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. -const CONTAINER_TIMEOUT = 120_000 +// Initialize containers and managers +const connections = new ConnectionManager() +const sandboxes: Record = {} // Load environment variables dotenv.config() @@ -48,118 +45,39 @@ 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 = {} -const connections: Record = {} -const fileManagers: Record = {} -const terminalManagers: Record = {} - // Middleware for socket authentication -io.use(async (socket, next) => { - // 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() +io.use(socketAuth) // Use the new socketAuth middleware // 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, - username: process.env.DOKKU_USERNAME, - privateKey: fs.readFileSync(process.env.DOKKU_KEY), - }) + host: process.env.DOKKU_HOST, + username: process.env.DOKKU_USERNAME, + 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}`, - process.env.DOKKU_KEY - ) + `dokku@${process.env.DOKKU_HOST}`, + process.env.DOKKU_KEY + ) : null // Add this near the top of the file, after other initializations @@ -170,364 +88,95 @@ 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) - if (data.isOwner) { - connections[data.sandboxId] = (connections[data.sandboxId] ?? 0) + 1 - } else { - if (!isOwnerConnected(data.sandboxId)) { - socket.emit("disableAccess", "The sandbox owner is not connected.") - return - } + // Register the connection + connections.addConnectionForSandbox(socket, data.sandboxId, data.isOwner) + + // Disable access unless the sandbox owner is connected + if (!data.isOwner && !connections.ownerIsConnected(data.sandboxId)) { + socket.emit("disableAccess", "The sandbox owner is not connected.") + return } - // Create or retrieve container - const createdContainer = await lockManager.acquireLock( - data.sandboxId, - async () => { - try { - // 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}`) - } - } - ) - - // Function to send loaded event - const sendLoadedEvent = (files: SandboxFiles) => { - socket.emit("loaded", files.files) - } - - // Initialize file and terminal managers if container was created - if (createdContainer) { - fileManagers[data.sandboxId] = new FileManager( + try { + // Create or retrieve the sandbox manager for the given sandbox ID + const sandbox = sandboxes[data.sandboxId] ?? new Sandbox( data.sandboxId, - containers[data.sandboxId], - sendLoadedEvent + { + aiWorker, dokkuClient, gitClient, + } ) - terminalManagers[data.sandboxId] = new TerminalManager( - containers[data.sandboxId] - ) - console.log(`terminal manager set up for ${data.sandboxId}`) - await fileManagers[data.sandboxId].initialize() - } + sandboxes[data.sandboxId] = sandbox - const fileManager = fileManagers[data.sandboxId] - const terminalManager = terminalManagers[data.sandboxId] + // This callback recieves an update when the file list changes, and notifies all relevant connections. + const sendFileNotifications = (files: (TFolder | TFile)[]) => { + connections.connectionsForSandbox(data.sandboxId).forEach((socket: Socket) => { + socket.emit("loaded", files); + }); + }; - // Load file list from the file manager into the editor - sendLoadedEvent(fileManager.sandboxFiles) + // Initialize the sandbox container + // The file manager and terminal managers will be set up if they have been closed + await sandbox.initialize(sendFileNotifications) + socket.emit("loaded", sandbox.fileManager?.files) - // Handle various socket events (heartbeat, file operations, terminal operations, etc.) - socket.on("heartbeat", async () => { - try { - // This keeps the container alive for another CONTAINER_TIMEOUT seconds. - // The E2B docs are unclear, but the timeout is relative to the time of this method call. - await containers[data.sandboxId].setTimeout(CONTAINER_TIMEOUT) - } catch (e: any) { - console.error("Error setting timeout:", e) - socket.emit("error", `Error: set timeout. ${e.message ?? e}`) - } - }) - - // Handle request to get file content - 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]) + // Register event handlers for the sandbox + // For each event handler, listen on the socket for that event + // Pass connection-specific information to the handlers + Object.entries(sandbox.handlers({ + userId: data.userId, + isOwner: data.isOwner, + socket + })).forEach(([event, handler]) => { + socket.on(event, async (options: any, callback?: (response: any) => void) => { + try { + const result = await handler(options) + callback?.(result); + } catch (e: any) { + handleErrors(`Error processing event "${event}":`, e, socket); } + }); + }); - 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 }) => { + // Handle disconnection event + socket.on("disconnect", async () => { try { - terminalManager.resizeTerminal(dimensions) + // Deregister the connection + connections.removeConnectionForSandbox(socket, data.sandboxId, data.isOwner) + + // If the owner has disconnected from all sockets, close open terminals and file watchers.o + // The sandbox itself will timeout after the heartbeat stops. + if (data.isOwner && !connections.ownerIsConnected(data.sandboxId)) { + await sandbox.disconnect() + socket.broadcast.emit( + "disableAccess", + "The sandbox owner has disconnected." + ) + } } catch (e: any) { - console.error("Error resizing terminal:", e) - socket.emit("error", `Error: terminal resizing. ${e.message ?? e}`) + handleErrors("Error disconnecting:", e, socket); } - } - ) + }) - // 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}`) - } - }) + } catch (e: any) { + handleErrors(`Error initializing sandbox ${data.sandboxId}:`, e, socket); + } - // 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) { - connections[data.sandboxId]-- - } - - await terminalManager.closeAllTerminals() - await fileManager.closeWatchers() - - if (data.isOwner && connections[data.sandboxId] <= 0) { - socket.broadcast.emit( - "disableAccess", - "The sandbox owner has disconnected." - ) - } - } catch (e: any) { - console.log("Error disconnecting:", e) - socket.emit("error", `Error: disconnecting. ${e.message ?? e}`) - } - }) } catch (e: any) { - console.error("Error connecting:", e) - socket.emit("error", `Error: connection. ${e.message ?? e}`) + handleErrors("Error connecting:", e, socket); } }) // Start the server httpServer.listen(port, () => { console.log(`Server running on port ${port}`) -}) +}) \ No newline at end of file diff --git a/backend/server/src/socketAuth.ts b/backend/server/src/socketAuth.ts new file mode 100644 index 0000000..3bd83b1 --- /dev/null +++ b/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() +} diff --git a/backend/server/src/types.ts b/backend/server/src/types.ts index 42ad6d0..93e45e6 100644 --- a/backend/server/src/types.ts +++ b/backend/server/src/types.ts @@ -68,3 +68,8 @@ export type R2FileBody = R2FileData & { json: Promise blob: Promise } +export interface DokkuResponse { + success: boolean + apps?: string[] + message?: string +} diff --git a/backend/server/src/utils.ts b/backend/server/src/utils.ts index 5ae1377..dd33984 100644 --- a/backend/server/src/utils.ts +++ b/backend/server/src/utils.ts @@ -20,4 +20,4 @@ export class LockManager { } return await this.locks[key] } -} +} \ No newline at end of file diff --git a/frontend/components/editor/generate.tsx b/frontend/components/editor/generate.tsx index 9e4bd09..0b5ff39 100644 --- a/frontend/components/editor/generate.tsx +++ b/frontend/components/editor/generate.tsx @@ -68,10 +68,12 @@ export default function GenerateInput({ setCurrentPrompt(input) socket.emit( "generateCode", - data.fileName, - data.code, - data.line, - regenerate ? currentPrompt : input, + { + fileName: data.fileName, + code: data.code, + line: data.line, + instructions: regenerate ? currentPrompt : input + }, (res: { response: string; success: boolean }) => { console.log("Generated code", res.response, res.success) // if (!res.success) { diff --git a/frontend/components/editor/index.tsx b/frontend/components/editor/index.tsx index e20a6d0..b426f9c 100644 --- a/frontend/components/editor/index.tsx +++ b/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() @@ -207,7 +207,7 @@ export default function CodeEditor({ ) const fetchFileContent = (fileId: string): Promise => { 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) => { @@ -777,8 +777,8 @@ export default function CodeEditor({ ? numTabs === 1 ? null : index < numTabs - 1 - ? tabs[index + 1].id - : tabs[index - 1].id + ? tabs[index + 1].id + : tabs[index - 1].id : activeFileId setTabs((prev) => prev.filter((t) => t.id !== id)) @@ -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({ {}} + setOpen={() => { }} /> @@ -944,8 +944,8 @@ export default function CodeEditor({ code: (isSelected && editorRef?.getSelection() ? editorRef - ?.getModel() - ?.getValueInRange(editorRef?.getSelection()!) + ?.getModel() + ?.getValueInRange(editorRef?.getSelection()!) : editorRef?.getValue()) ?? "", line: generate.line, }} @@ -1075,62 +1075,62 @@ export default function CodeEditor({ ) : // Note clerk.loaded is required here due to a bug: https://github.com/clerk/javascript/issues/1643 - clerk.loaded ? ( - <> - {provider && userInfo ? ( - - ) : null} - { - // If the new content is different from the cached content, update it - if (value !== fileContents[activeFileId]) { - setActiveFileContent(value ?? "") // Update the active file content - // Mark the file as unsaved by setting 'saved' to false - setTabs((prev) => - prev.map((tab) => - tab.id === activeFileId - ? { ...tab, saved: false } - : tab + clerk.loaded ? ( + <> + {provider && userInfo ? ( + + ) : null} + { + // If the new content is different from the cached content, update it + if (value !== fileContents[activeFileId]) { + setActiveFileContent(value ?? "") // Update the active file content + // Mark the file as unsaved by setting 'saved' to false + setTabs((prev) => + prev.map((tab) => + tab.id === activeFileId + ? { ...tab, saved: false } + : tab + ) ) - ) - } else { - // If the content matches the cached content, mark the file as saved - setTabs((prev) => - prev.map((tab) => - tab.id === activeFileId - ? { ...tab, saved: true } - : tab + } else { + // If the content matches the cached content, mark the file as saved + setTabs((prev) => + prev.map((tab) => + tab.id === activeFileId + ? { ...tab, saved: true } + : tab + ) ) - ) - } - }} - options={{ - tabSize: 2, - minimap: { - enabled: false, - }, - padding: { - bottom: 4, - top: 4, - }, - scrollBeyondLastLine: false, - fixedOverflowWidgets: true, - fontFamily: "var(--font-geist-mono)", - }} - theme={theme === "light" ? "vs" : "vs-dark"} - value={activeFileContent} - /> - - ) : ( -
- - Waiting for Clerk to load... -
- )} + } + }} + options={{ + tabSize: 2, + minimap: { + enabled: false, + }, + padding: { + bottom: 4, + top: 4, + }, + scrollBeyondLastLine: false, + fixedOverflowWidgets: true, + fontFamily: "var(--font-geist-mono)", + }} + theme={theme === "light" ? "vs" : "vs-dark"} + value={activeFileContent} + /> + + ) : ( +
+ + Waiting for Clerk to load... +
+ )} @@ -1140,10 +1140,10 @@ export default function CodeEditor({ isAIChatOpen && isHorizontalLayout ? "horizontal" : isAIChatOpen - ? "vertical" - : isHorizontalLayout - ? "horizontal" - : "vertical" + ? "vertical" + : isHorizontalLayout + ? "horizontal" + : "vertical" } > { setFiles(response) setMovingId("") diff --git a/frontend/components/editor/sidebar/new.tsx b/frontend/components/editor/sidebar/new.tsx index 7fec344..ca7dbc9 100644 --- a/frontend/components/editor/sidebar/new.tsx +++ b/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) }) } diff --git a/frontend/components/editor/terminals/terminal.tsx b/frontend/components/editor/terminals/terminal.tsx index 4790a44..19b3980 100644 --- a/frontend/components/editor/terminals/terminal.tsx +++ b/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) => { diff --git a/frontend/context/TerminalContext.tsx b/frontend/context/TerminalContext.tsx index a7f131a..5af74a8 100644 --- a/frontend/context/TerminalContext.tsx +++ b/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() }) } diff --git a/frontend/lib/terminal.ts b/frontend/lib/terminal.ts index a91db3c..1d0edbc 100644 --- a/frontend/lib/terminal.ts +++ b/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 = @@ -83,8 +83,8 @@ export const closeTerminal = ({ ? numTerminals === 1 ? null : index < numTerminals - 1 - ? terminals[index + 1].id - : terminals[index - 1].id + ? terminals[index + 1].id + : terminals[index - 1].id : activeTerminalId setTerminals((prev) => prev.filter((t) => t.id !== term.id)) diff --git a/tests/index.ts b/tests/index.ts index a36868c..951eefe 100644 --- a/tests/index.ts +++ b/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 {