summaryrefslogtreecommitdiffstats
path: root/internal/ratelimit/ratelimit_test.go
diff options
context:
space:
mode:
authorbndw <ben@bdw.to>2026-02-14 08:58:57 -0800
committerbndw <ben@bdw.to>2026-02-14 08:58:57 -0800
commitf0169fa1f9d2e2a5d1c292b9080da10ef0878953 (patch)
treec85d31dfbf270fe4ebbe2c53bdbb96c0a0a45ace /internal/ratelimit/ratelimit_test.go
parent44aa0591b0eed7851e961ea17bd1c9601570ac24 (diff)
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.
Diffstat (limited to 'internal/ratelimit/ratelimit_test.go')
-rw-r--r--internal/ratelimit/ratelimit_test.go438
1 files changed, 438 insertions, 0 deletions
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 @@
1package ratelimit
2
3import (
4 "context"
5 "testing"
6 "time"
7
8 "google.golang.org/grpc"
9 "google.golang.org/grpc/codes"
10 "google.golang.org/grpc/metadata"
11 "google.golang.org/grpc/status"
12)
13
14func TestBasicRateLimit(t *testing.T) {
15 config := &Config{
16 RequestsPerSecond: 10,
17 BurstSize: 10,
18 }
19
20 limiter := New(config)
21 defer limiter.Stop()
22
23 identifier := "test-user"
24 method := "/test.Service/Method"
25
26 // First 10 requests should succeed (burst)
27 for i := 0; i < 10; i++ {
28 if !limiter.Allow(identifier, method) {
29 t.Errorf("request %d should be allowed", i)
30 }
31 }
32
33 // 11th request should be denied (burst exhausted)
34 if limiter.Allow(identifier, method) {
35 t.Error("request 11 should be denied")
36 }
37
38 // Wait for tokens to refill
39 time.Sleep(150 * time.Millisecond)
40
41 // Should allow 1 more request (1 token refilled)
42 if !limiter.Allow(identifier, method) {
43 t.Error("request after refill should be allowed")
44 }
45}
46
47func TestPerUserLimits(t *testing.T) {
48 config := &Config{
49 RequestsPerSecond: 10,
50 BurstSize: 10,
51 }
52
53 limiter := New(config)
54 defer limiter.Stop()
55
56 method := "/test.Service/Method"
57
58 // Different users should have independent limits
59 user1 := "user1"
60 user2 := "user2"
61
62 // Exhaust user1's quota
63 for i := 0; i < 10; i++ {
64 limiter.Allow(user1, method)
65 }
66
67 // User1 should be denied
68 if limiter.Allow(user1, method) {
69 t.Error("user1 should be rate limited")
70 }
71
72 // User2 should still be allowed
73 if !limiter.Allow(user2, method) {
74 t.Error("user2 should not be rate limited")
75 }
76}
77
78func TestMethodSpecificLimits(t *testing.T) {
79 config := &Config{
80 RequestsPerSecond: 10,
81 BurstSize: 10,
82 MethodLimits: map[string]MethodLimit{
83 "/test.Service/StrictMethod": {
84 RequestsPerSecond: 2,
85 BurstSize: 2,
86 },
87 },
88 }
89
90 limiter := New(config)
91 defer limiter.Stop()
92
93 identifier := "test-user"
94
95 // Regular method should allow 10 requests
96 regularMethod := "/test.Service/RegularMethod"
97 for i := 0; i < 10; i++ {
98 if !limiter.Allow(identifier, regularMethod) {
99 t.Errorf("regular method request %d should be allowed", i)
100 }
101 }
102
103 // Strict method should only allow 2 requests
104 strictMethod := "/test.Service/StrictMethod"
105 for i := 0; i < 2; i++ {
106 if !limiter.Allow(identifier, strictMethod) {
107 t.Errorf("strict method request %d should be allowed", i)
108 }
109 }
110
111 // 3rd request should be denied
112 if limiter.Allow(identifier, strictMethod) {
113 t.Error("strict method request 3 should be denied")
114 }
115}
116
117func TestUserSpecificLimits(t *testing.T) {
118 config := &Config{
119 RequestsPerSecond: 10,
120 BurstSize: 10,
121 UserLimits: map[string]UserLimit{
122 "vip-user": {
123 RequestsPerSecond: 100,
124 BurstSize: 100,
125 },
126 },
127 }
128
129 limiter := New(config)
130 defer limiter.Stop()
131
132 method := "/test.Service/Method"
133
134 // Regular user should be limited to 10
135 regularUser := "regular-user"
136 for i := 0; i < 10; i++ {
137 limiter.Allow(regularUser, method)
138 }
139 if limiter.Allow(regularUser, method) {
140 t.Error("regular user should be rate limited")
141 }
142
143 // VIP user should allow 100
144 vipUser := "vip-user"
145 for i := 0; i < 100; i++ {
146 if !limiter.Allow(vipUser, method) {
147 t.Errorf("vip user request %d should be allowed", i)
148 }
149 }
150}
151
152func TestSkipMethods(t *testing.T) {
153 config := &Config{
154 RequestsPerSecond: 1,
155 BurstSize: 1,
156 SkipMethods: []string{
157 "/health/Check",
158 },
159 }
160
161 limiter := New(config)
162 defer limiter.Stop()
163
164 identifier := "test-user"
165
166 // Regular method should be rate limited
167 regularMethod := "/test.Service/Method"
168 limiter.Allow(identifier, regularMethod)
169 if limiter.Allow(identifier, regularMethod) {
170 t.Error("regular method should be rate limited")
171 }
172
173 // Skipped method should never be rate limited
174 skipMethod := "/health/Check"
175 for i := 0; i < 100; i++ {
176 if !limiter.Allow(identifier, skipMethod) {
177 t.Error("skipped method should never be rate limited")
178 }
179 }
180}
181
182func TestSkipUsers(t *testing.T) {
183 config := &Config{
184 RequestsPerSecond: 1,
185 BurstSize: 1,
186 SkipUsers: []string{
187 "admin-user",
188 },
189 }
190
191 limiter := New(config)
192 defer limiter.Stop()
193
194 method := "/test.Service/Method"
195
196 // Regular user should be rate limited
197 regularUser := "regular-user"
198 limiter.Allow(regularUser, method)
199 if limiter.Allow(regularUser, method) {
200 t.Error("regular user should be rate limited")
201 }
202
203 // Admin user should never be rate limited
204 adminUser := "admin-user"
205 for i := 0; i < 100; i++ {
206 if !limiter.Allow(adminUser, method) {
207 t.Error("admin user should never be rate limited")
208 }
209 }
210}
211
212func TestStats(t *testing.T) {
213 config := &Config{
214 RequestsPerSecond: 10,
215 BurstSize: 5,
216 }
217
218 limiter := New(config)
219 defer limiter.Stop()
220
221 identifier := "test-user"
222 method := "/test.Service/Method"
223
224 // Make some requests
225 for i := 0; i < 5; i++ {
226 limiter.Allow(identifier, method) // All allowed (within burst)
227 }
228 for i := 0; i < 3; i++ {
229 limiter.Allow(identifier, method) // All denied (burst exhausted)
230 }
231
232 stats := limiter.Stats()
233
234 if stats.Allowed != 5 {
235 t.Errorf("expected 5 allowed, got %d", stats.Allowed)
236 }
237 if stats.Denied != 3 {
238 t.Errorf("expected 3 denied, got %d", stats.Denied)
239 }
240 if stats.ActiveLimiters != 1 {
241 t.Errorf("expected 1 active limiter, got %d", stats.ActiveLimiters)
242 }
243
244 expectedDenialRate := 37.5 // 3/8 * 100
245 if stats.DenialRate() != expectedDenialRate {
246 t.Errorf("expected denial rate %.1f%%, got %.1f%%", expectedDenialRate, stats.DenialRate())
247 }
248}
249
250func TestCleanup(t *testing.T) {
251 config := &Config{
252 RequestsPerSecond: 10,
253 BurstSize: 10,
254 CleanupInterval: 100 * time.Millisecond,
255 MaxIdleTime: 200 * time.Millisecond,
256 }
257
258 limiter := New(config)
259 defer limiter.Stop()
260
261 // Create limiters for multiple users
262 for i := 0; i < 5; i++ {
263 limiter.Allow("user-"+string(rune('0'+i)), "/test")
264 }
265
266 stats := limiter.Stats()
267 if stats.ActiveLimiters != 5 {
268 t.Errorf("expected 5 active limiters, got %d", stats.ActiveLimiters)
269 }
270
271 // Wait for cleanup to run
272 time.Sleep(350 * time.Millisecond)
273
274 stats = limiter.Stats()
275 if stats.ActiveLimiters != 0 {
276 t.Errorf("expected 0 active limiters after cleanup, got %d", stats.ActiveLimiters)
277 }
278}
279
280func TestUnaryInterceptor(t *testing.T) {
281 config := &Config{
282 RequestsPerSecond: 2,
283 BurstSize: 2,
284 }
285
286 limiter := New(config)
287 defer limiter.Stop()
288
289 interceptor := UnaryInterceptor(limiter)
290
291 // Create a test handler
292 handler := func(ctx context.Context, req interface{}) (interface{}, error) {
293 return "success", nil
294 }
295
296 info := &grpc.UnaryServerInfo{
297 FullMethod: "/test.Service/Method",
298 }
299
300 // Create context with metadata (simulating IP)
301 md := metadata.Pairs("x-real-ip", "192.168.1.1")
302 ctx := metadata.NewIncomingContext(context.Background(), md)
303
304 // First 2 requests should succeed
305 for i := 0; i < 2; i++ {
306 _, err := interceptor(ctx, nil, info, handler)
307 if err != nil {
308 t.Errorf("request %d should succeed, got error: %v", i, err)
309 }
310 }
311
312 // 3rd request should be rate limited
313 _, err := interceptor(ctx, nil, info, handler)
314 if err == nil {
315 t.Error("expected rate limit error")
316 }
317
318 st, ok := status.FromError(err)
319 if !ok {
320 t.Error("expected gRPC status error")
321 }
322 if st.Code() != codes.ResourceExhausted {
323 t.Errorf("expected ResourceExhausted, got %v", st.Code())
324 }
325}
326
327func TestGetLimitForMethod(t *testing.T) {
328 config := &Config{
329 RequestsPerSecond: 10,
330 BurstSize: 20,
331 MethodLimits: map[string]MethodLimit{
332 "/test/Method1": {
333 RequestsPerSecond: 5,
334 BurstSize: 10,
335 },
336 },
337 UserLimits: map[string]UserLimit{
338 "vip-user": {
339 RequestsPerSecond: 50,
340 BurstSize: 100,
341 MethodLimits: map[string]MethodLimit{
342 "/test/Method1": {
343 RequestsPerSecond: 25,
344 BurstSize: 50,
345 },
346 },
347 },
348 },
349 }
350
351 tests := []struct {
352 name string
353 pubkey string
354 method string
355 expectedRPS float64
356 expectedBurst int
357 }{
358 {
359 name: "default for regular user",
360 pubkey: "regular-user",
361 method: "/test/Method2",
362 expectedRPS: 10,
363 expectedBurst: 20,
364 },
365 {
366 name: "method limit for regular user",
367 pubkey: "regular-user",
368 method: "/test/Method1",
369 expectedRPS: 5,
370 expectedBurst: 10,
371 },
372 {
373 name: "user limit default method",
374 pubkey: "vip-user",
375 method: "/test/Method2",
376 expectedRPS: 50,
377 expectedBurst: 100,
378 },
379 {
380 name: "user method limit (highest precedence)",
381 pubkey: "vip-user",
382 method: "/test/Method1",
383 expectedRPS: 25,
384 expectedBurst: 50,
385 },
386 }
387
388 for _, tt := range tests {
389 t.Run(tt.name, func(t *testing.T) {
390 rps, burst := config.GetLimitForMethod(tt.pubkey, tt.method)
391 if rps != tt.expectedRPS {
392 t.Errorf("expected RPS %.1f, got %.1f", tt.expectedRPS, rps)
393 }
394 if burst != tt.expectedBurst {
395 t.Errorf("expected burst %d, got %d", tt.expectedBurst, burst)
396 }
397 })
398 }
399}
400
401func BenchmarkRateLimitAllow(b *testing.B) {
402 config := &Config{
403 RequestsPerSecond: 1000,
404 BurstSize: 1000,
405 }
406
407 limiter := New(config)
408 defer limiter.Stop()
409
410 identifier := "bench-user"
411 method := "/test.Service/Method"
412
413 b.ResetTimer()
414 for i := 0; i < b.N; i++ {
415 limiter.Allow(identifier, method)
416 }
417}
418
419func BenchmarkRateLimitDeny(b *testing.B) {
420 config := &Config{
421 RequestsPerSecond: 1,
422 BurstSize: 1,
423 }
424
425 limiter := New(config)
426 defer limiter.Stop()
427
428 identifier := "bench-user"
429 method := "/test.Service/Method"
430
431 // Exhaust quota
432 limiter.Allow(identifier, method)
433
434 b.ResetTimer()
435 for i := 0; i < b.N; i++ {
436 limiter.Allow(identifier, method)
437 }
438}