sandbox-web-ide/backend/server/src/Sandbox.ts

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, TFileData, 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;
    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:
        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,
                })
            }
        })
        // 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)
        }

        // 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,
        };

    }

}