chore: format backend server code

This commit is contained in:
James Murdza 2024-11-16 21:35:06 -05:00
parent 87d311a5d1
commit 7ecbd02fef
8 changed files with 448 additions and 400 deletions

View File

@ -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<string, Counter> = {}
// Stores all sockets connected to a given sandbox
private sockets: Record<string, Set<Socket>> = {}
// Counts how many times the owner is connected to a sandbox
private ownerConnections: Record<string, Counter> = {}
// Stores all sockets connected to a given sandbox
private sockets: Record<string, Set<Socket>> = {}
// Checks if the owner of a sandbox is connected
ownerIsConnected(sandboxId: string): boolean {
return this.ownerConnections[sandboxId]?.getValue() > 0
// 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()
}
}
// Removes a connection for a sandbox
removeConnectionForSandbox(socket: Socket, sandboxId: string, isOwner: boolean) {
this.sockets[sandboxId]?.delete(socket)
// If the connection being removed is for the owner, decrements the owner connection counter
if (isOwner) {
this.ownerConnections[sandboxId]?.decrement()
}
}
// Returns the set of sockets connected to a given sandbox
connectionsForSandbox(sandboxId: string): Set<Socket> {
return this.sockets[sandboxId] ?? new Set();
// If the connection being removed is for the owner, decrements the owner connection counter
if (isOwner) {
this.ownerConnections[sandboxId]?.decrement()
}
}
// Returns the set of sockets connected to a given sandbox
connectionsForSandbox(sandboxId: string): Set<Socket> {
return this.sockets[sandboxId] ?? new Set()
}
}

View File

@ -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<string[]> {
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))
})
)

View File

@ -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

View File

@ -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<T = Record<string, any>> = (args: T) => any;
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
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)
)
}
}
)
})
}
// 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,
}
}
}

View File

@ -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,16 +176,14 @@ 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)
}
})

View File

@ -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()
}