From 4b986ffd86beaf650d78040eaf56e6467d1fc5d5 Mon Sep 17 00:00:00 2001
From: nyne <me@nyne.dev>
Date: Wed, 14 May 2025 22:06:22 +0800
Subject: [PATCH] Add download limit.

---
 server/config/config.go |  4 ++++
 server/service/file.go  | 34 +++++++++++++++++++++++++++++++---
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/server/config/config.go b/server/config/config.go
index 4b8e8c3..88110be 100644
--- a/server/config/config.go
+++ b/server/config/config.go
@@ -74,3 +74,7 @@ func MaxFileSize() int64 {
 func AllowRegister() bool {
 	return config.AllowRegister
 }
+
+func MaxDownloadsPerDayForSingleIP() int {
+	return config.MaxDownloadsPerDayForSingleIP
+}
diff --git a/server/service/file.go b/server/service/file.go
index 1ddb862..98437c6 100644
--- a/server/service/file.go
+++ b/server/service/file.go
@@ -9,6 +9,7 @@ import (
 	"os"
 	"path/filepath"
 	"strconv"
+	"sync"
 	"time"
 
 	"github.com/gofiber/fiber/v3/log"
@@ -19,7 +20,24 @@ const (
 	blockSize = 4 * 1024 * 1024 // 4MB
 )
 
-func getUploadingSize(uid uint) int64 {
+var (
+	ipDownloads = sync.Map{}
+)
+
+func init() {
+	go func() {
+		for {
+			// Clean up old IP download records every 24 hours
+			time.Sleep(24 * time.Hour)
+			ipDownloads.Range(func(key, value interface{}) bool {
+				ipDownloads.Delete(key)
+				return true
+			})
+		}
+	}()
+}
+
+func getUploadingSize() int64 {
 	return dao.GetStatistic("uploading_size")
 }
 
@@ -84,7 +102,7 @@ func CreateUploadingFile(uid uint, filename string, description string, fileSize
 		return nil, model.NewRequestError("file size exceeds the limit")
 	}
 
-	currentUploadingSize := getUploadingSize(uid)
+	currentUploadingSize := getUploadingSize()
 	if currentUploadingSize+fileSize > config.MaxUploadingSize() {
 		log.Info("A new uploading file is rejected due to max uploading size limit")
 		return nil, model.NewRequestError("server is busy, please try again later")
@@ -367,7 +385,17 @@ func GetFile(fid string) (*model.FileView, error) {
 	return file.ToView(), nil
 }
 
-func DownloadFile(fid string) (string, string, error) {
+func DownloadFile(ip string, fid string) (string, string, error) {
+	downloads, _ := ipDownloads.Load(ip)
+	if downloads == nil {
+		ipDownloads.Store(ip, 1)
+	} else {
+		count := downloads.(int)
+		if count >= config.MaxDownloadsPerDayForSingleIP() {
+			return "", "", model.NewRequestError("Too many requests, please try again later")
+		}
+		ipDownloads.Store(ip, count+1)
+	}
 	file, err := dao.GetFile(fid)
 	if err != nil {
 		log.Error("failed to get file: ", err)