From 09ab81f5bd702c1275d75a78476d611d2b6d5b9c Mon Sep 17 00:00:00 2001
From: James Murdza <james@jamesmurdza.com>
Date: Thu, 24 Oct 2024 23:36:04 -0600
Subject: [PATCH] refactor: move rate limiting to handler functions

---
 backend/server/src/SocketHandlers.ts | 22 +++++++++++++++++-----
 backend/server/src/index.ts          | 23 ++++++-----------------
 2 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/backend/server/src/SocketHandlers.ts b/backend/server/src/SocketHandlers.ts
index 9169dc8..5b2f663 100644
--- a/backend/server/src/SocketHandlers.ts
+++ b/backend/server/src/SocketHandlers.ts
@@ -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)
 }
 
diff --git a/backend/server/src/index.ts b/backend/server/src/index.ts
index 4850077..bda426d 100644
--- a/backend/server/src/index.ts
+++ b/backend/server/src/index.ts
@@ -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);