TNT-Flo/sandbox-web-ide
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: move rate limiting to handler functions

Browse Source
This commit is contained in:
James Murdza 2024-10-24 23:36:04 -06:00
parent 3bc555ca47
commit d3e987b0ab
2 changed files with 23 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
backend/server/src/SocketHandlers.ts
View File

@ -3,6 +3,13 @@ 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"
@ -42,7 +49,8 @@ export const handleGetFolder: SocketHandler = ({ folderId }: any, context: Handl
}
// Handle saving a file
export const handleSaveFile: SocketHandler = ({ fileId, body }: any, context: HandlerContext) => {
export const handleSaveFile: SocketHandler = async ({ fileId, body, userId }: any, context: HandlerContext) => {
await saveFileRL.consume(userId, 1);
return context.fileManager.saveFile(fileId, body)
}
@ -69,22 +77,26 @@ export const handleDeploy: SocketHandler = async ({ sandboxId }: any, context: H
}
// Handle creating a file
export const handleCreateFile: SocketHandler = async ({ name }: any, context: HandlerContext) => {
export 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
export const handleCreateFolder: SocketHandler = async ({ name }: any, context: HandlerContext) => {
export 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
export const handleRenameFile: SocketHandler = ({ fileId, newName }: any, context: HandlerContext) => {
export 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
export const handleDeleteFile: SocketHandler = ({ fileId }: any, context: HandlerContext) => {
export const handleDeleteFile: SocketHandler = async ({ fileId, userId }: any, context: HandlerContext) => {
await deleteFileRL.consume(userId, 1)
return context.fileManager.deleteFile(fileId)
}

23
backend/server/src/index.ts
View File

@ -9,13 +9,6 @@ import { AIWorker } from "./AIWorker"
import { CONTAINER_TIMEOUT } from "./constants"
import { DokkuClient } from "./DokkuClient"
import { FileManager, SandboxFiles } from "./FileManager"
import {
createFileRL,
createFolderRL,
deleteFileRL,
renameFileRL,
saveFileRL,
} from "./ratelimit"
import { SecureGitClient } from "./SecureGitClient"
import { socketAuth } from "./socketAuth"; // Import the new socketAuth middleware
import { handleCloseTerminal, handleCreateFile, handleCreateFolder, handleCreateTerminal, handleDeleteFile, handleDeleteFolder, handleDeploy, handleGenerateCode, handleGetFile, handleGetFolder, handleHeartbeat, handleListApps, handleMoveFile, HandlerContext, handleRenameFile, handleResizeTerminal, handleSaveFile, handleTerminalData } from "./SocketHandlers"
@ -184,15 +177,11 @@ io.on("connection", async (socket) => {
// Helper function to handle socket events with error handling and optional rate limiting
const handleSocketEvent = (
event: string,
handler: any,
rateLimiter: any | null = null
handler: any
) => {
socket.on(event, async (options: any, callback?: (response: any) => void) => {
try {
// Consume rate limiter if provided
if (rateLimiter) {
await rateLimiter.consume(data.userId, 1); // Adjust as needed for the specific rate limiter
}
const response = await handler({ ...options, ...data }, handlerContext)
callback?.(response);
} catch (e: any) {
@ -206,14 +195,14 @@ io.on("connection", async (socket) => {
handleSocketEvent("heartbeat", handleHeartbeat);
handleSocketEvent("getFile", handleGetFile);
handleSocketEvent("getFolder", handleGetFolder);
handleSocketEvent("saveFile", handleSaveFile, saveFileRL);
handleSocketEvent("saveFile", handleSaveFile);
handleSocketEvent("moveFile", handleMoveFile);
handleSocketEvent("list", handleListApps);
handleSocketEvent("deploy", handleDeploy);
handleSocketEvent("createFile", handleCreateFile, createFileRL);
handleSocketEvent("createFolder", handleCreateFolder, createFolderRL);
handleSocketEvent("renameFile", handleRenameFile, renameFileRL);
handleSocketEvent("deleteFile", handleDeleteFile, deleteFileRL);
handleSocketEvent("createFile", handleCreateFile);
handleSocketEvent("createFolder", handleCreateFolder);
handleSocketEvent("renameFile", handleRenameFile);
handleSocketEvent("deleteFile", handleDeleteFile);
handleSocketEvent("deleteFolder", handleDeleteFolder);
handleSocketEvent("createTerminal", handleCreateTerminal);
handleSocketEvent("resizeTerminal", handleResizeTerminal);