arete_server/websocket/
rate_limiter.rs

1use std::collections::HashMap;
2use std::net::SocketAddr;
3use std::sync::Arc;
4use std::time::{Duration, Instant};
5use tokio::sync::RwLock;
6use tracing::{debug, warn};
7
8/// Rate limit window configuration
9#[derive(Debug, Clone, Copy)]
10pub struct RateLimitWindow {
11    /// Maximum number of requests allowed in the window
12    pub max_requests: u32,
13    /// Window duration
14    pub window_duration: Duration,
15    /// Burst allowance (extra requests allowed temporarily)
16    pub burst: u32,
17}
18
19impl RateLimitWindow {
20    /// Create a new rate limit window
21    pub fn new(max_requests: u32, window_duration: Duration) -> Self {
22        Self {
23            max_requests,
24            window_duration,
25            burst: 0,
26        }
27    }
28
29    /// Add burst allowance
30    pub fn with_burst(mut self, burst: u32) -> Self {
31        self.burst = burst;
32        self
33    }
34}
35
36impl Default for RateLimitWindow {
37    fn default() -> Self {
38        Self {
39            max_requests: 100,
40            window_duration: Duration::from_secs(60),
41            burst: 10,
42        }
43    }
44}
45
46/// Rate limit result
47#[derive(Debug, Clone)]
48pub enum RateLimitResult {
49    /// Request is allowed
50    Allowed { remaining: u32, reset_at: Instant },
51    /// Request is denied due to rate limiting
52    Denied { retry_after: Duration, limit: u32 },
53}
54
55/// A single rate limit bucket using sliding window algorithm
56#[derive(Debug)]
57struct RateLimitBucket {
58    /// Request timestamps in the current window
59    requests: Vec<Instant>,
60    /// Window configuration
61    window: RateLimitWindow,
62}
63
64impl RateLimitBucket {
65    fn new(window: RateLimitWindow) -> Self {
66        Self {
67            requests: Vec::with_capacity((window.max_requests + window.burst) as usize),
68            window,
69        }
70    }
71
72    fn prune_expired(&mut self, now: Instant) {
73        let cutoff = now - self.window.window_duration;
74        self.requests.retain(|&t| t > cutoff);
75    }
76
77    /// Check if a request is allowed and record it.
78    ///
79    /// `limit_override` replaces the configured max+burst with a signed
80    /// per-token limit (no additional burst) when present.
81    fn check_and_record(&mut self, now: Instant, limit_override: Option<u32>) -> RateLimitResult {
82        self.prune_expired(now);
83
84        let limit = limit_override.unwrap_or(self.window.max_requests + self.window.burst);
85        let current_count = self.requests.len() as u32;
86
87        if current_count >= limit {
88            let reported_limit = limit_override.unwrap_or(self.window.max_requests);
89            // Calculate retry after time
90            if let Some(oldest) = self.requests.first() {
91                let retry_after =
92                    (*oldest + self.window.window_duration).saturating_duration_since(now);
93                RateLimitResult::Denied {
94                    retry_after,
95                    limit: reported_limit,
96                }
97            } else {
98                RateLimitResult::Denied {
99                    retry_after: self.window.window_duration,
100                    limit: reported_limit,
101                }
102            }
103        } else {
104            self.requests.push(now);
105            let reset_at = now + self.window.window_duration;
106            RateLimitResult::Allowed {
107                remaining: limit - current_count - 1,
108                reset_at,
109            }
110        }
111    }
112}
113
114/// Rate limiter configuration per key type
115#[derive(Debug, Clone)]
116pub struct RateLimiterConfig {
117    /// Rate limit for handshake attempts per IP
118    pub handshake_per_ip: RateLimitWindow,
119    /// Rate limit for connection attempts per resolved consumer
120    /// (falls back to the token subject for legacy tokens)
121    pub connections_per_consumer: RateLimitWindow,
122    /// Rate limit for connection attempts per resolved account
123    /// (falls back to the metering key for legacy tokens)
124    pub connections_per_account: RateLimitWindow,
125    /// Rate limit for subscription requests per connection
126    pub subscriptions_per_connection: RateLimitWindow,
127    /// Rate limit for messages per connection
128    pub messages_per_connection: RateLimitWindow,
129    /// Rate limit for snapshot requests per connection
130    pub snapshots_per_connection: RateLimitWindow,
131    /// Enable rate limiting (can be disabled for testing)
132    pub enabled: bool,
133}
134
135impl Default for RateLimiterConfig {
136    fn default() -> Self {
137        Self {
138            handshake_per_ip: RateLimitWindow::new(60, Duration::from_secs(60)).with_burst(10),
139            connections_per_consumer: RateLimitWindow::new(30, Duration::from_secs(60))
140                .with_burst(5),
141            connections_per_account: RateLimitWindow::new(100, Duration::from_secs(60))
142                .with_burst(20),
143            subscriptions_per_connection: RateLimitWindow::new(120, Duration::from_secs(60))
144                .with_burst(10),
145            messages_per_connection: RateLimitWindow::new(1000, Duration::from_secs(60))
146                .with_burst(100),
147            snapshots_per_connection: RateLimitWindow::new(30, Duration::from_secs(60))
148                .with_burst(5),
149            enabled: true,
150        }
151    }
152}
153
154impl RateLimiterConfig {
155    /// Load a rate limit window from `{prefix}_MAX` and `{prefix}_WINDOW_SECS`.
156    fn window_from_env(prefix: &str) -> Option<RateLimitWindow> {
157        let max = std::env::var(format!("{prefix}_MAX")).ok()?.parse().ok()?;
158        let secs = std::env::var(format!("{prefix}_WINDOW_SECS"))
159            .ok()?
160            .parse()
161            .ok()?;
162        Some(RateLimitWindow::new(max, Duration::from_secs(secs)))
163    }
164
165    /// Load a window from its current env prefix, falling back to a
166    /// deprecated alias with a startup warning.
167    fn window_from_env_with_alias(prefix: &str, deprecated: &str) -> Option<RateLimitWindow> {
168        if let Some(window) = Self::window_from_env(prefix) {
169            return Some(window);
170        }
171        let window = Self::window_from_env(deprecated)?;
172        tracing::warn!(
173            deprecated_prefix = deprecated,
174            replacement_prefix = prefix,
175            "deprecated rate-limit environment variables are set; \
176             rename them before the alias is removed"
177        );
178        Some(window)
179    }
180
181    /// Load configuration from environment variables
182    pub fn from_env() -> Self {
183        let mut config = Self::default();
184
185        // Handshake rate limit
186        if let Some(window) = Self::window_from_env("ARETE_RATE_LIMIT_HANDSHAKE_PER_IP") {
187            config.handshake_per_ip = window;
188        }
189
190        // Connection attempts per resolved consumer
191        // (deprecated alias: per-subject variables)
192        if let Some(window) = Self::window_from_env_with_alias(
193            "ARETE_RATE_LIMIT_CONNECTIONS_PER_CONSUMER",
194            "ARETE_RATE_LIMIT_CONNECTIONS_PER_SUBJECT",
195        ) {
196            config.connections_per_consumer = window;
197        }
198
199        // Connection attempts per resolved account
200        // (deprecated alias: per-metering-key variables)
201        if let Some(window) = Self::window_from_env_with_alias(
202            "ARETE_RATE_LIMIT_CONNECTIONS_PER_ACCOUNT",
203            "ARETE_RATE_LIMIT_CONNECTIONS_PER_METERING_KEY",
204        ) {
205            config.connections_per_account = window;
206        }
207
208        // Subscriptions per connection
209        if let Some(window) = Self::window_from_env("ARETE_RATE_LIMIT_SUBSCRIPTIONS_PER_CONNECTION")
210        {
211            config.subscriptions_per_connection = window;
212        }
213
214        // Messages per connection
215        if let Some(window) = Self::window_from_env("ARETE_RATE_LIMIT_MESSAGES_PER_CONNECTION") {
216            config.messages_per_connection = window;
217        }
218
219        // Snapshots per connection
220        if let Some(window) = Self::window_from_env("ARETE_RATE_LIMIT_SNAPSHOTS_PER_CONNECTION") {
221            config.snapshots_per_connection = window;
222        }
223
224        // Enable/disable
225        if let Ok(enabled) = std::env::var("ARETE_RATE_LIMITING_ENABLED") {
226            config.enabled = enabled.parse().unwrap_or(true);
227        }
228
229        config
230    }
231
232    /// Disable rate limiting (useful for testing)
233    pub fn disabled() -> Self {
234        Self {
235            enabled: false,
236            ..Default::default()
237        }
238    }
239}
240
241/// Multi-tenant rate limiter with per-key tracking
242#[derive(Debug)]
243pub struct WebSocketRateLimiter {
244    config: RateLimiterConfig,
245    /// Per-IP handshake rate limits
246    ip_buckets: Arc<RwLock<HashMap<String, RateLimitBucket>>>,
247    /// Per-consumer connection rate limits (subject for legacy tokens)
248    consumer_buckets: Arc<RwLock<HashMap<String, RateLimitBucket>>>,
249    /// Per-account connection rate limits (metering key for legacy tokens)
250    account_buckets: Arc<RwLock<HashMap<String, RateLimitBucket>>>,
251    /// Per-consumer subscription-create rate limits (signed limit only)
252    consumer_subscription_buckets: Arc<RwLock<HashMap<String, RateLimitBucket>>>,
253    /// Per-account subscription-create rate limits (signed limit only)
254    account_subscription_buckets: Arc<RwLock<HashMap<String, RateLimitBucket>>>,
255    /// Per-connection subscription rate limits
256    subscription_buckets: Arc<RwLock<HashMap<uuid::Uuid, RateLimitBucket>>>,
257    /// Per-connection message rate limits
258    message_buckets: Arc<RwLock<HashMap<uuid::Uuid, RateLimitBucket>>>,
259    /// Per-connection snapshot rate limits
260    snapshot_buckets: Arc<RwLock<HashMap<uuid::Uuid, RateLimitBucket>>>,
261}
262
263impl WebSocketRateLimiter {
264    /// Create a new rate limiter with the given configuration
265    pub fn new(config: RateLimiterConfig) -> Self {
266        Self {
267            config,
268            ip_buckets: Arc::new(RwLock::new(HashMap::new())),
269            consumer_buckets: Arc::new(RwLock::new(HashMap::new())),
270            account_buckets: Arc::new(RwLock::new(HashMap::new())),
271            consumer_subscription_buckets: Arc::new(RwLock::new(HashMap::new())),
272            account_subscription_buckets: Arc::new(RwLock::new(HashMap::new())),
273            subscription_buckets: Arc::new(RwLock::new(HashMap::new())),
274            message_buckets: Arc::new(RwLock::new(HashMap::new())),
275            snapshot_buckets: Arc::new(RwLock::new(HashMap::new())),
276        }
277    }
278
279    /// Check if handshake is allowed from the given IP
280    pub async fn check_handshake(&self, addr: SocketAddr) -> RateLimitResult {
281        if !self.config.enabled {
282            return RateLimitResult::Allowed {
283                remaining: u32::MAX,
284                reset_at: Instant::now() + Duration::from_secs(60),
285            };
286        }
287
288        let ip = addr.ip().to_string();
289        let mut buckets = self.ip_buckets.write().await;
290        let bucket = buckets
291            .entry(ip.clone())
292            .or_insert_with(|| RateLimitBucket::new(self.config.handshake_per_ip));
293
294        let result = bucket.check_and_record(Instant::now(), None);
295
296        match &result {
297            RateLimitResult::Denied { retry_after, limit } => {
298                warn!(
299                    ip = %ip,
300                    retry_after_secs = retry_after.as_secs(),
301                    limit = limit,
302                    "Rate limit exceeded for handshake"
303                );
304            }
305            RateLimitResult::Allowed { remaining, .. } => {
306                debug!(
307                    ip = %ip,
308                    remaining = remaining,
309                    "Handshake rate limit check passed"
310                );
311            }
312        }
313
314        result
315    }
316
317    fn allowed_unlimited() -> RateLimitResult {
318        RateLimitResult::Allowed {
319            remaining: u32::MAX,
320            reset_at: Instant::now() + Duration::from_secs(60),
321        }
322    }
323
324    async fn check_keyed_bucket(
325        &self,
326        buckets: &RwLock<HashMap<String, RateLimitBucket>>,
327        window: RateLimitWindow,
328        key: &str,
329        limit_override: Option<u32>,
330    ) -> RateLimitResult {
331        let mut buckets = buckets.write().await;
332        let bucket = buckets
333            .entry(key.to_string())
334            .or_insert_with(|| RateLimitBucket::new(window));
335        bucket.check_and_record(Instant::now(), limit_override)
336    }
337
338    /// Check if a connection attempt is allowed for the resolved consumer.
339    ///
340    /// `limit_override` carries the signed
341    /// `limits.max_connection_attempts_per_minute` when present; the
342    /// configured window applies otherwise.
343    pub async fn check_connection_for_consumer(
344        &self,
345        consumer: &str,
346        limit_override: Option<u32>,
347    ) -> RateLimitResult {
348        if !self.config.enabled {
349            return Self::allowed_unlimited();
350        }
351        self.check_keyed_bucket(
352            &self.consumer_buckets,
353            self.config.connections_per_consumer,
354            consumer,
355            limit_override,
356        )
357        .await
358    }
359
360    /// Check if a connection attempt is allowed for the resolved account.
361    ///
362    /// `limit_override` carries the signed
363    /// `account_limits.max_connection_attempts_per_minute` when present; the
364    /// configured window applies otherwise.
365    pub async fn check_connection_for_account(
366        &self,
367        account: &str,
368        limit_override: Option<u32>,
369    ) -> RateLimitResult {
370        if !self.config.enabled {
371            return Self::allowed_unlimited();
372        }
373        self.check_keyed_bucket(
374            &self.account_buckets,
375            self.config.connections_per_account,
376            account,
377            limit_override,
378        )
379        .await
380    }
381
382    /// Check the signed per-consumer subscription-create rate.
383    ///
384    /// Enforced only when the token carries
385    /// `limits.max_subscription_creates_per_minute`; a `None` limit is
386    /// allowed without creating bucket state.
387    pub async fn check_subscription_create_for_consumer(
388        &self,
389        consumer: &str,
390        limit: Option<u32>,
391    ) -> RateLimitResult {
392        let Some(limit) = limit else {
393            return Self::allowed_unlimited();
394        };
395        if !self.config.enabled {
396            return Self::allowed_unlimited();
397        }
398        self.check_keyed_bucket(
399            &self.consumer_subscription_buckets,
400            RateLimitWindow::new(limit, Duration::from_secs(60)),
401            consumer,
402            Some(limit),
403        )
404        .await
405    }
406
407    /// Check the signed per-account subscription-create rate.
408    ///
409    /// Enforced only when the token carries
410    /// `account_limits.max_subscription_creates_per_minute`; a `None` limit
411    /// is allowed without creating bucket state.
412    pub async fn check_subscription_create_for_account(
413        &self,
414        account: &str,
415        limit: Option<u32>,
416    ) -> RateLimitResult {
417        let Some(limit) = limit else {
418            return Self::allowed_unlimited();
419        };
420        if !self.config.enabled {
421            return Self::allowed_unlimited();
422        }
423        self.check_keyed_bucket(
424            &self.account_subscription_buckets,
425            RateLimitWindow::new(limit, Duration::from_secs(60)),
426            account,
427            Some(limit),
428        )
429        .await
430    }
431
432    /// Check if connection is allowed for the given subject
433    #[deprecated(note = "use check_connection_for_consumer with the resolved consumer identity")]
434    pub async fn check_connection_for_subject(&self, subject: &str) -> RateLimitResult {
435        self.check_connection_for_consumer(subject, None).await
436    }
437
438    /// Check if connection is allowed for the given metering key
439    #[deprecated(note = "use check_connection_for_account with the resolved account identity")]
440    pub async fn check_connection_for_metering_key(&self, metering_key: &str) -> RateLimitResult {
441        self.check_connection_for_account(metering_key, None).await
442    }
443
444    /// Check if subscription is allowed for the given connection
445    pub async fn check_subscription(&self, client_id: uuid::Uuid) -> RateLimitResult {
446        if !self.config.enabled {
447            return RateLimitResult::Allowed {
448                remaining: u32::MAX,
449                reset_at: Instant::now() + Duration::from_secs(60),
450            };
451        }
452
453        let mut buckets = self.subscription_buckets.write().await;
454        let bucket = buckets
455            .entry(client_id)
456            .or_insert_with(|| RateLimitBucket::new(self.config.subscriptions_per_connection));
457
458        bucket.check_and_record(Instant::now(), None)
459    }
460
461    /// Check if message is allowed for the given connection
462    pub async fn check_message(&self, client_id: uuid::Uuid) -> RateLimitResult {
463        if !self.config.enabled {
464            return RateLimitResult::Allowed {
465                remaining: u32::MAX,
466                reset_at: Instant::now() + Duration::from_secs(60),
467            };
468        }
469
470        let mut buckets = self.message_buckets.write().await;
471        let bucket = buckets
472            .entry(client_id)
473            .or_insert_with(|| RateLimitBucket::new(self.config.messages_per_connection));
474
475        bucket.check_and_record(Instant::now(), None)
476    }
477
478    /// Check if snapshot is allowed for the given connection
479    pub async fn check_snapshot(&self, client_id: uuid::Uuid) -> RateLimitResult {
480        if !self.config.enabled {
481            return RateLimitResult::Allowed {
482                remaining: u32::MAX,
483                reset_at: Instant::now() + Duration::from_secs(60),
484            };
485        }
486
487        let mut buckets = self.snapshot_buckets.write().await;
488        let bucket = buckets
489            .entry(client_id)
490            .or_insert_with(|| RateLimitBucket::new(self.config.snapshots_per_connection));
491
492        bucket.check_and_record(Instant::now(), None)
493    }
494
495    /// Clean up stale buckets to prevent memory growth
496    pub async fn cleanup_stale_buckets(&self) {
497        let now = Instant::now();
498
499        // Clean up IP buckets
500        {
501            let mut buckets = self.ip_buckets.write().await;
502            buckets.retain(|_, bucket| {
503                bucket.prune_expired(now);
504                !bucket.requests.is_empty()
505            });
506        }
507
508        // Clean up consumer/account connection and subscription-create buckets
509        for buckets in [
510            &self.consumer_buckets,
511            &self.account_buckets,
512            &self.consumer_subscription_buckets,
513            &self.account_subscription_buckets,
514        ] {
515            let mut buckets = buckets.write().await;
516            buckets.retain(|_, bucket| {
517                bucket.prune_expired(now);
518                !bucket.requests.is_empty()
519            });
520        }
521
522        // Clean up connection-specific buckets for disconnected clients
523        // These should be explicitly removed when clients disconnect
524    }
525
526    /// Remove all rate limit buckets for a disconnected client
527    pub async fn remove_client_buckets(&self, client_id: uuid::Uuid) {
528        let mut sub_buckets = self.subscription_buckets.write().await;
529        sub_buckets.remove(&client_id);
530        drop(sub_buckets);
531
532        let mut msg_buckets = self.message_buckets.write().await;
533        msg_buckets.remove(&client_id);
534        drop(msg_buckets);
535
536        let mut snap_buckets = self.snapshot_buckets.write().await;
537        snap_buckets.remove(&client_id);
538    }
539
540    /// Start a background task to periodically clean up stale buckets
541    pub fn start_cleanup_task(&self) {
542        let limiter = self.clone();
543        tokio::spawn(async move {
544            let mut interval = tokio::time::interval(Duration::from_secs(60));
545            loop {
546                interval.tick().await;
547                limiter.cleanup_stale_buckets().await;
548            }
549        });
550    }
551}
552
553impl Clone for WebSocketRateLimiter {
554    fn clone(&self) -> Self {
555        Self {
556            config: self.config.clone(),
557            ip_buckets: Arc::clone(&self.ip_buckets),
558            consumer_buckets: Arc::clone(&self.consumer_buckets),
559            account_buckets: Arc::clone(&self.account_buckets),
560            consumer_subscription_buckets: Arc::clone(&self.consumer_subscription_buckets),
561            account_subscription_buckets: Arc::clone(&self.account_subscription_buckets),
562            subscription_buckets: Arc::clone(&self.subscription_buckets),
563            message_buckets: Arc::clone(&self.message_buckets),
564            snapshot_buckets: Arc::clone(&self.snapshot_buckets),
565        }
566    }
567}
568
569#[cfg(test)]
570mod tests {
571    use super::*;
572
573    fn test_config() -> RateLimiterConfig {
574        RateLimiterConfig {
575            enabled: true,
576            handshake_per_ip: RateLimitWindow::new(60, Duration::from_secs(60)).with_burst(10),
577            connections_per_consumer: RateLimitWindow::new(30, Duration::from_secs(60))
578                .with_burst(5),
579            connections_per_account: RateLimitWindow::new(100, Duration::from_secs(60))
580                .with_burst(20),
581            subscriptions_per_connection: RateLimitWindow::new(120, Duration::from_secs(60))
582                .with_burst(10),
583            messages_per_connection: RateLimitWindow::new(1000, Duration::from_secs(60))
584                .with_burst(100),
585            snapshots_per_connection: RateLimitWindow::new(30, Duration::from_secs(60))
586                .with_burst(5),
587        }
588    }
589
590    #[tokio::test]
591    async fn test_rate_limiter_allows_within_limit() {
592        let config = RateLimiterConfig {
593            handshake_per_ip: RateLimitWindow::new(5, Duration::from_secs(60)),
594            ..test_config()
595        };
596        let limiter = WebSocketRateLimiter::new(config);
597
598        let addr: SocketAddr = "127.0.0.1:12345".parse().unwrap();
599
600        // Should allow first 5 requests
601        for i in 0..5 {
602            let result = limiter.check_handshake(addr).await;
603            match result {
604                RateLimitResult::Allowed { remaining, .. } => {
605                    assert_eq!(
606                        remaining,
607                        4 - i,
608                        "Request {} should have {} remaining",
609                        i,
610                        4 - i
611                    );
612                }
613                RateLimitResult::Denied { .. } => {
614                    panic!("Request {} should be allowed", i);
615                }
616            }
617        }
618    }
619
620    #[tokio::test]
621    async fn test_rate_limiter_denies_over_limit() {
622        let config = RateLimiterConfig {
623            handshake_per_ip: RateLimitWindow::new(2, Duration::from_secs(60)),
624            ..test_config()
625        };
626        let limiter = WebSocketRateLimiter::new(config);
627
628        let addr: SocketAddr = "127.0.0.1:12345".parse().unwrap();
629
630        // First 2 should be allowed
631        limiter.check_handshake(addr).await;
632        limiter.check_handshake(addr).await;
633
634        // Third should be denied
635        let result = limiter.check_handshake(addr).await;
636        assert!(
637            matches!(result, RateLimitResult::Denied { .. }),
638            "Third request should be denied"
639        );
640    }
641
642    #[tokio::test]
643    async fn test_rate_limiter_with_burst() {
644        let config = RateLimiterConfig {
645            handshake_per_ip: RateLimitWindow::new(2, Duration::from_secs(60)).with_burst(2),
646            ..test_config()
647        };
648        let limiter = WebSocketRateLimiter::new(config);
649
650        let addr: SocketAddr = "127.0.0.1:12345".parse().unwrap();
651
652        // First 4 should be allowed (2 base + 2 burst)
653        for i in 0..4 {
654            let result = limiter.check_handshake(addr).await;
655            assert!(
656                matches!(result, RateLimitResult::Allowed { .. }),
657                "Request {} should be allowed with burst",
658                i
659            );
660        }
661
662        // Fifth should be denied
663        let result = limiter.check_handshake(addr).await;
664        assert!(
665            matches!(result, RateLimitResult::Denied { .. }),
666            "Fifth request should be denied"
667        );
668    }
669
670    #[tokio::test]
671    async fn test_rate_limiter_disabled() {
672        let limiter = WebSocketRateLimiter::new(RateLimiterConfig::disabled());
673
674        let addr: SocketAddr = "127.0.0.1:12345".parse().unwrap();
675
676        // Should allow unlimited when disabled
677        for _ in 0..100 {
678            let result = limiter.check_handshake(addr).await;
679            assert!(
680                matches!(result, RateLimitResult::Allowed { .. }),
681                "Should be allowed when disabled"
682            );
683        }
684    }
685
686    #[tokio::test]
687    async fn test_consumer_rate_limiting() {
688        let config = RateLimiterConfig {
689            connections_per_consumer: RateLimitWindow::new(3, Duration::from_secs(60)),
690            ..test_config()
691        };
692        let limiter = WebSocketRateLimiter::new(config);
693
694        // First 3 connections allowed
695        for i in 0..3 {
696            let result = limiter
697                .check_connection_for_consumer("user-123", None)
698                .await;
699            assert!(
700                matches!(result, RateLimitResult::Allowed { remaining, .. } if remaining == 2 - i),
701                "Connection {} should be allowed",
702                i
703            );
704        }
705
706        // Fourth denied
707        let result = limiter
708            .check_connection_for_consumer("user-123", None)
709            .await;
710        assert!(
711            matches!(result, RateLimitResult::Denied { .. }),
712            "Fourth connection should be denied"
713        );
714
715        // Different consumer should still work
716        let result = limiter
717            .check_connection_for_consumer("user-456", None)
718            .await;
719        assert!(
720            matches!(result, RateLimitResult::Allowed { .. }),
721            "Different consumer should be allowed"
722        );
723    }
724
725    #[tokio::test]
726    async fn signed_limits_override_configured_connection_windows() {
727        let limiter = WebSocketRateLimiter::new(test_config());
728
729        // The signed account limit (2/min) wins over the configured window.
730        for _ in 0..2 {
731            assert!(matches!(
732                limiter
733                    .check_connection_for_account("account:42", Some(2))
734                    .await,
735                RateLimitResult::Allowed { .. }
736            ));
737        }
738        assert!(matches!(
739            limiter
740                .check_connection_for_account("account:42", Some(2))
741                .await,
742            RateLimitResult::Denied { .. }
743        ));
744
745        // A different account is unaffected.
746        assert!(matches!(
747            limiter
748                .check_connection_for_account("account:43", Some(2))
749                .await,
750            RateLimitResult::Allowed { .. }
751        ));
752    }
753
754    #[tokio::test]
755    async fn subscription_create_limits_apply_only_when_signed() {
756        let limiter = WebSocketRateLimiter::new(test_config());
757
758        // No signed limit: allowed and no state is created.
759        for _ in 0..10 {
760            assert!(matches!(
761                limiter
762                    .check_subscription_create_for_account("account:42", None)
763                    .await,
764                RateLimitResult::Allowed { .. }
765            ));
766        }
767        assert!(limiter.account_subscription_buckets.read().await.is_empty());
768
769        // Signed limit of 1/min: second create denied, other accounts fine.
770        assert!(matches!(
771            limiter
772                .check_subscription_create_for_account("account:42", Some(1))
773                .await,
774            RateLimitResult::Allowed { .. }
775        ));
776        assert!(matches!(
777            limiter
778                .check_subscription_create_for_account("account:42", Some(1))
779                .await,
780            RateLimitResult::Denied { .. }
781        ));
782        assert!(matches!(
783            limiter
784                .check_subscription_create_for_consumer("consumer:a", Some(1))
785                .await,
786            RateLimitResult::Allowed { .. }
787        ));
788    }
789
790    #[tokio::test]
791    async fn test_cleanup_stale_buckets_removes_expired_buckets() {
792        let limiter = WebSocketRateLimiter::new(test_config());
793        let stale_request = Instant::now() - Duration::from_secs(600);
794
795        {
796            let mut buckets = limiter.ip_buckets.write().await;
797            let mut bucket = RateLimitBucket::new(limiter.config.handshake_per_ip);
798            bucket.requests.push(stale_request);
799            buckets.insert("127.0.0.1".to_string(), bucket);
800        }
801
802        {
803            let mut buckets = limiter.consumer_buckets.write().await;
804            let mut bucket = RateLimitBucket::new(limiter.config.connections_per_consumer);
805            bucket.requests.push(stale_request);
806            buckets.insert("user-123".to_string(), bucket);
807        }
808
809        {
810            let mut buckets = limiter.account_buckets.write().await;
811            let mut bucket = RateLimitBucket::new(limiter.config.connections_per_account);
812            bucket.requests.push(stale_request);
813            buckets.insert("account-123".to_string(), bucket);
814        }
815
816        limiter.cleanup_stale_buckets().await;
817
818        assert!(limiter.ip_buckets.read().await.is_empty());
819        assert!(limiter.consumer_buckets.read().await.is_empty());
820        assert!(limiter.account_buckets.read().await.is_empty());
821    }
822
823    #[tokio::test]
824    async fn test_remove_client_buckets_clears_connection_specific_state() {
825        let limiter = WebSocketRateLimiter::new(test_config());
826        let client_id = uuid::Uuid::new_v4();
827
828        let _ = limiter.check_subscription(client_id).await;
829        let _ = limiter.check_message(client_id).await;
830        let _ = limiter.check_snapshot(client_id).await;
831
832        assert!(limiter
833            .subscription_buckets
834            .read()
835            .await
836            .contains_key(&client_id));
837        assert!(limiter
838            .message_buckets
839            .read()
840            .await
841            .contains_key(&client_id));
842        assert!(limiter
843            .snapshot_buckets
844            .read()
845            .await
846            .contains_key(&client_id));
847
848        limiter.remove_client_buckets(client_id).await;
849
850        assert!(!limiter
851            .subscription_buckets
852            .read()
853            .await
854            .contains_key(&client_id));
855        assert!(!limiter
856            .message_buckets
857            .read()
858            .await
859            .contains_key(&client_id));
860        assert!(!limiter
861            .snapshot_buckets
862            .read()
863            .await
864            .contains_key(&client_id));
865    }
866}
arete_server/websocket/rate_limiter.rs

arete_server/websocket/
rate_limiter.rs
