[management] move rate limiter to shared package (#7727)

shared/ratelimit/rate_limiter.go

@@ -0,0 +1,306 @@
+package ratelimit
+
+import (
+	"context"
+	"encoding/json"
+	"net"
+	"net/http"
+	"os"
+	"strconv"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"golang.org/x/time/rate"
+
+	"github.com/netbirdio/netbird/trustedproxy"
+)
+
+const (
+	RateLimitingEnabledEnv        = "NB_API_RATE_LIMITING_ENABLED"
+	RateLimitingBurstEnv          = "NB_API_RATE_LIMITING_BURST"
+	RateLimitingRPMEnv            = "NB_API_RATE_LIMITING_RPM"
+	RateLimitingTrustedProxiesEnv = "NB_API_RATE_LIMITING_TRUSTED_PROXIES"
+
+	defaultAPIRPM   = 6
+	defaultAPIBurst = 500
+)
+
+// RateLimiterConfig holds configuration for the API rate limiter
+type RateLimiterConfig struct {
+	// RequestsPerMinute defines the rate at which tokens are replenished
+	RequestsPerMinute float64
+	// Burst defines the maximum number of requests that can be made in a burst
+	Burst int
+	// CleanupInterval defines how often to clean up old limiters (how often garbage collection runs)
+	CleanupInterval time.Duration
+	// LimiterTTL defines how long a limiter should be kept after last use (age threshold for removal)
+	LimiterTTL time.Duration
+	// TrustedProxies lists the upstream proxies whose forwarding headers may be
+	// believed. Empty means requests are keyed by their direct peer address.
+	TrustedProxies *trustedproxy.List
+}
+
+// DefaultRateLimiterConfig returns a default configuration
+func DefaultRateLimiterConfig() *RateLimiterConfig {
+	return &RateLimiterConfig{
+		RequestsPerMinute: 100,
+		Burst:             120,
+		CleanupInterval:   5 * time.Minute,
+		LimiterTTL:        10 * time.Minute,
+	}
+}
+
+func RateLimiterConfigFromEnv() (cfg *RateLimiterConfig, enabled bool) {
+	rpm := defaultAPIRPM
+	if v := os.Getenv(RateLimitingRPMEnv); v != "" {
+		value, err := strconv.Atoi(v)
+		if err != nil {
+			log.Warnf("parsing %s env var: %v, using default %d", RateLimitingRPMEnv, err, rpm)
+		} else {
+			rpm = value
+		}
+	}
+	if rpm <= 0 {
+		log.Warnf("%s=%d is non-positive, using default %d", RateLimitingRPMEnv, rpm, defaultAPIRPM)
+		rpm = defaultAPIRPM
+	}
+
+	burst := defaultAPIBurst
+	if v := os.Getenv(RateLimitingBurstEnv); v != "" {
+		value, err := strconv.Atoi(v)
+		if err != nil {
+			log.Warnf("parsing %s env var: %v, using default %d", RateLimitingBurstEnv, err, burst)
+		} else {
+			burst = value
+		}
+	}
+	if burst <= 0 {
+		log.Warnf("%s=%d is non-positive, using default %d", RateLimitingBurstEnv, burst, defaultAPIBurst)
+		burst = defaultAPIBurst
+	}
+
+	trusted, err := trustedproxy.Parse(os.Getenv(RateLimitingTrustedProxiesEnv))
+	if err != nil {
+		log.Warnf("parsing %s env var: %v, trusting no proxies", RateLimitingTrustedProxiesEnv, err)
+		trusted = nil
+	}
+
+	return &RateLimiterConfig{
+		RequestsPerMinute: float64(rpm),
+		Burst:             burst,
+		CleanupInterval:   6 * time.Hour,
+		LimiterTTL:        24 * time.Hour,
+		TrustedProxies:    trusted,
+	}, os.Getenv(RateLimitingEnabledEnv) == "true"
+}
+
+// limiterEntry holds a rate limiter and its last access time
+type limiterEntry struct {
+	limiter    *rate.Limiter
+	lastAccess time.Time
+}
+
+// APIRateLimiter manages rate limiting for API tokens
+type APIRateLimiter struct {
+	config   *RateLimiterConfig
+	limiters map[string]*limiterEntry
+	mu       sync.RWMutex
+	stopChan chan struct{}
+	enabled  atomic.Bool
+}
+
+// NewAPIRateLimiter creates a new API rate limiter with the given configuration
+func NewAPIRateLimiter(config *RateLimiterConfig) *APIRateLimiter {
+	if config == nil {
+		config = DefaultRateLimiterConfig()
+	}
+
+	rl := &APIRateLimiter{
+		config:   config,
+		limiters: make(map[string]*limiterEntry),
+		stopChan: make(chan struct{}),
+	}
+	rl.enabled.Store(true)
+
+	go rl.cleanupLoop()
+
+	return rl
+}
+
+func (rl *APIRateLimiter) SetEnabled(enabled bool) {
+	rl.enabled.Store(enabled)
+}
+
+func (rl *APIRateLimiter) Enabled() bool {
+	return rl.enabled.Load()
+}
+
+func (rl *APIRateLimiter) UpdateConfig(config *RateLimiterConfig) {
+	if config == nil {
+		return
+	}
+	if config.RequestsPerMinute <= 0 || config.Burst <= 0 {
+		log.Warnf("UpdateConfig: ignoring invalid rpm=%v burst=%d", config.RequestsPerMinute, config.Burst)
+		return
+	}
+
+	newRPS := rate.Limit(config.RequestsPerMinute / 60.0)
+	newBurst := config.Burst
+
+	rl.mu.Lock()
+	rl.config.RequestsPerMinute = config.RequestsPerMinute
+	rl.config.Burst = newBurst
+	snapshot := make([]*rate.Limiter, 0, len(rl.limiters))
+	for _, entry := range rl.limiters {
+		snapshot = append(snapshot, entry.limiter)
+	}
+	rl.mu.Unlock()
+
+	for _, l := range snapshot {
+		l.SetLimit(newRPS)
+		l.SetBurst(newBurst)
+	}
+}
+
+// Allow checks if a request for the given key (token) is allowed
+func (rl *APIRateLimiter) Allow(key string) bool {
+	if !rl.enabled.Load() {
+		return true
+	}
+	limiter := rl.getLimiter(key)
+	return limiter.Allow()
+}
+
+// Wait blocks until the rate limiter allows another request for the given key
+// Returns an error if the context is canceled
+func (rl *APIRateLimiter) Wait(ctx context.Context, key string) error {
+	if !rl.enabled.Load() {
+		return nil
+	}
+	limiter := rl.getLimiter(key)
+	return limiter.Wait(ctx)
+}
+
+// getLimiter retrieves or creates a rate limiter for the given key
+func (rl *APIRateLimiter) getLimiter(key string) *rate.Limiter {
+	rl.mu.RLock()
+	entry, exists := rl.limiters[key]
+	rl.mu.RUnlock()
+
+	if exists {
+		rl.mu.Lock()
+		entry.lastAccess = time.Now()
+		rl.mu.Unlock()
+		return entry.limiter
+	}
+
+	rl.mu.Lock()
+	defer rl.mu.Unlock()
+
+	if entry, exists := rl.limiters[key]; exists {
+		entry.lastAccess = time.Now()
+		return entry.limiter
+	}
+
+	requestsPerSecond := rl.config.RequestsPerMinute / 60.0
+	limiter := rate.NewLimiter(rate.Limit(requestsPerSecond), rl.config.Burst)
+	rl.limiters[key] = &limiterEntry{
+		limiter:    limiter,
+		lastAccess: time.Now(),
+	}
+
+	return limiter
+}
+
+// cleanupLoop periodically removes old limiters that haven't been used recently
+func (rl *APIRateLimiter) cleanupLoop() {
+	ticker := time.NewTicker(rl.config.CleanupInterval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			rl.cleanup()
+		case <-rl.stopChan:
+			return
+		}
+	}
+}
+
+// cleanup removes limiters that haven't been used within the TTL period
+func (rl *APIRateLimiter) cleanup() {
+	rl.mu.Lock()
+	defer rl.mu.Unlock()
+
+	now := time.Now()
+	for key, entry := range rl.limiters {
+		if now.Sub(entry.lastAccess) > rl.config.LimiterTTL {
+			delete(rl.limiters, key)
+		}
+	}
+}
+
+// Stop stops the cleanup goroutine
+func (rl *APIRateLimiter) Stop() {
+	close(rl.stopChan)
+}
+
+// Reset removes the rate limiter for a specific key
+func (rl *APIRateLimiter) Reset(key string) {
+	rl.mu.Lock()
+	defer rl.mu.Unlock()
+	delete(rl.limiters, key)
+}
+
+// Middleware returns an HTTP middleware that rate limits requests by client IP.
+// Returns 429 Too Many Requests if the rate limit is exceeded.
+func (rl *APIRateLimiter) Middleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if !rl.enabled.Load() {
+			next.ServeHTTP(w, r)
+			return
+		}
+		clientIP := getClientIP(r, rl.config.TrustedProxies)
+		if !rl.Allow(clientIP) {
+			writeTooManyRequests(w)
+			return
+		}
+		next.ServeHTTP(w, r)
+	})
+}
+
+// errorResponse is the JSON body of a rejected request.
+type errorResponse struct {
+	Message string `json:"message"`
+	Code    int    `json:"code"`
+}
+
+// writeTooManyRequests writes a JSON error response with status 429 Too Many Requests.
+func writeTooManyRequests(w http.ResponseWriter) {
+	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
+	w.WriteHeader(http.StatusTooManyRequests)
+	if err := json.NewEncoder(w).Encode(errorResponse{
+		Message: "rate limit exceeded, please try again later",
+		Code:    http.StatusTooManyRequests,
+	}); err != nil {
+		log.Debugf("writing rate limit response: %v", err)
+	}
+}
+
+// getClientIP extracts the client IP address from the request. Forwarding headers
+// are used only when the request arrives from a trusted proxy.
+func getClientIP(r *http.Request, trusted *trustedproxy.List) string {
+	if !trusted.Empty() {
+		if addr := trusted.ResolveClientIP(r.RemoteAddr, r.Header.Get("X-Forwarded-For")); addr.IsValid() {
+			return addr.String()
+		}
+	}
+
+	ip, _, err := net.SplitHostPort(r.RemoteAddr)
+	if err != nil {
+		return r.RemoteAddr
+	}
+	return ip
+}

