From 1d0fe3048074b100df1cc67859b0accf8bd9d10b Mon Sep 17 00:00:00 2001
From: root <root@adhd.sh>
Date: Fri, 8 Aug 2025 15:41:02 +0200
Subject: [PATCH] Upload limit

---
 controllers/upload.go | 46 ++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 45 insertions(+), 1 deletion(-)

diff --git a/controllers/upload.go b/controllers/upload.go
index 05a2cf0..408e8c4 100644
--- a/controllers/upload.go
+++ b/controllers/upload.go
@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"strconv"
 
 	"github.com/gin-gonic/gin"
 )
@@ -25,13 +26,56 @@ func HandleFileUpload(c *gin.Context) {
 	}
 	defer file.Close()
 
+	// Resolve upload directory from environment
+	uploadDir := os.Getenv("UPLOAD_DIR")
+	if uploadDir == "" {
+		uploadDir = "./uploads"
+	}
+
+	// Ensure upload directory exists (so size calc works)
+	if _, err := os.Stat(uploadDir); os.IsNotExist(err) {
+		if err := os.MkdirAll(uploadDir, 0755); err != nil {
+			c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to create upload directory"})
+			return
+		}
+	}
+
+	// Read max storage from env and compute current usage
+	maxStorageEnv := os.Getenv("MAX_STORAGE_GB")
+	if maxStorageEnv == "" {
+		maxStorageEnv = "10"
+	}
+	maxStorageGB, err := strconv.ParseFloat(maxStorageEnv, 64)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "Invalid MAX_STORAGE_GB configuration"})
+		return
+	}
+
+	usedBytes, err := getDirectorySize(uploadDir)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to calculate storage usage"})
+		return
+	}
+	maxBytes := int64(maxStorageGB * 1024 * 1024 * 1024)
+
+	// If storage is already full, reject upload
+	if usedBytes >= maxBytes {
+		c.JSON(http.StatusInsufficientStorage, gin.H{"error": "Storage limit reached. Upload disabled."})
+		return
+	}
+	// If this file would exceed the limit (when size is known), reject upload
+	if header.Size > 0 && usedBytes+header.Size > maxBytes {
+		c.JSON(http.StatusInsufficientStorage, gin.H{"error": "Not enough storage available for this upload."})
+		return
+	}
+
 	fileID, err := generateID()
 	if err != nil {
 		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to generate file ID"})
 		return
 	}
 
-	fileDir := filepath.Join("./uploads", fileID)
+	fileDir := filepath.Join(uploadDir, fileID)
 	if err := os.MkdirAll(fileDir, 0755); err != nil {
 		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to create file directory"})
 		return