Merge pull request #80 from mysteriumnetwork/limit-max-requests

Limit number of allowed parallel requests

cmd/main.go

@@ -48,6 +48,8 @@ func main() {
 	r.Use(mlog.GinLogFunc())
 	r.Use(apierror.ErrorHandler)
 
+	r.Use(LimitMiddleware(cfg.MaxRequestsLimit))
+
 	r.GET("/", func(c *gin.Context) {
 		c.Redirect(http.StatusMovedPermanently, "/swagger/index.html")
 	})
@@ -100,3 +102,16 @@ func printBanner() {
 	log.Info().Msg(" and it will be our privilege to make that future bright.")
 	log.Info().Msg(strings.Repeat("▱", 60))
 }
+
+func LimitMiddleware(size int) gin.HandlerFunc {
+	limit := make(chan struct{}, size)
+	return func(c *gin.Context) {
+		select {
+		case limit <- struct{}{}:
+			c.Next()
+			<-limit
+		default:
+			c.AbortWithStatus(http.StatusTooManyRequests)
+		}
+	}
+}

config/options.go

@@ -31,6 +31,8 @@ type Options struct {
 
 	UniverseJWTSecret string
 	SentinelURL       string
+
+	MaxRequestsLimit int
 }
 
 func ReadDiscovery() (*Options, error) {
@@ -58,6 +60,12 @@ func ReadDiscovery() (*Options, error) {
 		return nil, err
 	}
 
+	maxRequestsLimit := OptionalEnv("MAX_REQUESTS_LIMIT", "1000")
+	limit, err := strconv.Atoi(maxRequestsLimit)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse max requests limit: %w", err)
+	}
+
 	return &Options{
 		QualityOracleURL: *qualityOracleURL,
 		QualityCacheTTL:  *qualityCacheTTL,
@@ -66,6 +74,7 @@ func ReadDiscovery() (*Options, error) {
 		LocationAddress:  *locationAddress,
 		LocationUser:     locationUser,
 		LocationPass:     locationPass,
+		MaxRequestsLimit: limit,
 	}, nil
 }