Resolve the rate limiter client IP behind trusted proxies

management/server/http/middleware/rate_limiter.go

@@ -14,12 +14,14 @@ import (
 	"golang.org/x/time/rate"
 
 	"github.com/netbirdio/netbird/shared/management/http/util"
+	"github.com/netbirdio/netbird/trustedproxy"
 )
 
 const (
-	RateLimitingEnabledEnv = "NB_API_RATE_LIMITING_ENABLED"
-	RateLimitingBurstEnv   = "NB_API_RATE_LIMITING_BURST"
-	RateLimitingRPMEnv     = "NB_API_RATE_LIMITING_RPM"
+	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
@@ -35,6 +37,9 @@ type RateLimiterConfig struct {
 	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
@@ -76,11 +81,18 @@ func RateLimiterConfigFromEnv() (cfg *RateLimiterConfig, enabled bool) {
 		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"
 }
 
@@ -250,7 +262,7 @@ func (rl *APIRateLimiter) Middleware(next http.Handler) http.Handler {
 			next.ServeHTTP(w, r)
 			return
 		}
-		clientIP := getClientIP(r)
+		clientIP := getClientIP(r, rl.config.TrustedProxies)
 		if !rl.Allow(clientIP) {
 			util.WriteErrorResponse("rate limit exceeded, please try again later", http.StatusTooManyRequests, w)
 			return
@@ -259,8 +271,15 @@ func (rl *APIRateLimiter) Middleware(next http.Handler) http.Handler {
 	})
 }
 
-// getClientIP extracts the client IP address from the request.
-func getClientIP(r *http.Request) string {
+// 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

management/server/http/middleware/rate_limiter_test.go

@@ -9,6 +9,9 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/trustedproxy"
 )
 
 func TestAPIRateLimiter_Allow(t *testing.T) {
@@ -134,7 +137,7 @@ func TestGetClientIP(t *testing.T) {
 		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))
+			assert.Equal(t, tc.expected, getClientIP(req, nil))
 		})
 	}
 }
@@ -327,3 +330,47 @@ func TestRateLimiterConfigFromEnv(t *testing.T) {
 	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))
+		})
+	}
+}