From 3325b20aed126c02b7625bd41b9427c9191a4e41 Mon Sep 17 00:00:00 2001
From: Ishaan Dey <ishaan1013@gmail.com>
Date: Thu, 9 May 2024 22:32:21 -0700
Subject: [PATCH] add project size ratelimiting

---
 backend/server/src/index.ts                | 13 +++++-
 backend/server/src/utils.ts                |  8 ++++
 backend/storage/src/index.ts               | 18 +++++++-
 frontend/components/editor/sidebar/new.tsx | 49 ++++++++++++----------
 4 files changed, 65 insertions(+), 23 deletions(-)

diff --git a/backend/server/src/index.ts b/backend/server/src/index.ts
index e2be665..f9450b4 100644
--- a/backend/server/src/index.ts
+++ b/backend/server/src/index.ts
@@ -13,6 +13,7 @@ import {
   createFile,
   deleteFile,
   generateCode,
+  getProjectSize,
   getSandboxFiles,
   renameFile,
   saveFile,
@@ -155,8 +156,16 @@ io.on("connection", async (socket) => {
     }
   })
 
-  socket.on("createFile", async (name: string) => {
+  socket.on("createFile", async (name: string, callback) => {
     try {
+
+      const size: number = await getProjectSize(data.sandboxId)
+      // limit is 200mb
+      if (size > 200 * 1024 * 1024) {
+        io.emit("rateLimit", "Rate limited: project size exceeded. Please delete some files.")
+        callback({success: false})
+      }
+
       await createFileRL.consume(data.userId, 1)
 
       const id = `projects/${data.sandboxId}/${name}`
@@ -177,6 +186,8 @@ io.on("connection", async (socket) => {
       })
 
       await createFile(id)
+
+      callback({success: true})
     } catch (e) {
       io.emit("rateLimit", "Rate limited: file creation. Please slow down.")
     }
diff --git a/backend/server/src/utils.ts b/backend/server/src/utils.ts
index 80bd9b8..73a9aba 100644
--- a/backend/server/src/utils.ts
+++ b/backend/server/src/utils.ts
@@ -1,3 +1,4 @@
+import e from "cors"
 import {
   R2FileBody,
   R2Files,
@@ -134,6 +135,13 @@ export const deleteFile = async (fileId: string) => {
   return res.ok
 }
 
+export const getProjectSize = async (id: string) => {
+  const res = await fetch(
+    `https://storage.ishaan1013.workers.dev/api/size?sandboxId=${id}`
+  )
+  return (await res.json()).size
+}
+
 export const generateCode = async ({
   fileName,
   code,
diff --git a/backend/storage/src/index.ts b/backend/storage/src/index.ts
index ac36a8f..5050be1 100644
--- a/backend/storage/src/index.ts
+++ b/backend/storage/src/index.ts
@@ -16,7 +16,23 @@ export default {
 		const path = url.pathname;
 		const method = request.method;
 
-		if (path === '/api') {
+    if (path === '/api/size' && method === 'GET') {
+      const params = url.searchParams;
+      const sandboxId = params.get('sandboxId');
+
+      if (sandboxId) {
+        const res = await env.R2.list({ prefix: `projects/${sandboxId}` });
+
+        // sum up the size of all files
+        let size = 0;
+        for (const file of res.objects) {
+          size += file.size;
+        }
+
+        return new Response(JSON.stringify({ size }), { status: 200 });
+      } else return invalidRequest;
+    }
+		else if (path === '/api') {
 			if (method === 'GET') {
 				const params = url.searchParams;
 				const sandboxId = params.get('sandboxId');
diff --git a/frontend/components/editor/sidebar/new.tsx b/frontend/components/editor/sidebar/new.tsx
index 3f90117..491fc25 100644
--- a/frontend/components/editor/sidebar/new.tsx
+++ b/frontend/components/editor/sidebar/new.tsx
@@ -1,9 +1,9 @@
-"use client"
+"use client";
 
-import { validateName } from "@/lib/utils"
-import Image from "next/image"
-import { useEffect, useRef } from "react"
-import { Socket } from "socket.io-client"
+import { validateName } from "@/lib/utils";
+import Image from "next/image";
+import { useEffect, useRef } from "react";
+import { Socket } from "socket.io-client";
 
 export default function New({
   socket,
@@ -11,31 +11,38 @@ export default function New({
   stopEditing,
   addNew,
 }: {
-  socket: Socket
-  type: "file" | "folder"
-  stopEditing: () => void
-  addNew: (name: string, type: "file" | "folder") => void
+  socket: Socket;
+  type: "file" | "folder";
+  stopEditing: () => void;
+  addNew: (name: string, type: "file" | "folder") => void;
 }) {
-  const inputRef = useRef<HTMLInputElement>(null)
+  const inputRef = useRef<HTMLInputElement>(null);
 
   const createNew = () => {
-    const name = inputRef.current?.value
+    const name = inputRef.current?.value;
 
     if (name) {
-      const valid = validateName(name, "", type)
+      const valid = validateName(name, "", type);
       if (valid.status) {
         if (type === "file") {
-          socket.emit("createFile", name)
+          socket.emit(
+            "createFile",
+            name,
+            ({ success }: { success: boolean }) => {
+              if (success) {
+                addNew(name, type);
+              }
+            }
+          );
         }
-        addNew(name, type)
       }
     }
-    stopEditing()
-  }
+    stopEditing();
+  };
 
   useEffect(() => {
-    inputRef.current?.focus()
-  }, [])
+    inputRef.current?.focus();
+  }, []);
 
   return (
     <div className="w-full flex items-center h-7 px-1 hover:bg-secondary rounded-sm cursor-pointer transition-colors focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring">
@@ -52,8 +59,8 @@ export default function New({
       />
       <form
         onSubmit={(e) => {
-          e.preventDefault()
-          createNew()
+          e.preventDefault();
+          createNew();
         }}
       >
         <input
@@ -63,5 +70,5 @@ export default function New({
         />
       </form>
     </div>
-  )
+  );
 }