feat: limit max connection count (#2701)

Author: BoYanZh

internal/conf/config.go

@@ -47,6 +47,7 @@ type Config struct {
 	TempDir        string    `json:"temp_dir" env:"TEMP_DIR"`
 	BleveDir       string    `json:"bleve_dir" env:"BLEVE_DIR"`
 	Log            LogConfig `json:"log"`
+	MaxConnections int       `json:"max_connections" env:"MAX_CONNECTIONS"`
 }
 
 func DefaultConfig() *Config {
@@ -74,5 +75,6 @@ func DefaultConfig() *Config {
 			MaxBackups: 5,
 			MaxAge:     28,
 		},
+		MaxConnections: 0,
 	}
 }

server/middlewares/limit.go

@@ -0,0 +1,16 @@
+package middlewares
+
+import (
+	"github.com/gin-gonic/gin"
+)
+
+func MaxAllowed(n int) gin.HandlerFunc {
+	sem := make(chan struct{}, n)
+	acquire := func() { sem <- struct{}{} }
+	release := func() { <-sem }
+	return func(c *gin.Context) {
+		acquire()
+		defer release()
+		c.Next()
+	}
+}

server/router.go

@@ -16,6 +16,9 @@ func Init(r *gin.Engine) {
 	common.SecretKey = []byte(conf.Conf.JwtSecret)
 	Cors(r)
 	r.Use(middlewares.StoragesLoaded)
+	if conf.Conf.MaxConnections > 0 {
+		r.Use(middlewares.MaxAllowed(conf.Conf.MaxConnections))
+	}
 	WebDav(r.Group("/dav"))
 
 	r.GET("/favicon.ico", handles.Favicon)