feat: rate-limit

Author: ijo42

go.mod

@@ -1,3 +1,5 @@
 module github.com/psumaps/icalmiddleware
 
 go 1.19
+
+require golang.org/x/time v0.10.0

go.sum

@@ -0,0 +1,2 @@
+golang.org/x/time v0.10.0 h1:3usCWA8tQn0L8+hFJQNgzpWbd89begxN66o1Ojdn5L4=
+golang.org/x/time v0.10.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=

main.go

@@ -8,24 +8,33 @@ import (
 	"net/http"
 	"net/netip"
 	"strings"
+	"sync"
 	"time"
+
+	"golang.org/x/time/rate"
 )
 
 type Config struct {
-	ForwardToken bool     `json:"forwardToken,omitempty"`
-	Freshness    int64    `json:"freshness,omitempty"`
-	HeaderName   string   `json:"headerName,omitempty"`
-	AllowSubnet  []string `json:"allowSubnet,omitempty"`
-	Timeout      int64    `json:"timeout,omitempty"`
+	ForwardToken    bool     `json:"forwardToken,omitempty"`
+	Freshness       int64    `json:"freshness,omitempty"`
+	HeaderName      string   `json:"headerName,omitempty"`
+	AllowSubnet     []string `json:"allowSubnet,omitempty"`
+	Timeout         int64    `json:"timeout,omitempty"`
+	GlobalRateLimit int      `json:"globalRateLimit,omitempty"`
+	PerIPRateLimit  int      `json:"perIPRateLimit,omitempty"`
+	RateLimitWindow int64    `json:"rateLimitWindow,omitempty"`
 }
 
 func CreateConfig() *Config {
 	return &Config{
-		HeaderName:   "Authorization",
-		ForwardToken: false,
-		Freshness:    3600,
-		AllowSubnet:  []string{"0.0.0.0/24"},
-		Timeout:      5,
+		HeaderName:      "Authorization",
+		ForwardToken:    false,
+		Freshness:       3600,
+		AllowSubnet:     []string{"0.0.0.0/24"},
+		Timeout:         10,
+		GlobalRateLimit: 100,
+		PerIPRateLimit:  10,
+		RateLimitWindow: 60,
 	}
 }
 
@@ -38,6 +47,13 @@ type ICalMiddleware struct {
 	allowSubnet  []netip.Prefix
 	timeout      time.Duration
 	name         string
+
+	// Поля для rate limiting
+	globalLimiter   *rate.Limiter
+	ipLimiters      map[string]*rate.Limiter
+	ipLimiterMutex  sync.Mutex
+	perIPRateLimit  int
+	rateLimitWindow time.Duration
 }
 
 func New(_ context.Context, next http.Handler, config *Config, name string) (http.Handler, error) {
@@ -60,15 +76,25 @@ func New(_ context.Context, next http.Handler, config *Config, name string) (htt
 	}
 	cache := NewCache(time.Duration(config.Freshness)*time.Second, 8*time.Hour)
 
+	var globalLimiter *rate.Limiter
+	if config.GlobalRateLimit > 0 && config.RateLimitWindow > 0 {
+		ratePerSec := float64(config.GlobalRateLimit) / float64(config.RateLimitWindow)
+		globalLimiter = rate.NewLimiter(rate.Limit(ratePerSec), config.GlobalRateLimit)
+	}
+
 	return &ICalMiddleware{
-		headerName:   config.HeaderName,
-		forwardToken: config.ForwardToken,
-		freshness:    config.Freshness,
-		allowSubnet:  cidrs,
-		next:         next,
-		cache:        cache,
-		timeout:      timeout,
-		name:         name,
+		headerName:      config.HeaderName,
+		forwardToken:    config.ForwardToken,
+		freshness:       config.Freshness,
+		allowSubnet:     cidrs,
+		next:            next,
+		cache:           cache,
+		timeout:         timeout,
+		name:            name,
+		globalLimiter:   globalLimiter,
+		ipLimiters:      make(map[string]*rate.Limiter),
+		perIPRateLimit:  config.PerIPRateLimit,
+		rateLimitWindow: time.Duration(config.RateLimitWindow) * time.Second,
 	}, nil
 }
 
@@ -110,7 +136,6 @@ func (plugin *ICalMiddleware) httpRequestAndCache(url string) error {
 	return nil
 }
 
