feat: add proxy max concurrent requests per IP limit

Author: Roberta001

internal/bootstrap/data/setting.go

@@ -238,6 +238,7 @@ func InitialSettings() []model.SettingItem {
 		{Key: conf.TaskCopyThreadsNum, Value: strconv.Itoa(conf.Conf.Tasks.Copy.Workers), Type: conf.TypeNumber, Group: model.TRAFFIC, Flag: model.PRIVATE},
 		{Key: conf.TaskDecompressDownloadThreadsNum, Value: strconv.Itoa(conf.Conf.Tasks.Decompress.Workers), Type: conf.TypeNumber, Group: model.TRAFFIC, Flag: model.PRIVATE},
 		{Key: conf.TaskDecompressUploadThreadsNum, Value: strconv.Itoa(conf.Conf.Tasks.DecompressUpload.Workers), Type: conf.TypeNumber, Group: model.TRAFFIC, Flag: model.PRIVATE},
+		{Key: conf.ProxyMaxConcurrentRequestsPerIP, Value: "-1", Type: conf.TypeNumber, Group: model.TRAFFIC, Flag: model.PRIVATE},
 		{Key: conf.StreamMaxClientDownloadSpeed, Value: "-1", Type: conf.TypeNumber, Group: model.TRAFFIC, Flag: model.PRIVATE},
 		{Key: conf.StreamMaxClientUploadSpeed, Value: "-1", Type: conf.TypeNumber, Group: model.TRAFFIC, Flag: model.PRIVATE},
 		{Key: conf.StreamMaxServerDownloadSpeed, Value: "-1", Type: conf.TypeNumber, Group: model.TRAFFIC, Flag: model.PRIVATE},

internal/conf/const.go

@@ -157,6 +157,7 @@ const (
 	TaskMoveThreadsNum                    = "move_task_threads_num"
 	TaskDecompressDownloadThreadsNum      = "decompress_download_task_threads_num"
 	TaskDecompressUploadThreadsNum        = "decompress_upload_task_threads_num"
+	ProxyMaxConcurrentRequestsPerIP       = "proxy_max_concurrent_requests_per_ip"
 	StreamMaxClientDownloadSpeed          = "max_client_download_speed"
 	StreamMaxClientUploadSpeed            = "max_client_upload_speed"
 	StreamMaxServerDownloadSpeed          = "max_server_download_speed"

server/middlewares/ip_limit.go

@@ -0,0 +1,65 @@
+package middlewares
+
+import (
+	"net/http"
+	"sync"
+
+	"github.com/OpenListTeam/OpenList/v4/internal/conf"
+	"github.com/OpenListTeam/OpenList/v4/internal/setting"
+	"github.com/gin-gonic/gin"
+)
+
+var (
+	proxyIPCounts   = make(map[string]int)
+	proxyIPCountsMu sync.Mutex
+)
+
+// ProxyIPConcurrencyLimit limits the maximum concurrent server proxy requests per IP Address
+func ProxyIPConcurrencyLimit() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		limit := setting.GetInt(conf.ProxyMaxConcurrentRequestsPerIP, -1)
+
+		if limit < 0 {
+			// -1 or less means unlimited
+			c.Next()
+			return
+		}
+
+		if limit == 0 {
+			// 0 means completely disabled
+			c.AbortWithStatusJSON(http.StatusForbidden, gin.H{
+				"code":    http.StatusForbidden,
+				"message": "Proxy is disabled",
+				"data":    nil,
+			})
+			return
+		}
+
+		ip := c.ClientIP()
+
+		proxyIPCountsMu.Lock()
+		count := proxyIPCounts[ip]
+		if count >= limit {
+			proxyIPCountsMu.Unlock()
+			c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
+				"code":    http.StatusTooManyRequests,
+				"message": "Too Many Proxy Requests from this IP",
+				"data":    nil,
+			})
+			return
+		}
+		proxyIPCounts[ip] = count + 1
+		proxyIPCountsMu.Unlock()
+
+		defer func() {
+			proxyIPCountsMu.Lock()
+			proxyIPCounts[ip]--
+			if proxyIPCounts[ip] <= 0 {
+				delete(proxyIPCounts, ip)
+			}
+			proxyIPCountsMu.Unlock()
+		}()
+
+		c.Next()
+	}
+}

server/router.go

@@ -43,17 +43,18 @@ func Init(e *gin.Engine) {
 	S3(g.Group("/s3"))
 
 	downloadLimiter := middlewares.DownloadRateLimiter(stream.ClientDownloadLimit)
+	proxyConcurrencyLimiter := middlewares.ProxyIPConcurrencyLimit()
 	signCheck := middlewares.Down(sign.Verify)
 	g.GET("/d/*path", middlewares.PathParse, signCheck, downloadLimiter, handles.Down)
-	g.GET("/p/*path", middlewares.PathParse, signCheck, downloadLimiter, handles.Proxy)
+	g.GET("/p/*path", middlewares.PathParse, signCheck, downloadLimiter, proxyConcurrencyLimiter, handles.Proxy)
 	g.HEAD("/d/*path", middlewares.PathParse, signCheck, handles.Down)
-	g.HEAD("/p/*path", middlewares.PathParse, signCheck, handles.Proxy)
+	g.HEAD("/p/*path", middlewares.PathParse, signCheck, proxyConcurrencyLimiter, handles.Proxy)
 	archiveSignCheck := middlewares.Down(sign.VerifyArchive)
 	g.GET("/ad/*path", middlewares.PathParse, archiveSignCheck, downloadLimiter, handles.ArchiveDown)
-	g.GET("/ap/*path", middlewares.PathParse, archiveSignCheck, downloadLimiter, handles.ArchiveProxy)
+	g.GET("/ap/*path", middlewares.PathParse, archiveSignCheck, downloadLimiter, proxyConcurrencyLimiter, handles.ArchiveProxy)
 	g.GET("/ae/*path", middlewares.PathParse, archiveSignCheck, downloadLimiter, handles.ArchiveInternalExtract)
 	g.HEAD("/ad/*path", middlewares.PathParse, archiveSignCheck, handles.ArchiveDown)
-	g.HEAD("/ap/*path", middlewares.PathParse, archiveSignCheck, handles.ArchiveProxy)
+	g.HEAD("/ap/*path", middlewares.PathParse, archiveSignCheck, proxyConcurrencyLimiter, handles.ArchiveProxy)
 	g.HEAD("/ae/*path", middlewares.PathParse, archiveSignCheck, handles.ArchiveInternalExtract)
 
 	g.GET("/sd/:sid", middlewares.EmptyPathParse, middlewares.SharingIdParse, downloadLimiter, handles.SharingDown)