arete_server/websocket/
client_manager.rs

1use crate::compression::CompressedPayload;
2use crate::websocket::auth::{AuthContext, AuthDeny};
3use crate::websocket::rate_limiter::{RateLimitResult, WebSocketRateLimiter};
4use arete_auth::Limits;
5use bytes::Bytes;
6use dashmap::DashMap;
7use futures_util::stream::SplitSink;
8use futures_util::SinkExt;
9use std::collections::HashMap;
10use std::net::SocketAddr;
11use std::sync::Arc;
12use std::time::{Duration, SystemTime};
13use tokio::net::TcpStream;
14use tokio::sync::{mpsc, RwLock};
15use tokio_tungstenite::{tungstenite::Message, WebSocketStream};
16use tokio_util::sync::CancellationToken;
17use tracing::{debug, info, warn};
18use uuid::Uuid;
19
20pub type WebSocketSender = SplitSink<WebSocketStream<TcpStream>, Message>;
21
22/// Error type for send operations
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum SendError {
25    /// Client not found in registry
26    ClientNotFound,
27    /// Client's message queue is full - client was disconnected
28    ClientBackpressured,
29    /// Client's channel is closed - client was disconnected
30    ClientDisconnected,
31}
32
33impl std::fmt::Display for SendError {
34    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
35        match self {
36            SendError::ClientNotFound => write!(f, "client not found"),
37            SendError::ClientBackpressured => write!(f, "client backpressured and disconnected"),
38            SendError::ClientDisconnected => write!(f, "client disconnected"),
39        }
40    }
41}
42
43impl std::error::Error for SendError {}
44
45/// Egress tracking for a client
46#[derive(Debug)]
47struct EgressTracker {
48    /// Bytes sent in the current minute window
49    bytes_this_minute: u64,
50    /// Start of the current minute window
51    window_start: SystemTime,
52}
53
54/// Inbound message-rate tracking for a client
55#[derive(Debug)]
56struct MessageRateTracker {
57    messages_this_minute: u32,
58    window_start: SystemTime,
59}
60
61impl MessageRateTracker {
62    fn new() -> Self {
63        Self {
64            messages_this_minute: 0,
65            window_start: SystemTime::now(),
66        }
67    }
68
69    fn maybe_reset_window(&mut self) {
70        let now = SystemTime::now();
71        if now.duration_since(self.window_start).unwrap_or_default() >= Duration::from_secs(60) {
72            self.messages_this_minute = 0;
73            self.window_start = now;
74        }
75    }
76
77    fn record_message(&mut self, limit: u32) -> bool {
78        self.maybe_reset_window();
79        if self.messages_this_minute + 1 > limit {
80            false
81        } else {
82            self.messages_this_minute += 1;
83            true
84        }
85    }
86
87    fn current_usage(&mut self) -> u32 {
88        self.maybe_reset_window();
89        self.messages_this_minute
90    }
91}
92
93impl EgressTracker {
94    fn new() -> Self {
95        Self {
96            bytes_this_minute: 0,
97            window_start: SystemTime::now(),
98        }
99    }
100
101    /// Check if we need to reset the window (new minute)
102    fn maybe_reset_window(&mut self) {
103        let now = SystemTime::now();
104        if now.duration_since(self.window_start).unwrap_or_default() >= Duration::from_secs(60) {
105            self.bytes_this_minute = 0;
106            self.window_start = now;
107        }
108    }
109
110    /// Record bytes sent, returning true if within limit
111    fn record_bytes(&mut self, bytes: usize, limit: u64) -> bool {
112        self.maybe_reset_window();
113        let bytes_u64 = bytes as u64;
114        if self.bytes_this_minute + bytes_u64 > limit {
115            false
116        } else {
117            self.bytes_this_minute += bytes_u64;
118            true
119        }
120    }
121
122    /// Get current usage
123    fn current_usage(&mut self) -> u64 {
124        self.maybe_reset_window();
125        self.bytes_this_minute
126    }
127}
128
129/// Information about a connected client
130#[derive(Debug)]
131pub struct ClientInfo {
132    pub id: Uuid,
133    pub last_seen: SystemTime,
134    pub sender: mpsc::Sender<Message>,
135    subscriptions: Arc<RwLock<HashMap<String, CancellationToken>>>,
136    /// Authentication context for this client
137    pub auth_context: Option<AuthContext>,
138    /// Client's IP address for rate limiting
139    pub remote_addr: SocketAddr,
140    /// Egress tracking for rate limiting
141    egress_tracker: std::sync::Mutex<EgressTracker>,
142    /// Inbound message-rate tracking for rate limiting
143    message_rate_tracker: std::sync::Mutex<MessageRateTracker>,
144}
145
146impl ClientInfo {
147    pub fn new(
148        id: Uuid,
149        sender: mpsc::Sender<Message>,
150        auth_context: Option<AuthContext>,
151        remote_addr: SocketAddr,
152    ) -> Self {
153        Self {
154            id,
155            last_seen: SystemTime::now(),
156            sender,
157            subscriptions: Arc::new(RwLock::new(HashMap::new())),
158            auth_context,
159            remote_addr,
160            egress_tracker: std::sync::Mutex::new(EgressTracker::new()),
161            message_rate_tracker: std::sync::Mutex::new(MessageRateTracker::new()),
162        }
163    }
164
165    /// Record bytes sent, returning true if within limit
166    pub fn record_egress(&self, bytes: usize) -> Option<u64> {
167        if let Ok(mut tracker) = self.egress_tracker.lock() {
168            if let Some(ref ctx) = self.auth_context {
169                if let Some(limit) = ctx.limits.max_bytes_per_minute {
170                    if tracker.record_bytes(bytes, limit) {
171                        return Some(tracker.current_usage());
172                    } else {
173                        return None; // Limit exceeded
174                    }
175                }
176            }
177            // No limit set, return current usage (0)
178            return Some(tracker.current_usage());
179        }
180        None
181    }
182
183    /// Record an inbound client message, returning true if within limit.
184    pub fn record_inbound_message(&self) -> Option<u32> {
185        if let Ok(mut tracker) = self.message_rate_tracker.lock() {
186            if let Some(ref ctx) = self.auth_context {
187                if let Some(limit) = ctx.limits.max_messages_per_minute {
188                    if tracker.record_message(limit) {
189                        return Some(tracker.current_usage());
190                    } else {
191                        return None;
192                    }
193                }
194            }
195
196            return Some(tracker.current_usage());
197        }
198
199        None
200    }
201
202    pub fn update_last_seen(&mut self) {
203        self.last_seen = SystemTime::now();
204    }
205
206    pub fn is_stale(&self, timeout: Duration) -> bool {
207        self.last_seen.elapsed().unwrap_or(Duration::MAX) > timeout
208    }
209
210    pub async fn add_subscription(
211        &self,
212        subscription_id: String,
213        token: CancellationToken,
214    ) -> bool {
215        let mut subs = self.subscriptions.write().await;
216        match subs.entry(subscription_id) {
217            std::collections::hash_map::Entry::Vacant(entry) => {
218                entry.insert(token);
219                true
220            }
221            std::collections::hash_map::Entry::Occupied(_) => false,
222        }
223    }
224
225    pub async fn remove_subscription(&self, subscription_id: &str) -> bool {
226        let mut subs = self.subscriptions.write().await;
227        if let Some(token) = subs.remove(subscription_id) {
228            token.cancel();
229            debug!("Cancelled subscription: {}", subscription_id);
230            true
231        } else {
232            debug!(
233                "Subscription not found for cancellation: {}",
234                subscription_id
235            );
236            false
237        }
238    }
239
240    pub async fn cancel_all_subscriptions(&self) {
241        let subs = self.subscriptions.read().await;
242        for (subscription_id, token) in subs.iter() {
243            token.cancel();
244            debug!("Cancelled subscription on disconnect: {}", subscription_id);
245        }
246    }
247
248    pub async fn subscription_count(&self) -> usize {
249        self.subscriptions.read().await.len()
250    }
251}
252
253/// Configuration for rate limiting in ClientManager
254///
255/// These settings control various rate limits at the connection level.
256/// Per-subject limits are controlled via AuthContext.Limits.
257#[derive(Debug, Clone)]
258pub struct RateLimitConfig {
259    /// Global maximum connections per IP address
260    pub max_connections_per_ip: Option<usize>,
261    /// Global maximum connections per metering key
262    pub max_connections_per_metering_key: Option<usize>,
263    /// Global maximum connections per origin
264    pub max_connections_per_origin: Option<usize>,
265    /// Default connection timeout for stale client cleanup
266    pub client_timeout: Duration,
267    /// Message queue size per client
268    pub message_queue_size: usize,
269    /// Maximum reconnect attempts per client (optional global default)
270    pub max_reconnect_attempts: Option<u32>,
271    /// Rate limit window duration for message counting
272    pub message_rate_window: Duration,
273    /// Rate limit window duration for egress tracking
274    pub egress_rate_window: Duration,
275    /// Default limits applied when auth token doesn't specify limits
276    /// These act as server-wide fallback limits for all connections
277    pub default_limits: Option<Limits>,
278}
279
280impl Default for RateLimitConfig {
281    fn default() -> Self {
282        Self {
283            max_connections_per_ip: None,
284            max_connections_per_metering_key: None,
285            max_connections_per_origin: None,
286            client_timeout: Duration::from_secs(300),
287            message_queue_size: 512,
288            max_reconnect_attempts: None,
289            message_rate_window: Duration::from_secs(60),
290            egress_rate_window: Duration::from_secs(60),
291            default_limits: None,
292        }
293    }
294}
295
296impl RateLimitConfig {
297    /// Load configuration from environment variables
298    ///
299    /// Environment variables:
300    /// - `ARETE_WS_MAX_CONNECTIONS_PER_IP` - Max connections per IP (default: unlimited)
301    /// - `ARETE_WS_MAX_CONNECTIONS_PER_METERING_KEY` - Max connections per metering key (default: unlimited)
302    /// - `ARETE_WS_MAX_CONNECTIONS_PER_ORIGIN` - Max connections per origin (default: unlimited)
303    /// - `ARETE_WS_CLIENT_TIMEOUT_SECS` - Client timeout in seconds (default: 300)
304    /// - `ARETE_WS_MESSAGE_QUEUE_SIZE` - Message queue size per client (default: 512)
305    /// - `ARETE_WS_RATE_LIMIT_WINDOW_SECS` - Rate limit window in seconds (default: 60)
306    /// - `ARETE_WS_DEFAULT_MAX_CONNECTIONS` - Default max connections per subject (fallback when token has no limit)
307    /// - `ARETE_WS_DEFAULT_MAX_SUBSCRIPTIONS` - Default max subscriptions per connection (fallback when token has no limit)
308    /// - `ARETE_WS_DEFAULT_MAX_SNAPSHOT_ROWS` - Default max snapshot rows per request (fallback when token has no limit)
309    /// - `ARETE_WS_DEFAULT_MAX_MESSAGES_PER_MINUTE` - Default max messages per minute (fallback when token has no limit)
310    /// - `ARETE_WS_DEFAULT_MAX_BYTES_PER_MINUTE` - Default max bytes per minute (fallback when token has no limit)
311    pub fn from_env() -> Self {
312        let mut config = Self::default();
313
314        if let Ok(val) = std::env::var("ARETE_WS_MAX_CONNECTIONS_PER_IP") {
315            if let Ok(max) = val.parse() {
316                config.max_connections_per_ip = Some(max);
317            }
318        }
319
320        if let Ok(val) = std::env::var("ARETE_WS_MAX_CONNECTIONS_PER_METERING_KEY") {
321            if let Ok(max) = val.parse() {
322                config.max_connections_per_metering_key = Some(max);
323            }
324        }
325
326        if let Ok(val) = std::env::var("ARETE_WS_MAX_CONNECTIONS_PER_ORIGIN") {
327            if let Ok(max) = val.parse() {
328                config.max_connections_per_origin = Some(max);
329            }
330        }
331
332        if let Ok(val) = std::env::var("ARETE_WS_CLIENT_TIMEOUT_SECS") {
333            if let Ok(secs) = val.parse() {
334                config.client_timeout = Duration::from_secs(secs);
335            }
336        }
337
338        if let Ok(val) = std::env::var("ARETE_WS_MESSAGE_QUEUE_SIZE") {
339            if let Ok(size) = val.parse() {
340                config.message_queue_size = size;
341            }
342        }
343
344        if let Ok(val) = std::env::var("ARETE_WS_RATE_LIMIT_WINDOW_SECS") {
345            if let Ok(secs) = val.parse() {
346                config.message_rate_window = Duration::from_secs(secs);
347                config.egress_rate_window = Duration::from_secs(secs);
348            }
349        }
350
351        // Load default limits from environment (fallback when auth token doesn't specify limits)
352        let mut default_limits = Limits::default();
353        let mut has_default_limits = false;
354
355        if let Ok(val) = std::env::var("ARETE_WS_DEFAULT_MAX_CONNECTIONS") {
356            if let Ok(max) = val.parse() {
357                default_limits.max_connections = Some(max);
358                has_default_limits = true;
359            }
360        }
361
362        if let Ok(val) = std::env::var("ARETE_WS_DEFAULT_MAX_SUBSCRIPTIONS") {
363            if let Ok(max) = val.parse() {
364                default_limits.max_subscriptions = Some(max);
365                has_default_limits = true;
366            }
367        }
368
369        if let Ok(val) = std::env::var("ARETE_WS_DEFAULT_MAX_SNAPSHOT_ROWS") {
370            if let Ok(max) = val.parse() {
371                default_limits.max_snapshot_rows = Some(max);
372                has_default_limits = true;
373            }
374        }
375
376        if let Ok(val) = std::env::var("ARETE_WS_DEFAULT_MAX_MESSAGES_PER_MINUTE") {
377            if let Ok(max) = val.parse() {
378                default_limits.max_messages_per_minute = Some(max);
379                has_default_limits = true;
380            }
381        }
382
383        if let Ok(val) = std::env::var("ARETE_WS_DEFAULT_MAX_BYTES_PER_MINUTE") {
384            if let Ok(max) = val.parse() {
385                default_limits.max_bytes_per_minute = Some(max);
386                has_default_limits = true;
387            }
388        }
389
390        if has_default_limits {
391            config.default_limits = Some(default_limits);
392        }
393
394        config
395    }
396
397    /// Set the maximum connections per IP
398    pub fn with_max_connections_per_ip(mut self, max: usize) -> Self {
399        self.max_connections_per_ip = Some(max);
400        self
401    }
402
403    /// Set the client timeout
404    pub fn with_timeout(mut self, timeout: Duration) -> Self {
405        self.client_timeout = timeout;
406        self
407    }
408
409    /// Set the message queue size
410    pub fn with_message_queue_size(mut self, size: usize) -> Self {
411        self.message_queue_size = size;
412        self
413    }
414
415    /// Set the rate limit window (applies to both message and egress windows)
416    pub fn with_rate_limit_window(mut self, window: Duration) -> Self {
417        self.message_rate_window = window;
418        self.egress_rate_window = window;
419        self
420    }
421
422    /// Set default limits applied when auth token doesn't specify limits
423    ///
424    /// These limits act as server-wide fallbacks for connections
425    /// where the authentication token doesn't include explicit limits.
426    pub fn with_default_limits(mut self, limits: Limits) -> Self {
427        self.default_limits = Some(limits);
428        self
429    }
430}
431
432/// Manages all connected WebSocket clients using lock-free DashMap.
433///
434/// Key design decisions:
435/// - Uses DashMap for lock-free concurrent access to client registry
436/// - Uses try_send instead of send to never block on slow clients
437/// - Disconnects clients that are backpressured (queue full) to prevent cascade failures
438/// - All public methods are non-blocking or use fine-grained per-key locks
439/// - Supports configurable rate limiting per IP, subject, and global defaults
440#[derive(Clone)]
441pub struct ClientManager {
442    clients: Arc<DashMap<Uuid, ClientInfo>>,
443    rate_limit_config: RateLimitConfig,
444    /// Optional WebSocket rate limiter for granular rate control
445    rate_limiter: Option<Arc<WebSocketRateLimiter>>,
446}
447
448impl ClientManager {
449    pub fn new() -> Self {
450        Self::with_config(RateLimitConfig::default())
451    }
452
453    /// Create a new ClientManager with the given rate limit configuration
454    pub fn with_config(config: RateLimitConfig) -> Self {
455        Self {
456            clients: Arc::new(DashMap::new()),
457            rate_limit_config: config,
458            rate_limiter: None,
459        }
460    }
461
462    /// Load configuration from environment variables
463    ///
464    /// See `RateLimitConfig::from_env` for supported variables.
465    pub fn from_env() -> Self {
466        Self::with_config(RateLimitConfig::from_env())
467    }
468
469    /// Set the client timeout for stale client cleanup
470    pub fn with_timeout(mut self, timeout: Duration) -> Self {
471        self.rate_limit_config.client_timeout = timeout;
472        self
473    }
474
475    /// Set the message queue size per client
476    pub fn with_message_queue_size(mut self, queue_size: usize) -> Self {
477        self.rate_limit_config.message_queue_size = queue_size;
478        self
479    }
480
481    /// Set a global limit on connections per IP address
482    pub fn with_max_connections_per_ip(mut self, max: usize) -> Self {
483        self.rate_limit_config.max_connections_per_ip = Some(max);
484        self
485    }
486
487    /// Set the rate limit window duration
488    pub fn with_rate_limit_window(mut self, window: Duration) -> Self {
489        self.rate_limit_config.message_rate_window = window;
490        self.rate_limit_config.egress_rate_window = window;
491        self
492    }
493
494    /// Set default limits applied when auth token doesn't specify limits
495    ///
496    /// These limits act as server-wide fallbacks for connections
497    /// where the authentication token doesn't include explicit limits.
498    pub fn with_default_limits(mut self, limits: Limits) -> Self {
499        self.rate_limit_config.default_limits = Some(limits);
500        self
501    }
502
503    /// Set a WebSocket rate limiter for granular rate control
504    pub fn with_rate_limiter(mut self, rate_limiter: Arc<WebSocketRateLimiter>) -> Self {
505        self.rate_limiter = Some(rate_limiter);
506        self
507    }
508
509    /// Get the rate limiter if configured
510    pub fn rate_limiter(&self) -> Option<&WebSocketRateLimiter> {
511        self.rate_limiter.as_ref().map(|r| r.as_ref())
512    }
513
514    /// Get the current rate limit configuration
515    pub fn rate_limit_config(&self) -> &RateLimitConfig {
516        &self.rate_limit_config
517    }
518
519    /// Add a new client connection.
520    ///
521    /// Spawns a dedicated sender task for this client that reads from its mpsc channel
522    /// and writes to the WebSocket. If the WebSocket write fails, the client is automatically
523    /// removed from the registry.
524    pub fn add_client(
525        &self,
526        client_id: Uuid,
527        mut ws_sender: WebSocketSender,
528        auth_context: Option<AuthContext>,
529        remote_addr: SocketAddr,
530    ) {
531        let (client_tx, mut client_rx) =
532            mpsc::channel::<Message>(self.rate_limit_config.message_queue_size);
533        let client_info = ClientInfo::new(client_id, client_tx, auth_context, remote_addr);
534
535        let clients_ref = self.clients.clone();
536        tokio::spawn(async move {
537            while let Some(message) = client_rx.recv().await {
538                if let Err(e) = ws_sender.send(message).await {
539                    warn!("Failed to send message to client {}: {}", client_id, e);
540                    break;
541                }
542            }
543            clients_ref.remove(&client_id);
544            debug!("WebSocket sender task for client {} stopped", client_id);
545        });
546
547        self.clients.insert(client_id, client_info);
548        info!("Client {} registered from {}", client_id, remote_addr);
549    }
550
551    /// Remove a client from the registry.
552    pub fn remove_client(&self, client_id: Uuid) {
553        if self.clients.remove(&client_id).is_some() {
554            info!("Client {} removed", client_id);
555        }
556    }
557
558    /// Update the auth context for a client.
559    ///
560    /// Used for in-band auth refresh without reconnecting.
561    pub fn update_client_auth(&self, client_id: Uuid, auth_context: AuthContext) -> bool {
562        if let Some(mut client) = self.clients.get_mut(&client_id) {
563            client.auth_context = Some(auth_context);
564            debug!("Updated auth context for client {}", client_id);
565            true
566        } else {
567            false
568        }
569    }
570
571    /// Check if a client's token has expired.
572    ///
573    /// Returns true if the client has an auth context and it has expired.
574    /// If expired, the client is removed from the registry.
575    pub fn check_and_remove_expired(&self, client_id: Uuid) -> bool {
576        if let Some(client) = self.clients.get(&client_id) {
577            if let Some(ref ctx) = client.auth_context {
578                let now = std::time::SystemTime::now()
579                    .duration_since(std::time::UNIX_EPOCH)
580                    .unwrap_or_default()
581                    .as_secs();
582                if ctx.expires_at <= now {
583                    warn!(
584                        "Client {} token expired (expired at {}), disconnecting",
585                        client_id, ctx.expires_at
586                    );
587                    self.clients.remove(&client_id);
588                    return true;
589                }
590            }
591        }
592        false
593    }
594
595    /// Get the current number of connected clients.
596    ///
597    /// This is lock-free and returns an approximate count (may be slightly stale
598    /// under high concurrency, which is fine for max_clients checks).
599    pub fn client_count(&self) -> usize {
600        self.clients.len()
601    }
602
603    /// Send data to a specific client (non-blocking).
604    ///
605    /// This method NEVER blocks. If the client's queue is full, the client is
606    /// considered too slow and is disconnected to prevent cascade failures.
607    /// Use this for live streaming updates.
608    ///
609    /// For initial snapshots where you expect to send many messages at once,
610    /// use `send_to_client_async` instead which will wait for queue space.
611    pub fn send_to_client(&self, client_id: Uuid, data: Arc<Bytes>) -> Result<(), SendError> {
612        // Check if client token has expired before sending
613        if self.check_and_remove_expired(client_id) {
614            return Err(SendError::ClientDisconnected);
615        }
616
617        // Check egress limits
618        if let Some(client) = self.clients.get(&client_id) {
619            if client.record_egress(data.len()).is_none() {
620                warn!("Client {} exceeded egress limit, disconnecting", client_id);
621                self.clients.remove(&client_id);
622                return Err(SendError::ClientDisconnected);
623            }
624        } else {
625            return Err(SendError::ClientNotFound);
626        }
627
628        let sender = {
629            let client = self
630                .clients
631                .get(&client_id)
632                .ok_or(SendError::ClientNotFound)?;
633            client.sender.clone()
634        };
635
636        let msg = Message::Binary((*data).clone());
637        match sender.try_send(msg) {
638            Ok(()) => Ok(()),
639            Err(mpsc::error::TrySendError::Full(_)) => {
640                warn!(
641                    "Client {} backpressured (queue full), disconnecting",
642                    client_id
643                );
644                self.clients.remove(&client_id);
645                Err(SendError::ClientBackpressured)
646            }
647            Err(mpsc::error::TrySendError::Closed(_)) => {
648                debug!("Client {} channel closed", client_id);
649                self.clients.remove(&client_id);
650                Err(SendError::ClientDisconnected)
651            }
652        }
653    }
654
655    /// Send data to a specific client (async, waits for queue space).
656    ///
657    /// This method will wait if the client's queue is full, allowing the client
658    /// time to catch up. Use this for initial snapshots where you need to send
659    /// many messages at once.
660    ///
661    /// For live streaming updates, use `send_to_client` instead which will
662    /// disconnect slow clients rather than blocking.
663    pub async fn send_to_client_async(
664        &self,
665        client_id: Uuid,
666        data: Arc<Bytes>,
667    ) -> Result<(), SendError> {
668        // Check if client token has expired before sending
669        if self.check_and_remove_expired(client_id) {
670            return Err(SendError::ClientDisconnected);
671        }
672
673        // Check egress limits
674        if let Some(client) = self.clients.get(&client_id) {
675            if client.record_egress(data.len()).is_none() {
676                warn!("Client {} exceeded egress limit, disconnecting", client_id);
677                self.clients.remove(&client_id);
678                return Err(SendError::ClientDisconnected);
679            }
680        } else {
681            return Err(SendError::ClientNotFound);
682        }
683
684        let sender = {
685            let client = self
686                .clients
687                .get(&client_id)
688                .ok_or(SendError::ClientNotFound)?;
689            client.sender.clone()
690        };
691
692        let msg = Message::Binary((*data).clone());
693        sender
694            .send(msg)
695            .await
696            .map_err(|_| SendError::ClientDisconnected)
697    }
698
699    /// Send a text message to a specific client (async).
700    ///
701    /// This method sends a text message directly to the client's WebSocket.
702    /// Used for control messages like auth refresh responses.
703    pub async fn send_text_to_client(
704        &self,
705        client_id: Uuid,
706        text: String,
707    ) -> Result<(), SendError> {
708        // Check if client token has expired before sending
709        if self.check_and_remove_expired(client_id) {
710            return Err(SendError::ClientDisconnected);
711        }
712
713        let sender = {
714            let client = self
715                .clients
716                .get(&client_id)
717                .ok_or(SendError::ClientNotFound)?;
718            client.sender.clone()
719        };
720
721        let msg = Message::Text(text.into());
722        sender
723            .send(msg)
724            .await
725            .map_err(|_| SendError::ClientDisconnected)
726    }
727
728    /// Send a potentially compressed payload to a client (async).
729    ///
730    /// Compressed payloads are sent as binary frames (raw gzip).
731    /// Uncompressed payloads are sent as text frames (JSON).
732    pub async fn send_compressed_async(
733        &self,
734        client_id: Uuid,
735        payload: CompressedPayload,
736    ) -> Result<(), SendError> {
737        // Check if client token has expired before sending
738        if self.check_and_remove_expired(client_id) {
739            return Err(SendError::ClientDisconnected);
740        }
741
742        let (sender, bytes_to_record) = {
743            let client = self
744                .clients
745                .get(&client_id)
746                .ok_or(SendError::ClientNotFound)?;
747
748            let bytes = match &payload {
749                CompressedPayload::Compressed(bytes) => bytes.len(),
750                CompressedPayload::Uncompressed(bytes) => bytes.len(),
751            };
752
753            (client.sender.clone(), bytes)
754        };
755
756        // Check egress limits
757        if let Some(client) = self.clients.get(&client_id) {
758            if client.record_egress(bytes_to_record).is_none() {
759                warn!("Client {} exceeded egress limit, disconnecting", client_id);
760                self.clients.remove(&client_id);
761                return Err(SendError::ClientDisconnected);
762            }
763        }
764
765        let msg = match payload {
766            CompressedPayload::Compressed(bytes) => Message::Binary(bytes),
767            CompressedPayload::Uncompressed(bytes) => Message::Binary(bytes),
768        };
769        sender
770            .send(msg)
771            .await
772            .map_err(|_| SendError::ClientDisconnected)
773    }
774
775    /// Update the last_seen timestamp for a client.
776    pub fn update_client_last_seen(&self, client_id: Uuid) {
777        if let Some(mut client) = self.clients.get_mut(&client_id) {
778            client.update_last_seen();
779        }
780    }
781
782    /// Check whether an inbound message is allowed for a client.
783    #[allow(clippy::result_large_err)]
784    pub fn check_inbound_message_allowed(&self, client_id: Uuid) -> Result<(), AuthDeny> {
785        if self.check_and_remove_expired(client_id) {
786            return Err(AuthDeny::new(
787                crate::websocket::auth::AuthErrorCode::TokenExpired,
788                "Authentication token expired",
789            ));
790        }
791
792        let Some(client) = self.clients.get(&client_id) else {
793            return Err(AuthDeny::new(
794                crate::websocket::auth::AuthErrorCode::InternalError,
795                "Client not found",
796            ));
797        };
798
799        if client.record_inbound_message().is_some() {
800            Ok(())
801        } else {
802            self.clients.remove(&client_id);
803            Err(AuthDeny::rate_limited(
804                self.rate_limit_config.message_rate_window,
805                "inbound websocket messages",
806            )
807            .with_context(format!(
808                "client {} exceeded the inbound message budget",
809                client_id
810            )))
811        }
812    }
813
814    /// Check if a client exists.
815    pub fn has_client(&self, client_id: Uuid) -> bool {
816        self.clients.contains_key(&client_id)
817    }
818
819    pub async fn add_client_subscription(
820        &self,
821        client_id: Uuid,
822        subscription_id: String,
823        token: CancellationToken,
824    ) -> bool {
825        if let Some(client) = self.clients.get(&client_id) {
826            client.add_subscription(subscription_id, token).await
827        } else {
828            false
829        }
830    }
831
832    pub async fn remove_client_subscription(&self, client_id: Uuid, subscription_id: &str) -> bool {
833        if let Some(client) = self.clients.get(&client_id) {
834            client.remove_subscription(subscription_id).await
835        } else {
836            false
837        }
838    }
839
840    pub async fn cancel_all_client_subscriptions(&self, client_id: Uuid) {
841        if let Some(client) = self.clients.get(&client_id) {
842            client.cancel_all_subscriptions().await;
843        }
844    }
845
846    /// Remove stale clients that haven't been seen within the timeout period.
847    pub fn cleanup_stale_clients(&self) -> usize {
848        let timeout = self.rate_limit_config.client_timeout;
849        let mut stale_clients = Vec::new();
850
851        for entry in self.clients.iter() {
852            if entry.value().is_stale(timeout) {
853                stale_clients.push(*entry.key());
854            }
855        }
856
857        let removed_count = stale_clients.len();
858        for client_id in stale_clients {
859            self.clients.remove(&client_id);
860            info!("Removed stale client {}", client_id);
861        }
862
863        removed_count
864    }
865
866    /// Start a background task that periodically cleans up stale clients.
867    pub fn start_cleanup_task(&self) {
868        let client_manager = self.clone();
869
870        tokio::spawn(async move {
871            let mut interval = tokio::time::interval(Duration::from_secs(30));
872
873            loop {
874                interval.tick().await;
875                let removed = client_manager.cleanup_stale_clients();
876                if removed > 0 {
877                    info!("Cleaned up {} stale clients", removed);
878                }
879            }
880        });
881    }
882
883    /// ENFORCEMENT HOOKS
884    ///
885    /// These methods provide hooks for enforcing limits based on auth context.
886    /// They check limits before allowing operations and return errors if limits are exceeded.
887    /// Check if a connection is allowed for the given auth context.
888    ///
889    /// Returns Ok(()) if the connection is allowed, or an error with a reason if not.
890    pub async fn check_connection_allowed(
891        &self,
892        remote_addr: SocketAddr,
893        auth_context: &Option<AuthContext>,
894    ) -> Result<(), AuthDeny> {
895        // Check rate limiter first if configured
896        if let Some(ref rate_limiter) = self.rate_limiter {
897            // Check handshake rate limit for IP
898            match rate_limiter.check_handshake(remote_addr).await {
899                RateLimitResult::Allowed { .. } => {}
900                RateLimitResult::Denied { retry_after, limit } => {
901                    return Err(AuthDeny::rate_limited(retry_after, "websocket handshakes")
902                        .with_context(format!(
903                            "handshake rate limit of {} per minute exceeded for {}",
904                            limit, remote_addr
905                        )));
906                }
907            }
908
909            // Check connection rate limit for subject
910            if let Some(ref ctx) = auth_context {
911                match rate_limiter
912                    .check_connection_for_subject(&ctx.subject)
913                    .await
914                {
915                    RateLimitResult::Allowed { .. } => {}
916                    RateLimitResult::Denied { retry_after, limit } => {
917                        return Err(AuthDeny::rate_limited(retry_after, "websocket connections")
918                            .with_context(format!(
919                                "connection rate limit for subject {} of {} per minute exceeded",
920                                ctx.subject, limit
921                            )));
922                    }
923                }
924
925                // Check connection rate limit for metering key
926                match rate_limiter
927                    .check_connection_for_metering_key(&ctx.metering_key)
928                    .await
929                {
930                    RateLimitResult::Allowed { .. } => {}
931                    RateLimitResult::Denied { retry_after, limit } => {
932                        return Err(AuthDeny::rate_limited(
933                            retry_after,
934                            "metered websocket connections",
935                        )
936                        .with_context(format!(
937                            "connection rate limit for metering key {} of {} per minute exceeded",
938                            ctx.metering_key, limit
939                        )));
940                    }
941                }
942            }
943        }
944
945        // Check global per-IP connection limit
946        if let Some(max_per_ip) = self.rate_limit_config.max_connections_per_ip {
947            let current_ip_connections = self.count_connections_for_ip(&remote_addr);
948            if current_ip_connections >= max_per_ip {
949                return Err(AuthDeny::connection_limit_exceeded(
950                    &format!("ip {}", remote_addr.ip()),
951                    current_ip_connections,
952                    max_per_ip,
953                ));
954            }
955        }
956
957        if let Some(ctx) = auth_context {
958            // Check max connections per subject (use token limits, fallback to default limits)
959            let max_connections = ctx.limits.max_connections.or_else(|| {
960                self.rate_limit_config
961                    .default_limits
962                    .as_ref()
963                    .and_then(|l| l.max_connections)
964            });
965            if let Some(max_connections) = max_connections {
966                let current_connections = self.count_connections_for_subject(&ctx.subject);
967                if current_connections >= max_connections as usize {
968                    return Err(AuthDeny::connection_limit_exceeded(
969                        &format!("subject {}", ctx.subject),
970                        current_connections,
971                        max_connections as usize,
972                    ));
973                }
974            }
975
976            // Check global max connections per metering key
977            if let Some(max_per_metering_key) =
978                self.rate_limit_config.max_connections_per_metering_key
979            {
980                let current_metering_connections =
981                    self.count_connections_for_metering_key(&ctx.metering_key);
982                if current_metering_connections >= max_per_metering_key {
983                    return Err(AuthDeny::connection_limit_exceeded(
984                        &format!("metering key {}", ctx.metering_key),
985                        current_metering_connections,
986                        max_per_metering_key,
987                    ));
988                }
989            }
990
991            // Check global max connections per origin
992            if let Some(max_per_origin) = self.rate_limit_config.max_connections_per_origin {
993                if let Some(ref origin) = ctx.origin {
994                    let current_origin_connections = self.count_connections_for_origin(origin);
995                    if current_origin_connections >= max_per_origin {
996                        return Err(AuthDeny::connection_limit_exceeded(
997                            &format!("origin {}", origin),
998                            current_origin_connections,
999                            max_per_origin,
1000                        ));
1001                    }
1002                }
1003            }
1004        }
1005        Ok(())
1006    }
1007
1008    /// Count connections from a specific IP address
1009    fn count_connections_for_ip(&self, remote_addr: &SocketAddr) -> usize {
1010        let ip = remote_addr.ip();
1011        self.clients
1012            .iter()
1013            .filter(|entry| entry.value().remote_addr.ip() == ip)
1014            .count()
1015    }
1016
1017    /// Count connections for a specific subject
1018    fn count_connections_for_subject(&self, subject: &str) -> usize {
1019        self.clients
1020            .iter()
1021            .filter(|entry| {
1022                entry
1023                    .value()
1024                    .auth_context
1025                    .as_ref()
1026                    .map(|ctx| ctx.subject == subject)
1027                    .unwrap_or(false)
1028            })
1029            .count()
1030    }
1031
1032    /// Count connections for a specific metering key
1033    fn count_connections_for_metering_key(&self, metering_key: &str) -> usize {
1034        self.clients
1035            .iter()
1036            .filter(|entry| {
1037                entry
1038                    .value()
1039                    .auth_context
1040                    .as_ref()
1041                    .map(|ctx| ctx.metering_key == metering_key)
1042                    .unwrap_or(false)
1043            })
1044            .count()
1045    }
1046
1047    /// Count connections for a specific origin
1048    fn count_connections_for_origin(&self, origin: &str) -> usize {
1049        self.clients
1050            .iter()
1051            .filter(|entry| {
1052                entry
1053                    .value()
1054                    .auth_context
1055                    .as_ref()
1056                    .and_then(|ctx| ctx.origin.as_ref())
1057                    .map(|o| o == origin)
1058                    .unwrap_or(false)
1059            })
1060            .count()
1061    }
1062
1063    /// Check if a subscription is allowed for the given client.
1064    ///
1065    /// Returns Ok(()) if the subscription is allowed, or an error with a reason if not.
1066    pub async fn check_subscription_allowed(&self, client_id: Uuid) -> Result<(), AuthDeny> {
1067        if let Some(client) = self.clients.get(&client_id) {
1068            let current_subs = client.subscription_count().await;
1069
1070            // Check max subscriptions per connection (use token limits, fallback to default limits)
1071            if let Some(ref ctx) = client.auth_context {
1072                let max_subs = ctx.limits.max_subscriptions.or_else(|| {
1073                    self.rate_limit_config
1074                        .default_limits
1075                        .as_ref()
1076                        .and_then(|l| l.max_subscriptions)
1077                });
1078                if let Some(max_subs) = max_subs {
1079                    if current_subs >= max_subs as usize {
1080                        return Err(AuthDeny::new(
1081                            crate::websocket::auth::AuthErrorCode::SubscriptionLimitExceeded,
1082                            format!(
1083                                "Subscription limit exceeded: {} of {} subscriptions for client {}",
1084                                current_subs, max_subs, client_id
1085                            ),
1086                        )
1087                        .with_suggested_action(
1088                            "Unsubscribe from an existing view before creating another subscription",
1089                        ));
1090                    }
1091                }
1092            }
1093        }
1094        Ok(())
1095    }
1096
1097    /// Get metering key for a client
1098    pub fn get_metering_key(&self, client_id: Uuid) -> Option<String> {
1099        self.clients.get(&client_id).and_then(|client| {
1100            client
1101                .auth_context
1102                .as_ref()
1103                .map(|ctx| ctx.metering_key.clone())
1104        })
1105    }
1106
1107    /// Get auth context for a client.
1108    pub fn get_auth_context(&self, client_id: Uuid) -> Option<AuthContext> {
1109        self.clients
1110            .get(&client_id)
1111            .and_then(|client| client.auth_context.clone())
1112    }
1113
1114    /// Check if a snapshot request is allowed (based on max_snapshot_rows limit)
1115    ///
1116    /// Uses token limits if available, falls back to default limits from RateLimitConfig.
1117    #[allow(clippy::result_large_err)]
1118    pub fn check_snapshot_allowed(
1119        &self,
1120        client_id: Uuid,
1121        requested_rows: u32,
1122    ) -> Result<(), AuthDeny> {
1123        if let Some(client) = self.clients.get(&client_id) {
1124            if let Some(ref ctx) = client.auth_context {
1125                let max_rows = ctx.limits.max_snapshot_rows.or_else(|| {
1126                    self.rate_limit_config
1127                        .default_limits
1128                        .as_ref()
1129                        .and_then(|l| l.max_snapshot_rows)
1130                });
1131                if let Some(max_rows) = max_rows {
1132                    if requested_rows > max_rows {
1133                        return Err(AuthDeny::new(
1134                            crate::websocket::auth::AuthErrorCode::SnapshotLimitExceeded,
1135                            format!(
1136                                "Snapshot limit exceeded: requested {} rows, max allowed is {} for client {}",
1137                                requested_rows, max_rows, client_id
1138                            ),
1139                        )
1140                        .with_suggested_action(
1141                            "Request fewer rows or lower the snapshotLimit on the subscription",
1142                        ));
1143                    }
1144                }
1145            }
1146        }
1147        Ok(())
1148    }
1149}
1150
1151impl Default for ClientManager {
1152    fn default() -> Self {
1153        Self::new()
1154    }
1155}
1156
1157#[cfg(test)]
1158mod tests {
1159    use super::*;
1160    use crate::websocket::auth::AuthContext;
1161    use arete_auth::{KeyClass, Limits};
1162    use std::net::{IpAddr, Ipv4Addr, SocketAddr};
1163
1164    fn create_test_auth_context(subject: &str, limits: Limits) -> AuthContext {
1165        AuthContext {
1166            subject: subject.to_string(),
1167            issuer: "test-issuer".to_string(),
1168            key_class: KeyClass::Publishable,
1169            metering_key: format!("meter-{}", subject),
1170            deployment_id: None,
1171            expires_at: u64::MAX,
1172            scope: "read".to_string(),
1173            limits,
1174            plan: None,
1175            origin: None,
1176            client_ip: None,
1177            jti: uuid::Uuid::new_v4().to_string(),
1178        }
1179    }
1180
1181    fn create_test_socket_addr(ip: &str) -> SocketAddr {
1182        SocketAddr::new(
1183            ip.parse::<IpAddr>()
1184                .unwrap_or(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1))),
1185            12345,
1186        )
1187    }
1188
1189    #[test]
1190    fn test_egress_tracker_basic() {
1191        let mut tracker = EgressTracker::new();
1192
1193        // Should allow bytes within limit
1194        assert!(tracker.record_bytes(500, 1000));
1195        assert_eq!(tracker.current_usage(), 500);
1196
1197        // Should allow more bytes within limit
1198        assert!(tracker.record_bytes(400, 1000));
1199        assert_eq!(tracker.current_usage(), 900);
1200
1201        // Should reject bytes over limit
1202        assert!(!tracker.record_bytes(200, 1000));
1203        assert_eq!(tracker.current_usage(), 900); // Usage shouldn't increase
1204    }
1205
1206    #[test]
1207    fn test_egress_tracker_window_reset() {
1208        let mut tracker = EgressTracker::new();
1209
1210        // Use up the limit
1211        assert!(tracker.record_bytes(100, 100));
1212        assert!(!tracker.record_bytes(1, 100));
1213
1214        // Reset the window
1215        tracker.bytes_this_minute = 0;
1216        tracker.window_start = SystemTime::now() - Duration::from_secs(61);
1217
1218        // Should allow after window reset
1219        assert!(tracker.record_bytes(50, 100));
1220    }
1221
1222    #[test]
1223    fn test_message_rate_tracker_basic() {
1224        let mut tracker = MessageRateTracker::new();
1225
1226        assert!(tracker.record_message(2));
1227        assert_eq!(tracker.current_usage(), 1);
1228
1229        assert!(tracker.record_message(2));
1230        assert_eq!(tracker.current_usage(), 2);
1231
1232        assert!(!tracker.record_message(2));
1233        assert_eq!(tracker.current_usage(), 2);
1234    }
1235
1236    #[tokio::test]
1237    async fn test_client_inbound_message_limit() {
1238        let (tx, _rx) = mpsc::channel(1);
1239        let client = ClientInfo::new(
1240            Uuid::new_v4(),
1241            tx,
1242            Some(create_test_auth_context(
1243                "user-1",
1244                Limits {
1245                    max_messages_per_minute: Some(2),
1246                    ..Default::default()
1247                },
1248            )),
1249            create_test_socket_addr("127.0.0.1"),
1250        );
1251
1252        assert_eq!(client.record_inbound_message(), Some(1));
1253        assert_eq!(client.record_inbound_message(), Some(2));
1254        assert_eq!(client.record_inbound_message(), None);
1255    }
1256
1257    #[tokio::test]
1258    async fn duplicate_subscription_id_is_rejected_without_replacement() {
1259        let (tx, _rx) = mpsc::channel(1);
1260        let client = ClientInfo::new(
1261            Uuid::new_v4(),
1262            tx,
1263            None,
1264            create_test_socket_addr("127.0.0.1"),
1265        );
1266        let first = CancellationToken::new();
1267        let duplicate = CancellationToken::new();
1268
1269        assert!(
1270            client
1271                .add_subscription("opaque-id".to_string(), first.clone())
1272                .await
1273        );
1274        assert!(
1275            !client
1276                .add_subscription("opaque-id".to_string(), duplicate.clone())
1277                .await
1278        );
1279        assert!(!first.is_cancelled());
1280        assert!(!duplicate.is_cancelled());
1281
1282        assert!(client.remove_subscription("opaque-id").await);
1283        assert!(first.is_cancelled());
1284        assert!(!duplicate.is_cancelled());
1285    }
1286
1287    #[tokio::test]
1288    async fn test_no_limits() {
1289        let manager = ClientManager::new();
1290        let addr = create_test_socket_addr("127.0.0.1");
1291
1292        // No auth context - should succeed
1293        assert!(manager.check_connection_allowed(addr, &None).await.is_ok());
1294
1295        // Auth context with no limits - should succeed
1296        let auth_context = create_test_auth_context("test", Limits::default());
1297        assert!(manager
1298            .check_connection_allowed(addr, &Some(auth_context))
1299            .await
1300            .is_ok());
1301    }
1302
1303    #[tokio::test]
1304    async fn test_per_subject_connection_limit() {
1305        let manager = ClientManager::new();
1306
1307        let limits = Limits {
1308            max_connections: Some(2),
1309            ..Default::default()
1310        };
1311
1312        let auth_context = create_test_auth_context("user-1", limits);
1313        let addr = create_test_socket_addr("127.0.0.1");
1314
1315        // First connection should succeed (no clients yet)
1316        assert!(manager
1317            .check_connection_allowed(addr, &Some(auth_context.clone()))
1318            .await
1319            .is_ok());
1320    }
1321
1322    #[tokio::test]
1323    async fn test_per_ip_connection_limit() {
1324        let manager = ClientManager::new().with_max_connections_per_ip(2);
1325        let addr = create_test_socket_addr("192.168.1.1");
1326
1327        // Should succeed when no connections from that IP
1328        assert!(manager.check_connection_allowed(addr, &None).await.is_ok());
1329    }
1330
1331    // Tests for RateLimitConfig
1332    #[test]
1333    fn rate_limit_config_default() {
1334        let config = RateLimitConfig::default();
1335        assert!(config.max_connections_per_ip.is_none());
1336        assert_eq!(config.client_timeout, Duration::from_secs(300));
1337        assert_eq!(config.message_queue_size, 512);
1338        assert!(config.max_reconnect_attempts.is_none());
1339        assert_eq!(config.message_rate_window, Duration::from_secs(60));
1340        assert_eq!(config.egress_rate_window, Duration::from_secs(60));
1341    }
1342
1343    #[test]
1344    fn rate_limit_config_builder_methods() {
1345        let config = RateLimitConfig::default()
1346            .with_max_connections_per_ip(10)
1347            .with_timeout(Duration::from_secs(600))
1348            .with_message_queue_size(1024)
1349            .with_rate_limit_window(Duration::from_secs(120));
1350
1351        assert_eq!(config.max_connections_per_ip, Some(10));
1352        assert_eq!(config.client_timeout, Duration::from_secs(600));
1353        assert_eq!(config.message_queue_size, 1024);
1354        assert_eq!(config.message_rate_window, Duration::from_secs(120));
1355        assert_eq!(config.egress_rate_window, Duration::from_secs(120));
1356    }
1357
1358    #[tokio::test]
1359    async fn client_manager_with_config() {
1360        let config = RateLimitConfig::default()
1361            .with_max_connections_per_ip(5)
1362            .with_timeout(Duration::from_secs(120))
1363            .with_message_queue_size(256);
1364
1365        let manager = ClientManager::with_config(config);
1366        let addr = create_test_socket_addr("10.0.0.1");
1367
1368        // Check that the configuration was applied
1369        assert_eq!(manager.rate_limit_config().max_connections_per_ip, Some(5));
1370        assert_eq!(
1371            manager.rate_limit_config().client_timeout,
1372            Duration::from_secs(120)
1373        );
1374        assert_eq!(manager.rate_limit_config().message_queue_size, 256);
1375
1376        // Should allow when under limit
1377        assert!(manager.check_connection_allowed(addr, &None).await.is_ok());
1378    }
1379
1380    #[tokio::test]
1381    async fn client_manager_builder_pattern() {
1382        let manager = ClientManager::new()
1383            .with_max_connections_per_ip(10)
1384            .with_timeout(Duration::from_secs(180))
1385            .with_message_queue_size(1024)
1386            .with_rate_limit_window(Duration::from_secs(90));
1387
1388        assert_eq!(manager.rate_limit_config().max_connections_per_ip, Some(10));
1389        assert_eq!(
1390            manager.rate_limit_config().client_timeout,
1391            Duration::from_secs(180)
1392        );
1393        assert_eq!(manager.rate_limit_config().message_queue_size, 1024);
1394        assert_eq!(
1395            manager.rate_limit_config().message_rate_window,
1396            Duration::from_secs(90)
1397        );
1398    }
1399
1400    // Integration test: Connection limits are enforced
1401    #[tokio::test]
1402    async fn connection_limit_enforcement_with_actual_clients() {
1403        let manager = ClientManager::new().with_max_connections_per_ip(2);
1404        let addr1 = create_test_socket_addr("192.168.1.1");
1405        let addr2 = create_test_socket_addr("192.168.1.2");
1406
1407        // First connection from IP1 should succeed
1408        let auth1 = create_test_auth_context("user-1", Limits::default());
1409        assert!(manager
1410            .check_connection_allowed(addr1, &Some(auth1.clone()))
1411            .await
1412            .is_ok());
1413
1414        // Simulate adding a client (we can't easily do this without a real WebSocket,
1415        // but we can verify the check logic works)
1416
1417        // Same IP, different auth context - should still count toward IP limit
1418        let auth2 = create_test_auth_context("user-2", Limits::default());
1419        assert!(manager
1420            .check_connection_allowed(addr1, &Some(auth2.clone()))
1421            .await
1422            .is_ok());
1423
1424        // Different IP - should succeed regardless
1425        let auth3 = create_test_auth_context("user-3", Limits::default());
1426        assert!(manager
1427            .check_connection_allowed(addr2, &Some(auth3.clone()))
1428            .await
1429            .is_ok());
1430    }
1431
1432    // Test subscription limit enforcement
1433    #[tokio::test]
1434    async fn subscription_limit_enforcement() {
1435        let manager = ClientManager::new();
1436        let addr = create_test_socket_addr("127.0.0.1");
1437
1438        // Create auth context with subscription limit
1439        let auth = create_test_auth_context(
1440            "user-1",
1441            Limits {
1442                max_subscriptions: Some(2),
1443                ..Default::default()
1444            },
1445        );
1446
1447        // Check should pass initially
1448        assert!(manager
1449            .check_connection_allowed(addr, &Some(auth.clone()))
1450            .await
1451            .is_ok());
1452
1453        // Note: We can't easily test the full subscription flow without a real connection,
1454        // but we verify the limit configuration is properly stored
1455        assert_eq!(auth.limits.max_subscriptions, Some(2));
1456    }
1457
1458    // Test snapshot limit enforcement
1459    #[tokio::test]
1460    async fn snapshot_limit_enforcement() {
1461        let manager = ClientManager::new();
1462        let addr = create_test_socket_addr("127.0.0.1");
1463
1464        let auth = create_test_auth_context(
1465            "user-1",
1466            Limits {
1467                max_snapshot_rows: Some(1000),
1468                ..Default::default()
1469            },
1470        );
1471
1472        assert!(manager
1473            .check_connection_allowed(addr, &Some(auth.clone()))
1474            .await
1475            .is_ok());
1476
1477        // Note: Actual snapshot limit checking happens in check_snapshot_allowed
1478        // which requires a connected client
1479    }
1480
1481    // Test WebSocketRateLimiter integration
1482    #[tokio::test]
1483    async fn test_rate_limiter_integration() {
1484        use crate::websocket::rate_limiter::{RateLimiterConfig, WebSocketRateLimiter};
1485
1486        let rate_limiter = Arc::new(WebSocketRateLimiter::new(RateLimiterConfig::default()));
1487        let manager = ClientManager::new().with_rate_limiter(rate_limiter);
1488        let addr = create_test_socket_addr("127.0.0.1");
1489
1490        // Should allow connections when rate limiter is configured
1491        let auth = create_test_auth_context("user-1", Limits::default());
1492        assert!(manager
1493            .check_connection_allowed(addr, &Some(auth))
1494            .await
1495            .is_ok());
1496    }
1497}
arete_server/websocket/client_manager.rs

arete_server/websocket/
client_manager.rs
