rustant_core/gateway/
connection.rs

1//! WebSocket connection management.
2
3use chrono::{DateTime, Utc};
4use std::collections::HashMap;
5use uuid::Uuid;
6
7/// Metadata about a connected WebSocket client.
8#[derive(Debug, Clone)]
9pub struct ConnectionInfo {
10    pub connection_id: Uuid,
11    pub authenticated: bool,
12    pub connected_at: DateTime<Utc>,
13    pub last_activity: DateTime<Utc>,
14}
15
16/// Manages active WebSocket connections.
17#[derive(Debug, Default)]
18pub struct ConnectionManager {
19    connections: HashMap<Uuid, ConnectionInfo>,
20    max_connections: usize,
21}
22
23impl ConnectionManager {
24    /// Create a new connection manager with the given capacity limit.
25    pub fn new(max_connections: usize) -> Self {
26        Self {
27            connections: HashMap::new(),
28            max_connections,
29        }
30    }
31
32    /// Register a new connection. Returns `None` if the limit is reached.
33    pub fn add_connection(&mut self) -> Option<Uuid> {
34        if self.connections.len() >= self.max_connections {
35            return None;
36        }
37
38        let id = Uuid::new_v4();
39        let now = Utc::now();
40        self.connections.insert(
41            id,
42            ConnectionInfo {
43                connection_id: id,
44                authenticated: false,
45                connected_at: now,
46                last_activity: now,
47            },
48        );
49        Some(id)
50    }
51
52    /// Remove a connection.
53    pub fn remove_connection(&mut self, id: &Uuid) -> bool {
54        self.connections.remove(id).is_some()
55    }
56
57    /// Mark a connection as authenticated.
58    pub fn authenticate(&mut self, id: &Uuid) -> bool {
59        if let Some(conn) = self.connections.get_mut(id) {
60            conn.authenticated = true;
61            conn.last_activity = Utc::now();
62            true
63        } else {
64            false
65        }
66    }
67
68    /// Update the last activity timestamp for a connection.
69    pub fn touch(&mut self, id: &Uuid) {
70        if let Some(conn) = self.connections.get_mut(id) {
71            conn.last_activity = Utc::now();
72        }
73    }
74
75    /// Get connection info.
76    pub fn get(&self, id: &Uuid) -> Option<&ConnectionInfo> {
77        self.connections.get(id)
78    }
79
80    /// Number of active connections.
81    pub fn active_count(&self) -> usize {
82        self.connections.len()
83    }
84
85    /// Number of authenticated connections.
86    pub fn authenticated_count(&self) -> usize {
87        self.connections
88            .values()
89            .filter(|c| c.authenticated)
90            .count()
91    }
92
93    /// List all connection IDs.
94    pub fn connection_ids(&self) -> Vec<Uuid> {
95        self.connections.keys().copied().collect()
96    }
97
98    /// Check if a connection is authenticated.
99    pub fn is_authenticated(&self, id: &Uuid) -> bool {
100        self.connections
101            .get(id)
102            .map(|c| c.authenticated)
103            .unwrap_or(false)
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    #[test]
112    fn test_add_connection() {
113        let mut mgr = ConnectionManager::new(10);
114        let id = mgr.add_connection();
115        assert!(id.is_some());
116        assert_eq!(mgr.active_count(), 1);
117    }
118
119    #[test]
120    fn test_connection_limit() {
121        let mut mgr = ConnectionManager::new(2);
122        assert!(mgr.add_connection().is_some());
123        assert!(mgr.add_connection().is_some());
124        assert!(mgr.add_connection().is_none()); // limit reached
125        assert_eq!(mgr.active_count(), 2);
126    }
127
128    #[test]
129    fn test_remove_connection() {
130        let mut mgr = ConnectionManager::new(10);
131        let id = mgr.add_connection().unwrap();
132        assert_eq!(mgr.active_count(), 1);
133
134        assert!(mgr.remove_connection(&id));
135        assert_eq!(mgr.active_count(), 0);
136
137        // Removing nonexistent returns false
138        assert!(!mgr.remove_connection(&Uuid::new_v4()));
139    }
140
141    #[test]
142    fn test_authenticate_connection() {
143        let mut mgr = ConnectionManager::new(10);
144        let id = mgr.add_connection().unwrap();
145
146        assert!(!mgr.is_authenticated(&id));
147        assert_eq!(mgr.authenticated_count(), 0);
148
149        assert!(mgr.authenticate(&id));
150        assert!(mgr.is_authenticated(&id));
151        assert_eq!(mgr.authenticated_count(), 1);
152    }
153
154    #[test]
155    fn test_authenticate_nonexistent() {
156        let mut mgr = ConnectionManager::new(10);
157        assert!(!mgr.authenticate(&Uuid::new_v4()));
158    }
159
160    #[test]
161    fn test_connection_ids() {
162        let mut mgr = ConnectionManager::new(10);
163        let id1 = mgr.add_connection().unwrap();
164        let id2 = mgr.add_connection().unwrap();
165
166        let ids = mgr.connection_ids();
167        assert_eq!(ids.len(), 2);
168        assert!(ids.contains(&id1));
169        assert!(ids.contains(&id2));
170    }
171
172    #[test]
173    fn test_get_connection_info() {
174        let mut mgr = ConnectionManager::new(10);
175        let id = mgr.add_connection().unwrap();
176
177        let info = mgr.get(&id).unwrap();
178        assert_eq!(info.connection_id, id);
179        assert!(!info.authenticated);
180
181        assert!(mgr.get(&Uuid::new_v4()).is_none());
182    }
183
184    #[test]
185    fn test_touch_updates_activity() {
186        let mut mgr = ConnectionManager::new(10);
187        let id = mgr.add_connection().unwrap();
188        let initial = mgr.get(&id).unwrap().last_activity;
189
190        // Touch to update
191        mgr.touch(&id);
192        let updated = mgr.get(&id).unwrap().last_activity;
193        assert!(updated >= initial);
194    }
195}
rustant_core/gateway/connection.rs

rustant_core/gateway/
connection.rs
