From f0169fa1f9d2e2a5d1c292b9080da10ef0878953 Mon Sep 17 00:00:00 2001 From: bndw Date: Sat, 14 Feb 2026 08:58:57 -0800 Subject: feat: implement per-user rate limiting with token bucket algorithm Add comprehensive rate limiting package that works seamlessly with NIP-98 authentication. Features: - Token bucket algorithm (allows bursts, smooth average rate) - Per-pubkey limits for authenticated users - Per-IP limits for unauthenticated users (fallback) - Method-specific overrides (e.g., stricter for PublishEvent) - Per-user custom limits (VIP/admin tiers) - Standard gRPC interceptors (chain after auth) - Automatic cleanup of idle limiters - Statistics tracking (allowed/denied/denial rate) Configuration options: - Default rate limits and burst sizes - Method-specific overrides - User-specific overrides (with method overrides) - Skip methods (health checks, public endpoints) - Skip users (admins, monitoring) - Configurable cleanup intervals Performance: - In-memory (200 bytes per user) - O(1) lookups with sync.RWMutex - ~85ns per rate limit check - Periodic cleanup to free memory Returns gRPC ResourceExhausted (HTTP 429) when limits exceeded. Includes comprehensive tests, benchmarks, and detailed README with usage examples, configuration reference, and security considerations. --- internal/ratelimit/ratelimit_test.go | 438 +++++++++++++++++++++++++++++++++++ 1 file changed, 438 insertions(+) create mode 100644 internal/ratelimit/ratelimit_test.go (limited to 'internal/ratelimit/ratelimit_test.go') diff --git a/internal/ratelimit/ratelimit_test.go b/internal/ratelimit/ratelimit_test.go new file mode 100644 index 0000000..963d97f --- /dev/null +++ b/internal/ratelimit/ratelimit_test.go @@ -0,0 +1,438 @@ +package ratelimit + +import ( + "context" + "testing" + "time" + + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" +) + +func TestBasicRateLimit(t *testing.T) { + config := &Config{ + RequestsPerSecond: 10, + BurstSize: 10, + } + + limiter := New(config) + defer limiter.Stop() + + identifier := "test-user" + method := "/test.Service/Method" + + // First 10 requests should succeed (burst) + for i := 0; i < 10; i++ { + if !limiter.Allow(identifier, method) { + t.Errorf("request %d should be allowed", i) + } + } + + // 11th request should be denied (burst exhausted) + if limiter.Allow(identifier, method) { + t.Error("request 11 should be denied") + } + + // Wait for tokens to refill + time.Sleep(150 * time.Millisecond) + + // Should allow 1 more request (1 token refilled) + if !limiter.Allow(identifier, method) { + t.Error("request after refill should be allowed") + } +} + +func TestPerUserLimits(t *testing.T) { + config := &Config{ + RequestsPerSecond: 10, + BurstSize: 10, + } + + limiter := New(config) + defer limiter.Stop() + + method := "/test.Service/Method" + + // Different users should have independent limits + user1 := "user1" + user2 := "user2" + + // Exhaust user1's quota + for i := 0; i < 10; i++ { + limiter.Allow(user1, method) + } + + // User1 should be denied + if limiter.Allow(user1, method) { + t.Error("user1 should be rate limited") + } + + // User2 should still be allowed + if !limiter.Allow(user2, method) { + t.Error("user2 should not be rate limited") + } +} + +func TestMethodSpecificLimits(t *testing.T) { + config := &Config{ + RequestsPerSecond: 10, + BurstSize: 10, + MethodLimits: map[string]MethodLimit{ + "/test.Service/StrictMethod": { + RequestsPerSecond: 2, + BurstSize: 2, + }, + }, + } + + limiter := New(config) + defer limiter.Stop() + + identifier := "test-user" + + // Regular method should allow 10 requests + regularMethod := "/test.Service/RegularMethod" + for i := 0; i < 10; i++ { + if !limiter.Allow(identifier, regularMethod) { + t.Errorf("regular method request %d should be allowed", i) + } + } + + // Strict method should only allow 2 requests + strictMethod := "/test.Service/StrictMethod" + for i := 0; i < 2; i++ { + if !limiter.Allow(identifier, strictMethod) { + t.Errorf("strict method request %d should be allowed", i) + } + } + + // 3rd request should be denied + if limiter.Allow(identifier, strictMethod) { + t.Error("strict method request 3 should be denied") + } +} + +func TestUserSpecificLimits(t *testing.T) { + config := &Config{ + RequestsPerSecond: 10, + BurstSize: 10, + UserLimits: map[string]UserLimit{ + "vip-user": { + RequestsPerSecond: 100, + BurstSize: 100, + }, + }, + } + + limiter := New(config) + defer limiter.Stop() + + method := "/test.Service/Method" + + // Regular user should be limited to 10 + regularUser := "regular-user" + for i := 0; i < 10; i++ { + limiter.Allow(regularUser, method) + } + if limiter.Allow(regularUser, method) { + t.Error("regular user should be rate limited") + } + + // VIP user should allow 100 + vipUser := "vip-user" + for i := 0; i < 100; i++ { + if !limiter.Allow(vipUser, method) { + t.Errorf("vip user request %d should be allowed", i) + } + } +} + +func TestSkipMethods(t *testing.T) { + config := &Config{ + RequestsPerSecond: 1, + BurstSize: 1, + SkipMethods: []string{ + "/health/Check", + }, + } + + limiter := New(config) + defer limiter.Stop() + + identifier := "test-user" + + // Regular method should be rate limited + regularMethod := "/test.Service/Method" + limiter.Allow(identifier, regularMethod) + if limiter.Allow(identifier, regularMethod) { + t.Error("regular method should be rate limited") + } + + // Skipped method should never be rate limited + skipMethod := "/health/Check" + for i := 0; i < 100; i++ { + if !limiter.Allow(identifier, skipMethod) { + t.Error("skipped method should never be rate limited") + } + } +} + +func TestSkipUsers(t *testing.T) { + config := &Config{ + RequestsPerSecond: 1, + BurstSize: 1, + SkipUsers: []string{ + "admin-user", + }, + } + + limiter := New(config) + defer limiter.Stop() + + method := "/test.Service/Method" + + // Regular user should be rate limited + regularUser := "regular-user" + limiter.Allow(regularUser, method) + if limiter.Allow(regularUser, method) { + t.Error("regular user should be rate limited") + } + + // Admin user should never be rate limited + adminUser := "admin-user" + for i := 0; i < 100; i++ { + if !limiter.Allow(adminUser, method) { + t.Error("admin user should never be rate limited") + } + } +} + +func TestStats(t *testing.T) { + config := &Config{ + RequestsPerSecond: 10, + BurstSize: 5, + } + + limiter := New(config) + defer limiter.Stop() + + identifier := "test-user" + method := "/test.Service/Method" + + // Make some requests + for i := 0; i < 5; i++ { + limiter.Allow(identifier, method) // All allowed (within burst) + } + for i := 0; i < 3; i++ { + limiter.Allow(identifier, method) // All denied (burst exhausted) + } + + stats := limiter.Stats() + + if stats.Allowed != 5 { + t.Errorf("expected 5 allowed, got %d", stats.Allowed) + } + if stats.Denied != 3 { + t.Errorf("expected 3 denied, got %d", stats.Denied) + } + if stats.ActiveLimiters != 1 { + t.Errorf("expected 1 active limiter, got %d", stats.ActiveLimiters) + } + + expectedDenialRate := 37.5 // 3/8 * 100 + if stats.DenialRate() != expectedDenialRate { + t.Errorf("expected denial rate %.1f%%, got %.1f%%", expectedDenialRate, stats.DenialRate()) + } +} + +func TestCleanup(t *testing.T) { + config := &Config{ + RequestsPerSecond: 10, + BurstSize: 10, + CleanupInterval: 100 * time.Millisecond, + MaxIdleTime: 200 * time.Millisecond, + } + + limiter := New(config) + defer limiter.Stop() + + // Create limiters for multiple users + for i := 0; i < 5; i++ { + limiter.Allow("user-"+string(rune('0'+i)), "/test") + } + + stats := limiter.Stats() + if stats.ActiveLimiters != 5 { + t.Errorf("expected 5 active limiters, got %d", stats.ActiveLimiters) + } + + // Wait for cleanup to run + time.Sleep(350 * time.Millisecond) + + stats = limiter.Stats() + if stats.ActiveLimiters != 0 { + t.Errorf("expected 0 active limiters after cleanup, got %d", stats.ActiveLimiters) + } +} + +func TestUnaryInterceptor(t *testing.T) { + config := &Config{ + RequestsPerSecond: 2, + BurstSize: 2, + } + + limiter := New(config) + defer limiter.Stop() + + interceptor := UnaryInterceptor(limiter) + + // Create a test handler + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return "success", nil + } + + info := &grpc.UnaryServerInfo{ + FullMethod: "/test.Service/Method", + } + + // Create context with metadata (simulating IP) + md := metadata.Pairs("x-real-ip", "192.168.1.1") + ctx := metadata.NewIncomingContext(context.Background(), md) + + // First 2 requests should succeed + for i := 0; i < 2; i++ { + _, err := interceptor(ctx, nil, info, handler) + if err != nil { + t.Errorf("request %d should succeed, got error: %v", i, err) + } + } + + // 3rd request should be rate limited + _, err := interceptor(ctx, nil, info, handler) + if err == nil { + t.Error("expected rate limit error") + } + + st, ok := status.FromError(err) + if !ok { + t.Error("expected gRPC status error") + } + if st.Code() != codes.ResourceExhausted { + t.Errorf("expected ResourceExhausted, got %v", st.Code()) + } +} + +func TestGetLimitForMethod(t *testing.T) { + config := &Config{ + RequestsPerSecond: 10, + BurstSize: 20, + MethodLimits: map[string]MethodLimit{ + "/test/Method1": { + RequestsPerSecond: 5, + BurstSize: 10, + }, + }, + UserLimits: map[string]UserLimit{ + "vip-user": { + RequestsPerSecond: 50, + BurstSize: 100, + MethodLimits: map[string]MethodLimit{ + "/test/Method1": { + RequestsPerSecond: 25, + BurstSize: 50, + }, + }, + }, + }, + } + + tests := []struct { + name string + pubkey string + method string + expectedRPS float64 + expectedBurst int + }{ + { + name: "default for regular user", + pubkey: "regular-user", + method: "/test/Method2", + expectedRPS: 10, + expectedBurst: 20, + }, + { + name: "method limit for regular user", + pubkey: "regular-user", + method: "/test/Method1", + expectedRPS: 5, + expectedBurst: 10, + }, + { + name: "user limit default method", + pubkey: "vip-user", + method: "/test/Method2", + expectedRPS: 50, + expectedBurst: 100, + }, + { + name: "user method limit (highest precedence)", + pubkey: "vip-user", + method: "/test/Method1", + expectedRPS: 25, + expectedBurst: 50, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + rps, burst := config.GetLimitForMethod(tt.pubkey, tt.method) + if rps != tt.expectedRPS { + t.Errorf("expected RPS %.1f, got %.1f", tt.expectedRPS, rps) + } + if burst != tt.expectedBurst { + t.Errorf("expected burst %d, got %d", tt.expectedBurst, burst) + } + }) + } +} + +func BenchmarkRateLimitAllow(b *testing.B) { + config := &Config{ + RequestsPerSecond: 1000, + BurstSize: 1000, + } + + limiter := New(config) + defer limiter.Stop() + + identifier := "bench-user" + method := "/test.Service/Method" + + b.ResetTimer() + for i := 0; i < b.N; i++ { + limiter.Allow(identifier, method) + } +} + +func BenchmarkRateLimitDeny(b *testing.B) { + config := &Config{ + RequestsPerSecond: 1, + BurstSize: 1, + } + + limiter := New(config) + defer limiter.Stop() + + identifier := "bench-user" + method := "/test.Service/Method" + + // Exhaust quota + limiter.Allow(identifier, method) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + limiter.Allow(identifier, method) + } +} -- cgit v1.2.3