-
 func (plugin *ICalMiddleware) extractTokenFromHeader(request *http.Request) string {
 	canonicalHeaderName := http.CanonicalHeaderKey(plugin.headerName)
 	token := request.Header.Get(canonicalHeaderName)
@@ -130,14 +155,14 @@ func (plugin *ICalMiddleware) extractTokenFromHeader(request *http.Request) stri
 }
 
 func ReadUserIP(r *http.Request) string {
-	IPAddress := r.Header.Get("X-Real-Ip")
-	if IPAddress == "" {
-		IPAddress = r.Header.Get("X-Forwarded-For")
+	ipAddress := r.Header.Get("X-Real-Ip")
+	if ipAddress == "" {
+		ipAddress = r.Header.Get("X-Forwarded-For")
 	}
-	if IPAddress == "" {
-		IPAddress, _, _ = net.SplitHostPort(r.RemoteAddr)
+	if ipAddress == "" {
+		ipAddress, _, _ = net.SplitHostPort(r.RemoteAddr)
 	}
-	return IPAddress
+	return ipAddress
 }
 
 func (plugin *ICalMiddleware) containsSubnet(address string) bool {
@@ -161,29 +186,51 @@ func (plugin *ICalMiddleware) validate(request *http.Request) (int, error) {
 	userIP := ReadUserIP(request)
 	fmt.Printf("[DEBUG] [%s] Обработка запроса от IP: %s, URL: %s\n", plugin.name, userIP, request.URL.String())
 
-	if !plugin.containsSubnet(userIP) {
-		token := plugin.extractTokenFromHeader(request)
-		if token == "" {
-			fmt.Printf("[ERROR] [%s] Токен не предоставлен в заголовке '%s' для запроса от IP %s\n", plugin.name, plugin.headerName, userIP)
-			return http.StatusUnauthorized, fmt.Errorf("no token provided")
+	if plugin.containsSubnet(userIP) {
+		fmt.Printf("[DEBUG] [%s] IP %s входит в разрешённую подсет\n", plugin.name, userIP)
+		return http.StatusOK, nil
+	}
+
+	if plugin.globalLimiter != nil && !plugin.globalLimiter.Allow() {
+		fmt.Printf("[ERROR] [%s] Превышен глобальный лимит запросов\n", plugin.name)
+		return http.StatusTooManyRequests, fmt.Errorf("global rate limit exceeded")
+	}
+
+	if plugin.perIPRateLimit > 0 && plugin.rateLimitWindow > 0 {
+		plugin.ipLimiterMutex.Lock()
+		limiter, exists := plugin.ipLimiters[userIP]
+		if !exists {
+			ratePerSec := float64(plugin.perIPRateLimit) / plugin.rateLimitWindow.Seconds()
+			limiter = rate.NewLimiter(rate.Limit(ratePerSec), plugin.perIPRateLimit)
+			plugin.ipLimiters[userIP] = limiter
 		}
-		if len(token) != 16 {
-			fmt.Printf("[ERROR] [%s] Неверная длина токена '%s' для запроса от IP %s\n", plugin.name, token, userIP)
-			return http.StatusUnauthorized, fmt.Errorf("incorrect token len")
+		plugin.ipLimiterMutex.Unlock()
+		if !limiter.Allow() {
+			fmt.Printf("[ERROR] [%s] Превышен лимит запросов для IP: %s\n", plugin.name, userIP)
+			return http.StatusTooManyRequests, fmt.Errorf("rate limit exceeded for IP %s", userIP)
 		}
-		if !plugin.cache.Has(token) {
-			err := plugin.httpRequestAndCache(token)
-			if err != nil {
-				fmt.Printf("[ERROR] [%s] Проверка токена '%s' не пройдена для запроса от IP %s: %v\n", plugin.name, token, userIP, err)
-				return http.StatusUnauthorized, err
-			}
-			fmt.Printf("[DEBUG] [%s] Токен '%s' валидирован и кэширован для IP %s\n", plugin.name, token, userIP)
-		} else {
-			fmt.Printf("[DEBUG] [%s] Токен '%s' найден в кэше для IP %s\n", plugin.name, token, userIP)
+	}
+
+	token := plugin.extractTokenFromHeader(request)
+	if token == "" {
+		fmt.Printf("[ERROR] [%s] Токен не предоставлен в заголовке '%s' для запроса от IP %s\n", plugin.name, plugin.headerName, userIP)
+		return http.StatusUnauthorized, fmt.Errorf("no token provided")
+	}
+	if len(token) != 16 {
+		fmt.Printf("[ERROR] [%s] Неверная длина токена '%s' для запроса от IP %s\n", plugin.name, token, userIP)
+		return http.StatusUnauthorized, fmt.Errorf("incorrect token len")
+	}
+	if !plugin.cache.Has(token) {
+		err := plugin.httpRequestAndCache(token)
+		if err != nil {
+			fmt.Printf("[ERROR] [%s] Проверка токена '%s' не пройдена для запроса от IP %s: %v\n", plugin.name, token, userIP, err)
+			return http.StatusUnauthorized, err
 		}
+		fmt.Printf("[DEBUG] [%s] Токен '%s' валидирован и кэширован для IP %s\n", plugin.name, token, userIP)
 	} else {
-		fmt.Printf("[DEBUG] [%s] Запрос от IP %s пропущен без проверки токена (разрешённая подсеть)\n", plugin.name, userIP)
+		fmt.Printf("[DEBUG] [%s] Токен '%s' найден в кэше для IP %s\n", plugin.name, token, userIP)
 	}
+
 	return http.StatusOK, nil
 }