add pat rate limiting

Author: pascal-fischer

management/server/http/handler.go

@@ -4,9 +4,13 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"os"
+	"strconv"
+	"time"
 
 	"github.com/gorilla/mux"
 	"github.com/rs/cors"
+	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/management-integrations/integrations"
 
@@ -38,7 +42,12 @@ import (
 	"github.com/netbirdio/netbird/management/server/telemetry"
 )
 
-const apiPrefix = "/api"
+const (
+	apiPrefix              = "/api"
+	rateLimitingEnabledKey = "NB_API_RATE_LIMITING_ENABLED"
+	rateLimitingBurstKey   = "NB_API_RATE_LIMITING_BURST"
+	rateLimitingRPMKey     = "NB_API_RATE_LIMITING_RPM"
+)
 
 // NewAPIHandler creates the Management service HTTP API handler registering all the available endpoints.
 func NewAPIHandler(
@@ -58,11 +67,42 @@ func NewAPIHandler(
 	settingsManager settings.Manager,
 ) (http.Handler, error) {
 
+	var rateLimitingConfig *middleware.RateLimiterConfig
+	if os.Getenv(rateLimitingEnabledKey) == "true" {
+		rpm := 6
+		if v := os.Getenv(rateLimitingRPMKey); v != "" {
+			value, err := strconv.Atoi(v)
+			if err != nil {
+				log.Warnf("parsing %s env var: %v, using default %d", rateLimitingRPMKey, err, rpm)
+			} else {
+				rpm = value
+			}
+		}
+
+		burst := 500
+		if v := os.Getenv(rateLimitingBurstKey); v != "" {
+			value, err := strconv.Atoi(v)
+			if err != nil {
+				log.Warnf("parsing %s env var: %v, using default %d", rateLimitingBurstKey, err, burst)
+			} else {
+				burst = value
+			}
+		}
+
+		rateLimitingConfig = &middleware.RateLimiterConfig{
+			RequestsPerMinute: float64(rpm),
+			Burst:             burst,
+			CleanupInterval:   6 * time.Hour,
+			LimiterTTL:        24 * time.Hour,
+		}
+	}
+
 	authMiddleware := middleware.NewAuthMiddleware(
 		authManager,
 		accountManager.GetAccountIDFromUserAuth,
 		accountManager.SyncUserJWTGroups,
 		accountManager.GetUserFromUserAuth,
+		rateLimitingConfig,
 	)
 
 	corsMiddleware := cors.AllowAll()

management/server/http/middleware/auth_middleware.go

@@ -29,6 +29,7 @@ type AuthMiddleware struct {
 	ensureAccount       EnsureAccountFunc
 	getUserFromUserAuth GetUserFromUserAuthFunc
 	syncUserJWTGroups   SyncUserJWTGroupsFunc
+	rateLimiter         *APIRateLimiter
 }
 
 // NewAuthMiddleware instance constructor
@@ -37,12 +38,19 @@ func NewAuthMiddleware(
 	ensureAccount EnsureAccountFunc,
 	syncUserJWTGroups SyncUserJWTGroupsFunc,
 	getUserFromUserAuth GetUserFromUserAuthFunc,
+	rateLimiterConfig *RateLimiterConfig,
 ) *AuthMiddleware {
+	var rateLimiter *APIRateLimiter
+	if rateLimiterConfig != nil {
+		rateLimiter = NewAPIRateLimiter(rateLimiterConfig)
+	}
+
 	return &AuthMiddleware{
 		authManager:         authManager,
 		ensureAccount:       ensureAccount,
 		syncUserJWTGroups:   syncUserJWTGroups,
 		getUserFromUserAuth: getUserFromUserAuth,
+		rateLimiter:         rateLimiter,
 	}
 }
 
@@ -145,6 +153,12 @@ func (m *AuthMiddleware) checkPATFromRequest(r *http.Request, auth []string) (*h
 		return r, fmt.Errorf("error extracting token: %w", err)
 	}
 
+	if m.rateLimiter != nil {
+		if !m.rateLimiter.Allow(token) {
+			return r, fmt.Errorf("too many requests")
+		}
+	}
+
 	ctx := r.Context()
 	user, pat, accDomain, accCategory, err := m.authManager.GetPATInfo(ctx, token)
 	if err != nil {