shared/ratelimit/rate_limiter_test.go

@@ -0,0 +1,378 @@
+package ratelimit
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/trustedproxy"
+)
+
+func TestAPIRateLimiter_Allow(t *testing.T) {
+	rl := NewAPIRateLimiter(&RateLimiterConfig{
+		RequestsPerMinute: 60, // 1 per second
+		Burst:             2,
+		CleanupInterval:   time.Minute,
+		LimiterTTL:        time.Minute,
+	})
+	defer rl.Stop()
+
+	// First two requests should be allowed (burst)
+	assert.True(t, rl.Allow("test-key"))
+	assert.True(t, rl.Allow("test-key"))
+
+	// Third request should be denied (exceeded burst)
+	assert.False(t, rl.Allow("test-key"))
+
+	// Different key should be allowed
+	assert.True(t, rl.Allow("different-key"))
+}
+
+func TestAPIRateLimiter_Middleware(t *testing.T) {
+	rl := NewAPIRateLimiter(&RateLimiterConfig{
+		RequestsPerMinute: 60, // 1 per second
+		Burst:             2,
+		CleanupInterval:   time.Minute,
+		LimiterTTL:        time.Minute,
+	})
+	defer rl.Stop()
+
+	// Create a simple handler that returns 200 OK
+	nextHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	})
+
+	// Wrap with rate limiter middleware
+	handler := rl.Middleware(nextHandler)
+
+	// First two requests should pass (burst)
+	for i := 0; i < 2; i++ {
+		req := httptest.NewRequest(http.MethodGet, "/test", nil)
+		req.RemoteAddr = "192.168.1.1:12345"
+		rr := httptest.NewRecorder()
+		handler.ServeHTTP(rr, req)
+		assert.Equal(t, http.StatusOK, rr.Code, "request %d should be allowed", i+1)
+	}
+
+	// Third request should be rate limited
+	req := httptest.NewRequest(http.MethodGet, "/test", nil)
+	req.RemoteAddr = "192.168.1.1:12345"
+	rr := httptest.NewRecorder()
+	handler.ServeHTTP(rr, req)
+	assert.Equal(t, http.StatusTooManyRequests, rr.Code)
+	assert.Equal(t, "application/json; charset=UTF-8", rr.Header().Get("Content-Type"))
+	assert.JSONEq(t, `{"message":"rate limit exceeded, please try again later","code":429}`, rr.Body.String())
+}
+
+func TestAPIRateLimiter_Middleware_DifferentIPs(t *testing.T) {
+	rl := NewAPIRateLimiter(&RateLimiterConfig{
+		RequestsPerMinute: 60,
+		Burst:             1,
+		CleanupInterval:   time.Minute,
+		LimiterTTL:        time.Minute,
+	})
+	defer rl.Stop()
+
+	nextHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	})
+
+	handler := rl.Middleware(nextHandler)
+
+	// Request from first IP
+	req1 := httptest.NewRequest(http.MethodGet, "/test", nil)
+	req1.RemoteAddr = "192.168.1.1:12345"
+	rr1 := httptest.NewRecorder()
+	handler.ServeHTTP(rr1, req1)
+	assert.Equal(t, http.StatusOK, rr1.Code)
+
+	// Second request from first IP should be rate limited
+	req2 := httptest.NewRequest(http.MethodGet, "/test", nil)
+	req2.RemoteAddr = "192.168.1.1:12345"
+	rr2 := httptest.NewRecorder()
+	handler.ServeHTTP(rr2, req2)
+	assert.Equal(t, http.StatusTooManyRequests, rr2.Code)
+
+	// Request from different IP should be allowed
+	req3 := httptest.NewRequest(http.MethodGet, "/test", nil)
+	req3.RemoteAddr = "192.168.1.2:12345"
+	rr3 := httptest.NewRecorder()
+	handler.ServeHTTP(rr3, req3)
+	assert.Equal(t, http.StatusOK, rr3.Code)
+}
+
+func TestGetClientIP(t *testing.T) {
+	tests := []struct {
+		name       string
+		remoteAddr string
+		expected   string
+	}{
+		{
+			name:       "remote addr with port",
+			remoteAddr: "192.168.1.1:12345",
+			expected:   "192.168.1.1",
+		},
+		{
+			name:       "remote addr without port",
+			remoteAddr: "192.168.1.1",
+			expected:   "192.168.1.1",
+		},
+		{
+			name:       "IPv6 with port",
+			remoteAddr: "[::1]:12345",
+			expected:   "::1",
+		},
+		{
+			name:       "IPv6 without port",
+			remoteAddr: "::1",
+			expected:   "::1",
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			req := httptest.NewRequest(http.MethodGet, "/test", nil)
+			req.RemoteAddr = tc.remoteAddr
+			assert.Equal(t, tc.expected, getClientIP(req, nil))
+		})
+	}
+}
+
+func TestAPIRateLimiter_Reset(t *testing.T) {
+	rl := NewAPIRateLimiter(&RateLimiterConfig{
+		RequestsPerMinute: 60,
+		Burst:             1,
+		CleanupInterval:   time.Minute,
+		LimiterTTL:        time.Minute,
+	})
+	defer rl.Stop()
+
+	// Use up the burst
+	assert.True(t, rl.Allow("test-key"))
+	assert.False(t, rl.Allow("test-key"))
+
+	// Reset the limiter
+	rl.Reset("test-key")
+
+	// Should be allowed again
+	assert.True(t, rl.Allow("test-key"))
+}
+
+func TestAPIRateLimiter_SetEnabled(t *testing.T) {
+	rl := NewAPIRateLimiter(&RateLimiterConfig{
+		RequestsPerMinute: 60,
+		Burst:             1,
+		CleanupInterval:   time.Minute,
+		LimiterTTL:        time.Minute,
+	})
+	defer rl.Stop()
+
+	assert.True(t, rl.Allow("key"))
+	assert.False(t, rl.Allow("key"), "burst exhausted while enabled")
+
+	rl.SetEnabled(false)
+	assert.False(t, rl.Enabled())
+	for i := 0; i < 5; i++ {
+		assert.True(t, rl.Allow("key"), "disabled limiter must always allow")
+	}
+
+	rl.SetEnabled(true)
+	assert.True(t, rl.Enabled())
+	assert.False(t, rl.Allow("key"), "re-enabled limiter retains prior bucket state")
+}
+
+func TestAPIRateLimiter_UpdateConfig(t *testing.T) {
+	rl := NewAPIRateLimiter(&RateLimiterConfig{
+		RequestsPerMinute: 60,
+		Burst:             2,
+		CleanupInterval:   time.Minute,
+		LimiterTTL:        time.Minute,
+	})
+	defer rl.Stop()
+
+	assert.True(t, rl.Allow("k1"))
+	assert.True(t, rl.Allow("k1"))
+	assert.False(t, rl.Allow("k1"), "burst=2 exhausted")
+
+	rl.UpdateConfig(&RateLimiterConfig{
+		RequestsPerMinute: 60,
+		Burst:             10,
+		CleanupInterval:   time.Minute,
+		LimiterTTL:        time.Minute,
+	})
+
+	// New burst applies to existing keys in place; bucket refills up to new burst over time,
+	// but importantly newly-added keys use the updated config immediately.
+	assert.True(t, rl.Allow("k2"))
+	for i := 0; i < 9; i++ {
+		assert.True(t, rl.Allow("k2"))
+	}
+	assert.False(t, rl.Allow("k2"), "new burst=10 exhausted")
+}
+
+func TestAPIRateLimiter_UpdateConfig_NilIgnored(t *testing.T) {
+	rl := NewAPIRateLimiter(&RateLimiterConfig{
+		RequestsPerMinute: 60,
+		Burst:             1,
+		CleanupInterval:   time.Minute,
+		LimiterTTL:        time.Minute,
+	})
+	defer rl.Stop()
+
+	rl.UpdateConfig(nil) // must not panic or zero the config
+
+	assert.True(t, rl.Allow("k"))
+	assert.False(t, rl.Allow("k"))
+}
+
+func TestAPIRateLimiter_UpdateConfig_NonPositiveIgnored(t *testing.T) {
+	rl := NewAPIRateLimiter(&RateLimiterConfig{
+		RequestsPerMinute: 60,
+		Burst:             1,
+		CleanupInterval:   time.Minute,
+		LimiterTTL:        time.Minute,
+	})
+	defer rl.Stop()
+
+	assert.True(t, rl.Allow("k"))
+	assert.False(t, rl.Allow("k"))
+
+	rl.UpdateConfig(&RateLimiterConfig{RequestsPerMinute: 0, Burst: 0, CleanupInterval: time.Minute, LimiterTTL: time.Minute})
+	rl.UpdateConfig(&RateLimiterConfig{RequestsPerMinute: -1, Burst: 5, CleanupInterval: time.Minute, LimiterTTL: time.Minute})
+	rl.UpdateConfig(&RateLimiterConfig{RequestsPerMinute: 60, Burst: -1, CleanupInterval: time.Minute, LimiterTTL: time.Minute})
+
+	rl.Reset("k")
+	assert.True(t, rl.Allow("k"))
+	assert.False(t, rl.Allow("k"), "burst should still be 1 — invalid UpdateConfig calls were ignored")
+}
+
+func TestAPIRateLimiter_ConcurrentAllowAndUpdate(t *testing.T) {
+	rl := NewAPIRateLimiter(&RateLimiterConfig{
+		RequestsPerMinute: 600,
+		Burst:             10,
+		CleanupInterval:   time.Minute,
+		LimiterTTL:        time.Minute,
+	})
+	defer rl.Stop()
+
+	var wg sync.WaitGroup
+	stop := make(chan struct{})
+
+	for i := 0; i < 8; i++ {
+		wg.Add(1)
+		go func(id int) {
+			defer wg.Done()
+			key := fmt.Sprintf("k%d", id)
+			for {
+				select {
+				case <-stop:
+					return
+				default:
+					rl.Allow(key)
+				}
+			}
+		}(i)
+	}
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for i := 0; i < 200; i++ {
+			select {
+			case <-stop:
+				return
+			default:
+				rl.UpdateConfig(&RateLimiterConfig{
+					RequestsPerMinute: float64(30 + (i % 90)),
+					Burst:             1 + (i % 20),
+					CleanupInterval:   time.Minute,
+					LimiterTTL:        time.Minute,
+				})
+				rl.SetEnabled(i%2 == 0)
+			}
+		}
+	}()
+
+	time.Sleep(100 * time.Millisecond)
+	close(stop)
+	wg.Wait()
+}
+
+func TestRateLimiterConfigFromEnv(t *testing.T) {
+	t.Setenv(RateLimitingEnabledEnv, "true")
+	t.Setenv(RateLimitingRPMEnv, "42")
+	t.Setenv(RateLimitingBurstEnv, "7")
+
+	cfg, enabled := RateLimiterConfigFromEnv()
+	assert.True(t, enabled)
+	assert.Equal(t, float64(42), cfg.RequestsPerMinute)
+	assert.Equal(t, 7, cfg.Burst)
+
+	t.Setenv(RateLimitingEnabledEnv, "false")
+	_, enabled = RateLimiterConfigFromEnv()
+	assert.False(t, enabled)
+
+	t.Setenv(RateLimitingEnabledEnv, "")
+	t.Setenv(RateLimitingRPMEnv, "")
+	t.Setenv(RateLimitingBurstEnv, "")
+	cfg, enabled = RateLimiterConfigFromEnv()
+	assert.False(t, enabled)
+	assert.Equal(t, float64(defaultAPIRPM), cfg.RequestsPerMinute)
+	assert.Equal(t, defaultAPIBurst, cfg.Burst)
+
+	t.Setenv(RateLimitingRPMEnv, "0")
+	t.Setenv(RateLimitingBurstEnv, "-5")
+	cfg, _ = RateLimiterConfigFromEnv()
+	assert.Equal(t, float64(defaultAPIRPM), cfg.RequestsPerMinute, "non-positive rpm must fall back to default")
+	assert.Equal(t, defaultAPIBurst, cfg.Burst, "non-positive burst must fall back to default")
+}
+
+func TestGetClientIP_TrustedProxies(t *testing.T) {
+	trusted, err := trustedproxy.Parse("10.0.0.0/8")
+	require.NoError(t, err)
+
+	tests := []struct {
+		name       string
+		list       *trustedproxy.List
+		remoteAddr string
+		xff        string
+		expected   string
+	}{
+		{
+			name:       "no trusted proxies ignores the header",
+			remoteAddr: "10.0.0.1:5555",
+			xff:        "1.1.1.1, 2.2.2.2",
+			expected:   "10.0.0.1",
+		},
+		{
+			name:       "behind a trusted proxy uses the right-most untrusted hop",
+			list:       trusted,
+			remoteAddr: "10.0.0.1:5555",
+			xff:        "1.1.1.1, 2.2.2.2",
+			expected:   "2.2.2.2",
+		},
+		{
+			name:       "a caller reaching us directly cannot forge the header",
+			list:       trusted,
+			remoteAddr: "203.0.113.5:5555",
+			xff:        "1.1.1.1",
+			expected:   "203.0.113.5",
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			req := httptest.NewRequest(http.MethodGet, "/test", nil)
+			req.RemoteAddr = tc.remoteAddr
+			req.Header.Set("X-Forwarded-For", tc.xff)
+
+			assert.Equal(t, tc.expected, getClientIP(req, tc.list))
+		})
+	}
+}