diff --git a/backend/server/src/ConnectionManager.ts b/backend/server/src/ConnectionManager.ts index 45b5432..a95f82e 100644 --- a/backend/server/src/ConnectionManager.ts +++ b/backend/server/src/ConnectionManager.ts @@ -1,58 +1,61 @@ import { Socket } from "socket.io" class Counter { - private count: number = 0 + private count: number = 0 - increment() { - this.count++ - } + increment() { + this.count++ + } - decrement() { - this.count = Math.max(0, this.count - 1) - } + decrement() { + this.count = Math.max(0, this.count - 1) + } - getValue(): number { - return this.count - } + 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> = {} + // 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 + // 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() } + } - // Adds a connection for a sandbox - addConnectionForSandbox(socket: Socket, sandboxId: string, isOwner: boolean) { - this.sockets[sandboxId] ??= new Set() - this.sockets[sandboxId].add(socket) + // Removes a connection for a sandbox + removeConnectionForSandbox( + socket: Socket, + sandboxId: string, + isOwner: boolean + ) { + this.sockets[sandboxId]?.delete(socket) - // If the connection is for the owner, increments the owner connection counter - if (isOwner) { - this.ownerConnections[sandboxId] ??= new Counter() - this.ownerConnections[sandboxId].increment() - } + // If the connection being removed is for the owner, decrements the owner connection counter + if (isOwner) { + this.ownerConnections[sandboxId]?.decrement() } + } - // 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 + // Returns the set of sockets connected to a given sandbox + connectionsForSandbox(sandboxId: string): Set { + return this.sockets[sandboxId] ?? new Set() + } +} diff --git a/backend/server/src/FileManager.ts b/backend/server/src/FileManager.ts index abc48f2..63093ae 100644 --- a/backend/server/src/FileManager.ts +++ b/backend/server/src/FileManager.ts @@ -23,7 +23,11 @@ function generateFileStructure(paths: string[]): (TFolder | TFile)[] { } } else { if (isFile) { - const file: TFile = { id: `/${parts.join("/")}`, type: "file", name: part } + const file: TFile = { + id: `/${parts.join("/")}`, + type: "file", + name: part, + } current.children.push(file) } else { const folder: TFolder = { @@ -75,7 +79,9 @@ export class FileManager { if (isFile) { const fileId = `/${parts.join("/")}` - const data = await RemoteFileStorage.fetchFileContent(`projects/${this.sandboxId}${fileId}`) + const data = await RemoteFileStorage.fetchFileContent( + `projects/${this.sandboxId}${fileId}` + ) fileData.push({ id: fileId, data }) } } @@ -91,7 +97,7 @@ export class FileManager { // Convert remote file path to local file path private getLocalFileId(remoteId: string): string | undefined { const allParts = remoteId.split("/") - if (allParts[1] !== this.sandboxId) return undefined; + if (allParts[1] !== this.sandboxId) return undefined return allParts.slice(2).join("/") } @@ -99,7 +105,7 @@ export class FileManager { private getLocalFileIds(remoteIds: string[]): string[] { return remoteIds .map(this.getLocalFileId.bind(this)) - .filter((id) => id !== undefined); + .filter((id) => id !== undefined) } // Download files from remote storage @@ -120,7 +126,6 @@ export class FileManager { // Initialize the FileManager async initialize() { - // Download files from remote file storage await this.updateFileStructure() await this.updateFileData() @@ -141,10 +146,14 @@ export class FileManager { await Promise.all(promises) // Reload file list from the container to include template files - const result = await this.sandbox.commands.run(`find "${this.dirName}" -type f`); // List all files recursively - const localPaths = result.stdout.split('\n').filter(path => path); // Split the output into an array and filter out empty strings - const relativePaths = localPaths.map(filePath => path.posix.relative(this.dirName, filePath)); // Convert absolute paths to relative paths - this.files = generateFileStructure(relativePaths); + const result = await this.sandbox.commands.run( + `find "${this.dirName}" -type f` + ) // List all files recursively + const localPaths = result.stdout.split("\n").filter((path) => path) // Split the output into an array and filter out empty strings + const relativePaths = localPaths.map((filePath) => + path.posix.relative(this.dirName, filePath) + ) // Convert absolute paths to relative paths + this.files = generateFileStructure(relativePaths) // Make the logged in user the owner of all project files this.fixPermissions() @@ -169,9 +178,7 @@ export class FileManager { // Change the owner of the project directory to user private async fixPermissions() { try { - await this.sandbox.commands.run( - `sudo chown -R user "${this.dirName}"` - ) + await this.sandbox.commands.run(`sudo chown -R user "${this.dirName}"`) } catch (e: any) { console.log("Failed to fix permissions: " + e) } @@ -193,7 +200,10 @@ export class FileManager { // This is the absolute file path in the container const containerFilePath = path.posix.join(directory, event.name) // This is the file path relative to the project directory - const sandboxFilePath = removeDirName(containerFilePath, this.dirName) + const sandboxFilePath = removeDirName( + containerFilePath, + this.dirName + ) // This is the directory being watched relative to the project directory const sandboxDirectory = removeDirName(directory, this.dirName) @@ -218,16 +228,16 @@ export class FileManager { const newItem = isDir ? ({ - id: sandboxFilePath, - name: event.name, - type: "folder", - children: [], - } as TFolder) + id: sandboxFilePath, + name: event.name, + type: "folder", + children: [], + } as TFolder) : ({ - id: sandboxFilePath, - name: event.name, - type: "file", - } as TFile) + id: sandboxFilePath, + name: event.name, + type: "file", + } as TFile) if (folder) { // If the folder exists, add the new item (file/folder) as a child @@ -361,7 +371,9 @@ export class FileManager { // Get folder content async getFolder(folderId: string): Promise { - const remotePaths = await RemoteFileStorage.getFolder(this.getRemoteFileId(folderId)) + const remotePaths = await RemoteFileStorage.getFolder( + this.getRemoteFileId(folderId) + ) return this.getLocalFileIds(remotePaths) } @@ -400,7 +412,11 @@ export class FileManager { fileData.id = newFileId file.id = newFileId - await RemoteFileStorage.renameFile(this.getRemoteFileId(fileId), this.getRemoteFileId(newFileId), fileData.data) + await RemoteFileStorage.renameFile( + this.getRemoteFileId(fileId), + this.getRemoteFileId(newFileId), + fileData.data + ) return this.updateFileStructure() } @@ -465,7 +481,11 @@ export class FileManager { await this.moveFileInContainer(fileId, newFileId) await this.fixPermissions() - await RemoteFileStorage.renameFile(this.getRemoteFileId(fileId), this.getRemoteFileId(newFileId), fileData.data) + await RemoteFileStorage.renameFile( + this.getRemoteFileId(fileId), + this.getRemoteFileId(newFileId), + fileData.data + ) fileData.id = newFileId file.id = newFileId @@ -477,9 +497,7 @@ export class FileManager { if (!file) return this.files await this.sandbox.files.remove(path.posix.join(this.dirName, fileId)) - this.fileData = this.fileData.filter( - (f) => f.id !== fileId - ) + this.fileData = this.fileData.filter((f) => f.id !== fileId) await RemoteFileStorage.deleteFile(this.getRemoteFileId(fileId)) return this.updateFileStructure() @@ -487,14 +505,14 @@ export class FileManager { // Delete a folder async deleteFolder(folderId: string): Promise<(TFolder | TFile)[]> { - const files = await RemoteFileStorage.getFolder(this.getRemoteFileId(folderId)) + const files = await RemoteFileStorage.getFolder( + this.getRemoteFileId(folderId) + ) await Promise.all( files.map(async (file) => { await this.sandbox.files.remove(path.posix.join(this.dirName, file)) - this.fileData = this.fileData.filter( - (f) => f.id !== file - ) + this.fileData = this.fileData.filter((f) => f.id !== file) await RemoteFileStorage.deleteFile(this.getRemoteFileId(file)) }) ) diff --git a/backend/server/src/RemoteFileStorage.ts b/backend/server/src/RemoteFileStorage.ts index e5ed4b2..ff705d2 100644 --- a/backend/server/src/RemoteFileStorage.ts +++ b/backend/server/src/RemoteFileStorage.ts @@ -61,11 +61,7 @@ export const RemoteFileStorage = { return res.ok }, - renameFile: async ( - fileId: string, - newFileId: string, - data: string - ) => { + renameFile: async (fileId: string, newFileId: string, data: string) => { const res = await fetch(`${process.env.STORAGE_WORKER_URL}/api/rename`, { method: "POST", headers: { @@ -111,7 +107,7 @@ export const RemoteFileStorage = { } ) return (await res.json()).size - } + }, } -export default RemoteFileStorage \ No newline at end of file +export default RemoteFileStorage diff --git a/backend/server/src/Sandbox.ts b/backend/server/src/Sandbox.ts index 5de7754..4657523 100644 --- a/backend/server/src/Sandbox.ts +++ b/backend/server/src/Sandbox.ts @@ -5,11 +5,11 @@ import { CONTAINER_TIMEOUT } from "./constants" import { DokkuClient } from "./DokkuClient" import { FileManager } from "./FileManager" import { - createFileRL, - createFolderRL, - deleteFileRL, - renameFileRL, - saveFileRL, + createFileRL, + createFolderRL, + deleteFileRL, + renameFileRL, + saveFileRL, } from "./ratelimit" import { SecureGitClient } from "./SecureGitClient" import { TerminalManager } from "./TerminalManager" @@ -18,245 +18,267 @@ import { LockManager } from "./utils" const lockManager = new LockManager() // Define a type for SocketHandler functions -type SocketHandler> = (args: T) => any; +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 + 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; -}; + aiWorker: AIWorker + dokkuClient: DokkuClient | null + gitClient: SecureGitClient | null +} export class Sandbox { + // Sandbox properties: + sandboxId: string + type: string + fileManager: FileManager | null + terminalManager: TerminalManager | null + container: E2BSandbox | null + // Server context: + dokkuClient: DokkuClient | null + gitClient: SecureGitClient | null + aiWorker: AIWorker + + constructor( + sandboxId: string, + type: string, + { aiWorker, dokkuClient, gitClient }: ServerContext + ) { // Sandbox properties: - sandboxId: string; - type: string; - fileManager: FileManager | null; - terminalManager: TerminalManager | null; - container: E2BSandbox | null; + this.sandboxId = sandboxId + this.type = type + this.fileManager = null + this.terminalManager = null + this.container = null // Server context: - dokkuClient: DokkuClient | null; - gitClient: SecureGitClient | null; - aiWorker: AIWorker; + this.aiWorker = aiWorker + this.dokkuClient = dokkuClient + this.gitClient = gitClient + } - constructor(sandboxId: string, type: string, { aiWorker, dokkuClient, gitClient }: ServerContext) { - // Sandbox properties: - this.sandboxId = sandboxId; - this.type = type; - 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 template and timeout - const templateTypes = ["vanillajs", "reactjs", "nextjs", "streamlit"]; - const template = templateTypes.includes(this.type) - ? `gitwit-${this.type}` - : `base`; - this.container = await E2BSandbox.create(template, { - timeoutMs: CONTAINER_TIMEOUT, - }) - } + // 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 template and timeout + const templateTypes = ["vanillajs", "reactjs", "nextjs", "streamlit"] + const template = templateTypes.includes(this.type) + ? `gitwit-${this.type}` + : `base` + this.container = await E2BSandbox.create(template, { + timeoutMs: CONTAINER_TIMEOUT, }) - // Ensure a container was successfully created - if (!this.container) throw new Error("Failed to create container") + } + }) + // 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() - } + // 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}`) } - // 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; + // 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) + } } - handlers(connection: { userId: string, isOwner: boolean, socket: Socket }) { + // Handle getting a file + const handleGetFile: SocketHandler = ({ fileId }: any) => { + return this.fileManager?.getFile(fileId) + } - // 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 folder + const handleGetFolder: SocketHandler = ({ folderId }: any) => { + return this.fileManager?.getFolder(folderId) + } - // Handle getting a file - const handleGetFile: SocketHandler = ({ fileId }: any) => { - return this.fileManager?.getFile(fileId) - } + // Handle saving a file + const handleSaveFile: SocketHandler = async ({ fileId, body }: any) => { + await saveFileRL.consume(connection.userId, 1) + return this.fileManager?.saveFile(fileId, body) + } - // Handle getting a folder - const handleGetFolder: SocketHandler = ({ folderId }: any) => { - return this.fileManager?.getFolder(folderId) - } + // Handle moving a file + const handleMoveFile: SocketHandler = ({ fileId, folderId }: any) => { + return this.fileManager?.moveFile(fileId, 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 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 moving a file - const handleMoveFile: SocketHandler = ({ fileId, folderId }: any) => { - return this.fileManager?.moveFile(fileId, folderId) - } + // 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 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 creating a file + const handleCreateFile: SocketHandler = async ({ name }: any) => { + await createFileRL.consume(connection.userId, 1) + return { success: await this.fileManager?.createFile(name) } + } - // 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 folder + const handleCreateFolder: SocketHandler = async ({ name }: any) => { + await createFolderRL.consume(connection.userId, 1) + return { success: await this.fileManager?.createFolder(name) } + } - // Handle creating a file - const handleCreateFile: SocketHandler = async ({ name }: any) => { - await createFileRL.consume(connection.userId, 1); - return { "success": await this.fileManager?.createFile(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 creating a folder - const handleCreateFolder: SocketHandler = async ({ name }: any) => { - await createFolderRL.consume(connection.userId, 1); - return { "success": await this.fileManager?.createFolder(name) } - } + // Handle deleting a file + const handleDeleteFile: SocketHandler = async ({ fileId }: any) => { + await deleteFileRL.consume(connection.userId, 1) + return this.fileManager?.deleteFile(fileId) + } - // 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 folder + const handleDeleteFolder: SocketHandler = ({ folderId }: any) => { + return this.fileManager?.deleteFolder(folderId) + } - // 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 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, }) - } - - // 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) - } - - // Handle downloading files by download button - const handleDownloadFiles: SocketHandler = async () => { - if (!this.fileManager) throw Error("No file manager") - - // Get all files with their data through fileManager - const files = this.fileManager.fileData.map((file: TFileData) => ({ - path: file.id.startsWith('/') ? file.id.slice(1) : file.id, - content: file.data - })) - - return { files } - } - - return { - "heartbeat": handleHeartbeat, - "getFile": handleGetFile, - "downloadFiles": handleDownloadFiles, - "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, - }; - + const port = extractPortNumber(responseString) + if (port) { + connection.socket.emit( + "previewURL", + "https://" + this.container?.getHost(port) + ) + } + } + ) + }) } -} \ No newline at end of file + // 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 + ) + } + + // Handle downloading files by download button + const handleDownloadFiles: SocketHandler = async () => { + if (!this.fileManager) throw Error("No file manager") + + // Get all files with their data through fileManager + const files = this.fileManager.fileData.map((file: TFileData) => ({ + path: file.id.startsWith("/") ? file.id.slice(1) : file.id, + content: file.data, + })) + + return { files } + } + + return { + heartbeat: handleHeartbeat, + getFile: handleGetFile, + downloadFiles: handleDownloadFiles, + 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, + } + } +} diff --git a/backend/server/src/constants.ts b/backend/server/src/constants.ts index dfd5ce3..3c7dcfb 100644 --- a/backend/server/src/constants.ts +++ b/backend/server/src/constants.ts @@ -1,2 +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 +export const CONTAINER_TIMEOUT = 120_000 diff --git a/backend/server/src/index.ts b/backend/server/src/index.ts index e1e5eda..bdde654 100644 --- a/backend/server/src/index.ts +++ b/backend/server/src/index.ts @@ -10,14 +10,14 @@ import { ConnectionManager } from "./ConnectionManager" import { DokkuClient } from "./DokkuClient" import { Sandbox } from "./Sandbox" import { SecureGitClient } from "./SecureGitClient" -import { socketAuth } from "./socketAuth"; // Import the new socketAuth middleware +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}`); -}; + console.error(message, error) + socket.emit("error", `${message} ${error.message ?? error}`) +} // Handle uncaught exceptions process.on("uncaughtException", (error) => { @@ -64,10 +64,10 @@ if (!process.env.DOKKU_KEY) 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 dokkuClient?.connect() @@ -75,9 +75,9 @@ dokkuClient?.connect() 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 @@ -110,21 +110,23 @@ io.on("connection", async (socket) => { try { // Create or retrieve the sandbox manager for the given sandbox ID - const sandbox = sandboxes[data.sandboxId] ?? new Sandbox( - data.sandboxId, - data.type, - { - aiWorker, dokkuClient, gitClient, - } - ) + const sandbox = + sandboxes[data.sandboxId] ?? + new Sandbox(data.sandboxId, data.type, { + aiWorker, + dokkuClient, + gitClient, + }) sandboxes[data.sandboxId] = sandbox // 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); - }); - }; + connections + .connectionsForSandbox(data.sandboxId) + .forEach((socket: Socket) => { + socket.emit("loaded", files) + }) + } // Initialize the sandbox container // The file manager and terminal managers will be set up if they have been closed @@ -134,26 +136,35 @@ io.on("connection", async (socket) => { // 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); + 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) + } } - }); - }); + ) + }) // Handle disconnection event socket.on("disconnect", async () => { try { // Deregister the connection - connections.removeConnectionForSandbox(socket, data.sandboxId, data.isOwner) + 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. @@ -165,20 +176,18 @@ io.on("connection", async (socket) => { ) } } catch (e: any) { - handleErrors("Error disconnecting:", e, socket); + handleErrors("Error disconnecting:", e, socket) } }) - } catch (e: any) { - handleErrors(`Error initializing sandbox ${data.sandboxId}:`, e, socket); + handleErrors(`Error initializing sandbox ${data.sandboxId}:`, e, socket) } - } catch (e: any) { - handleErrors("Error connecting:", e, socket); + 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 index 030a89b..532e462 100644 --- a/backend/server/src/socketAuth.ts +++ b/backend/server/src/socketAuth.ts @@ -4,72 +4,72 @@ import { Sandbox, 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(), - }) + // 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) + 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 + // 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 - 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 - - // Fetch sandbox data from the database - const dbSandbox = await fetch( - `${process.env.DATABASE_WORKER_URL}/api/sandbox?id=${sandboxId}`, - { - headers: { - Authorization: `${process.env.WORKERS_KEY}`, - }, - } - ) - const dbSandboxJSON = (await dbSandbox.json()) as Sandbox - - // Check if user data was retrieved successfully - if (!dbUserJSON) { - next(new Error("DB error.")) - return + // Fetch sandbox data from the database + const dbSandbox = await fetch( + `${process.env.DATABASE_WORKER_URL}/api/sandbox?id=${sandboxId}`, + { + headers: { + Authorization: `${process.env.WORKERS_KEY}`, + }, } + ) + const dbSandboxJSON = (await dbSandbox.json()) as Sandbox - // 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 - ) + // Check if user data was retrieved successfully + if (!dbUserJSON) { + next(new Error("DB error.")) + return + } - // If user doesn't own or have shared access to the sandbox, deny access - if (!sandbox && !sharedSandboxes) { - next(new Error("Invalid credentials.")) - 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 + ) - // Set socket data with user information - socket.data = { - userId, - sandboxId: sandboxId, - isOwner: sandbox !== undefined, - type: dbSandboxJSON.type - } + // If user doesn't own or have shared access to the sandbox, deny access + if (!sandbox && !sharedSandboxes) { + next(new Error("Invalid credentials.")) + return + } - // Allow the connection - next() + // Set socket data with user information + socket.data = { + userId, + sandboxId: sandboxId, + isOwner: sandbox !== undefined, + type: dbSandboxJSON.type, + } + + // Allow the connection + next() } diff --git a/backend/server/src/utils.ts b/backend/server/src/utils.ts index dd33984..5ae1377 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 +}