Add proxy <-> management authentication

management/internals/shared/grpc/proxy_auth_ratelimit_test.go

@@ -0,0 +1,98 @@
+package grpc
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/time/rate"
+)
+
+func TestAuthFailureLimiter_NotLimitedInitially(t *testing.T) {
+	l := newAuthFailureLimiter()
+	defer l.stop()
+
+	assert.False(t, l.isLimited("192.168.1.1"), "new IP should not be rate limited")
+}
+
+func TestAuthFailureLimiter_LimitedAfterBurst(t *testing.T) {
+	l := newAuthFailureLimiter()
+	defer l.stop()
+
+	ip := "192.168.1.1"
+	for i := 0; i < proxyAuthFailureBurst; i++ {
+		l.recordFailure(ip)
+	}
+
+	assert.True(t, l.isLimited(ip), "IP should be limited after exhausting burst")
+}
+
+func TestAuthFailureLimiter_DifferentIPsIndependent(t *testing.T) {
+	l := newAuthFailureLimiter()
+	defer l.stop()
+
+	for i := 0; i < proxyAuthFailureBurst; i++ {
+		l.recordFailure("192.168.1.1")
+	}
+
+	assert.True(t, l.isLimited("192.168.1.1"))
+	assert.False(t, l.isLimited("192.168.1.2"), "different IP should not be affected")
+}
+
+func TestAuthFailureLimiter_RecoveryOverTime(t *testing.T) {
+	l := newAuthFailureLimiterWithRate(rate.Limit(100)) // 100 tokens/sec for fast recovery
+	defer l.stop()
+
+	ip := "10.0.0.1"
+
+	// Exhaust burst
+	for i := 0; i < proxyAuthFailureBurst; i++ {
+		l.recordFailure(ip)
+	}
+	require.True(t, l.isLimited(ip))
+
+	// Wait for token replenishment
+	time.Sleep(50 * time.Millisecond)
+
+	assert.False(t, l.isLimited(ip), "should recover after tokens replenish")
+}
+
+func TestAuthFailureLimiter_Cleanup(t *testing.T) {
+	l := newAuthFailureLimiter()
+	defer l.stop()
+
+	l.recordFailure("10.0.0.1")
+
+	l.mu.Lock()
+	require.Len(t, l.limiters, 1)
+	// Backdate the entry so it looks stale
+	l.limiters["10.0.0.1"].lastAccess = time.Now().Add(-proxyAuthLimiterTTL - time.Minute)
+	l.mu.Unlock()
+
+	l.cleanup()
+
+	l.mu.Lock()
+	assert.Empty(t, l.limiters, "stale entries should be cleaned up")
+	l.mu.Unlock()
+}
+
+func TestAuthFailureLimiter_CleanupKeepsFresh(t *testing.T) {
+	l := newAuthFailureLimiter()
+	defer l.stop()
+
+	l.recordFailure("10.0.0.1")
+	l.recordFailure("10.0.0.2")
+
+	l.mu.Lock()
+	// Only backdate one entry
+	l.limiters["10.0.0.1"].lastAccess = time.Now().Add(-proxyAuthLimiterTTL - time.Minute)
+	l.mu.Unlock()
+
+	l.cleanup()
+
+	l.mu.Lock()
+	assert.Len(t, l.limiters, 1, "only stale entries should be removed")
+	assert.Contains(t, l.limiters, "10.0.0.2")
+	l.mu.Unlock()
+}