refactor: create sandboxManager class

This commit is contained in:
James Murdza 2024-10-24 23:58:28 -06:00
parent fcc7a836a6
commit 16e0c250d6
3 changed files with 187 additions and 179 deletions

View File

@ -0,0 +1,181 @@
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 { LockManager } from "./utils"
// 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
}
export class SandboxManager {
fileManager: FileManager;
terminalManager: TerminalManager;
container: any;
aiWorker: AIWorker;
dokkuClient: DokkuClient | null;
gitClient: SecureGitClient | null;
lockManager: LockManager;
socket: Socket;
constructor(fileManager: FileManager, terminalManager: TerminalManager, aiWorker: AIWorker, dokkuClient: DokkuClient | null, gitClient: SecureGitClient | null, lockManager: LockManager, sandboxManager: any, socket: Socket) {
this.fileManager = fileManager;
this.terminalManager = terminalManager;
this.aiWorker = aiWorker;
this.dokkuClient = dokkuClient;
this.gitClient = gitClient;
this.lockManager = lockManager;
this.socket = socket;
this.container = sandboxManager;
}
handlers() {
// Handle heartbeat from a socket connection
const handleHeartbeat: SocketHandler = (_: any) => {
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, userId }: any) => {
await saveFileRL.consume(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 ({ sandboxId }: any) => {
if (!this.gitClient) throw Error("Failed to retrieve apps list: No git client")
const fixedFilePaths = this.fileManager.sandboxFiles.fileData.map((file) => ({
...file,
id: file.id.split("/").slice(2).join("/"),
}))
await this.gitClient.pushFiles(fixedFilePaths, sandboxId)
return { success: true }
}
// Handle creating a file
const handleCreateFile: SocketHandler = async ({ name, userId }: any) => {
await createFileRL.consume(userId, 1);
return { "success": await this.fileManager.createFile(name) }
}
// Handle creating a folder
const handleCreateFolder: SocketHandler = async ({ name, userId }: any) => {
await createFolderRL.consume(userId, 1);
return { "success": await this.fileManager.createFolder(name) }
}
// Handle renaming a file
const handleRenameFile: SocketHandler = async ({ fileId, newName, userId }: any) => {
await renameFileRL.consume(userId, 1)
return this.fileManager.renameFile(fileId, newName)
}
// Handle deleting a file
const handleDeleteFile: SocketHandler = async ({ fileId, userId }: any) => {
await deleteFileRL.consume(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, sandboxId }: any) => {
await this.lockManager.acquireLock(sandboxId, async () => {
await this.terminalManager.createTerminal(id, (responseString: string) => {
this.socket.emit("terminalResponse", { id, data: responseString })
const port = extractPortNumber(responseString)
if (port) {
this.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 = ({ userId, fileName, code, line, instructions }: any) => {
return this.aiWorker.generateCode(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,
};
}
}

View File

@ -1,165 +0,0 @@
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 { LockManager } from "./utils"
export interface HandlerContext {
fileManager: FileManager;
terminalManager: TerminalManager;
sandboxManager: any;
aiWorker: AIWorker;
dokkuClient: DokkuClient | null;
gitClient: SecureGitClient | null;
lockManager: LockManager
socket: Socket
}
// 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
}
// Handle heartbeat from a socket connection
const handleHeartbeat: SocketHandler = (_: any, context: HandlerContext) => {
context.sandboxManager.setTimeout(CONTAINER_TIMEOUT)
}
// Handle getting a file
const handleGetFile: SocketHandler = ({ fileId }: any, context: HandlerContext) => {
return context.fileManager.getFile(fileId)
}
// Handle getting a folder
const handleGetFolder: SocketHandler = ({ folderId }: any, context: HandlerContext) => {
return context.fileManager.getFolder(folderId)
}
// Handle saving a file
const handleSaveFile: SocketHandler = async ({ fileId, body, userId }: any, context: HandlerContext) => {
await saveFileRL.consume(userId, 1);
return context.fileManager.saveFile(fileId, body)
}
// Handle moving a file
const handleMoveFile: SocketHandler = ({ fileId, folderId }: any, context: HandlerContext) => {
return context.fileManager.moveFile(fileId, folderId)
}
// Handle listing apps
const handleListApps: SocketHandler = async (_: any, context: HandlerContext) => {
if (!context.dokkuClient) throw Error("Failed to retrieve apps list: No Dokku client")
return { success: true, apps: await context.dokkuClient.listApps() }
}
// Handle deploying code
const handleDeploy: SocketHandler = async ({ sandboxId }: any, context: HandlerContext) => {
if (!context.gitClient) throw Error("Failed to retrieve apps list: No git client")
const fixedFilePaths = context.fileManager.sandboxFiles.fileData.map((file) => ({
...file,
id: file.id.split("/").slice(2).join("/"),
}))
await context.gitClient.pushFiles(fixedFilePaths, sandboxId)
return { success: true }
}
// Handle creating a file
const handleCreateFile: SocketHandler = async ({ name, userId }: any, context: HandlerContext) => {
await createFileRL.consume(userId, 1);
return { "success": await context.fileManager.createFile(name) }
}
// Handle creating a folder
const handleCreateFolder: SocketHandler = async ({ name, userId }: any, context: HandlerContext) => {
await createFolderRL.consume(userId, 1);
return { "success": await context.fileManager.createFolder(name) }
}
// Handle renaming a file
const handleRenameFile: SocketHandler = async ({ fileId, newName, userId }: any, context: HandlerContext) => {
await renameFileRL.consume(userId, 1)
return context.fileManager.renameFile(fileId, newName)
}
// Handle deleting a file
const handleDeleteFile: SocketHandler = async ({ fileId, userId }: any, context: HandlerContext) => {
await deleteFileRL.consume(userId, 1)
return context.fileManager.deleteFile(fileId)
}
// Handle deleting a folder
const handleDeleteFolder: SocketHandler = ({ folderId }: any, context: HandlerContext) => {
return context.fileManager.deleteFolder(folderId)
}
// Handle creating a terminal session
const handleCreateTerminal: SocketHandler = async ({ id, sandboxId }: any, context: HandlerContext) => {
await context.lockManager.acquireLock(sandboxId, async () => {
await context.terminalManager.createTerminal(id, (responseString: string) => {
context.socket.emit("terminalResponse", { id, data: responseString })
const port = extractPortNumber(responseString)
if (port) {
context.socket.emit(
"previewURL",
"https://" + context.sandboxManager.getHost(port)
)
}
})
})
}
// Handle resizing a terminal
const handleResizeTerminal: SocketHandler = ({ dimensions }: any, context: HandlerContext) => {
context.terminalManager.resizeTerminal(dimensions)
}
// Handle sending data to a terminal
const handleTerminalData: SocketHandler = ({ id, data }: any, context: HandlerContext) => {
return context.terminalManager.sendTerminalData(id, data)
}
// Handle closing a terminal
const handleCloseTerminal: SocketHandler = ({ id }: any, context: HandlerContext) => {
return context.terminalManager.closeTerminal(id)
}
// Handle generating code
const handleGenerateCode: SocketHandler = ({ userId, fileName, code, line, instructions }: any, context: HandlerContext) => {
return context.aiWorker.generateCode(userId, fileName, code, line, instructions)
}
// Define a type for SocketHandler functions
type SocketHandler<T = Record<string, any>> = (args: T, context: HandlerContext) => any;
export const eventHandlers = {
"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,
};

View File

@ -9,9 +9,9 @@ import { AIWorker } from "./AIWorker"
import { CONTAINER_TIMEOUT } from "./constants" import { CONTAINER_TIMEOUT } from "./constants"
import { DokkuClient } from "./DokkuClient" import { DokkuClient } from "./DokkuClient"
import { FileManager, SandboxFiles } from "./FileManager" import { FileManager, SandboxFiles } from "./FileManager"
import { SandboxManager } from "./SandboxManager"
import { SecureGitClient } from "./SecureGitClient" import { SecureGitClient } from "./SecureGitClient"
import { socketAuth } from "./socketAuth"; // Import the new socketAuth middleware import { socketAuth } from "./socketAuth"; // Import the new socketAuth middleware
import { eventHandlers, HandlerContext } from "./SocketHandlers"
import { TerminalManager } from "./TerminalManager" import { TerminalManager } from "./TerminalManager"
import { LockManager } from "./utils" import { LockManager } from "./utils"
@ -163,36 +163,28 @@ io.on("connection", async (socket) => {
// Load file list from the file manager into the editor // Load file list from the file manager into the editor
sendLoadedEvent(fileManager.sandboxFiles) sendLoadedEvent(fileManager.sandboxFiles)
const handlerContext: HandlerContext = { const sandboxManager = new SandboxManager(
fileManager, fileManager,
terminalManager, terminalManager,
aiWorker, aiWorker,
dokkuClient, dokkuClient,
gitClient, gitClient,
lockManager, lockManager,
sandboxManager: containers[data.sandboxId], containers[data.sandboxId],
socket socket
} )
// Helper function to handle socket events with error handling and optional rate limiting Object.entries(sandboxManager.handlers()).forEach(([event, handler]) => {
const handleSocketEvent = (
event: string,
handler: any
) => {
socket.on(event, async (options: any, callback?: (response: any) => void) => { socket.on(event, async (options: any, callback?: (response: any) => void) => {
try { try {
// Consume rate limiter if provided // Consume rate limiter if provided
const response = await handler({ ...options, ...data }, handlerContext) const response = await handler({ ...options, ...data })
callback?.(response); callback?.(response);
} catch (e: any) { } catch (e: any) {
console.error(`Error processing event "${event}":`, e); console.error(`Error processing event "${event}":`, e);
socket.emit("error", `Error: ${event}. ${e.message ?? e}`); socket.emit("error", `Error: ${event}. ${e.message ?? e}`);
} }
}); });
};
Object.entries(eventHandlers).forEach(([event, handler]) => {
handleSocketEvent(event, handler);
}); });
socket.on("disconnect", async () => { socket.on("disconnect", async () => {