1use crate::auth::AuthContext;
4use crate::connection::ConnectionState;
5use crate::service::ServiceManager;
6use std::collections::HashMap;
7use std::sync::Arc;
8use tokio::sync::RwLock;
9
10#[derive(Debug)]
12pub struct ClientState {
13 pub connection_state: ConnectionState,
15
16 pub auth_contexts: HashMap<String, AuthContext>,
18
19 pub services: ServiceManager,
21
22 pub session_id: Option<String>,
24
25 pub client_id: String,
27
28 pub server_version: Option<String>,
30
31 pub last_activity: chrono::DateTime<chrono::Utc>,
33}
34
35impl ClientState {
36 pub fn new(client_id: String) -> Self {
38 Self {
39 connection_state: ConnectionState::Disconnected,
40 auth_contexts: HashMap::new(),
41 services: ServiceManager::new(),
42 session_id: None,
43 client_id,
44 server_version: None,
45 last_activity: chrono::Utc::now(),
46 }
47 }
48
49 pub fn touch(&mut self) {
51 self.last_activity = chrono::Utc::now();
52 }
53
54 pub fn idle_seconds(&self) -> u64 {
56 (chrono::Utc::now() - self.last_activity)
57 .num_seconds()
58 .max(0) as u64
59 }
60
61 pub fn add_auth_context(&mut self, tenant_id: String, context: AuthContext) {
63 self.auth_contexts.insert(tenant_id, context);
64 }
65
66 pub fn get_auth_context(&self, tenant_id: &str) -> Option<&AuthContext> {
68 self.auth_contexts.get(tenant_id)
69 }
70
71 pub fn is_authenticated_to(&self, tenant_id: &str) -> bool {
73 self.auth_contexts.contains_key(tenant_id)
74 }
75
76 pub fn authenticated_tenants(&self) -> Vec<&str> {
78 self.auth_contexts.keys().map(|s| s.as_str()).collect()
79 }
80
81 pub fn clear_auth(&mut self, tenant_id: &str) {
83 self.auth_contexts.remove(tenant_id);
84 }
85
86 pub fn clear_all_auth(&mut self) {
88 self.auth_contexts.clear();
89 }
90
91 pub fn reset(&mut self) {
93 self.connection_state = ConnectionState::Disconnected;
94 self.session_id = None;
95 self.auth_contexts.clear();
96 self.services.clear();
97 self.last_activity = chrono::Utc::now();
98 }
99}
100
101pub type SharedState = Arc<RwLock<ClientState>>;
103
104pub fn create_shared_state(client_id: String) -> SharedState {
106 Arc::new(RwLock::new(ClientState::new(client_id)))
107}
108
109#[cfg(test)]
110mod tests {
111 use super::*;
112
113 #[test]
114 fn test_client_state_creation() {
115 let state = ClientState::new("client_1".to_string());
116 assert_eq!(state.client_id, "client_1");
117 assert!(state.session_id.is_none());
118 }
119
120 #[test]
121 fn test_idle_calculation() {
122 let mut state = ClientState::new("client_1".to_string());
123 state.last_activity = chrono::Utc::now() - chrono::Duration::seconds(10);
124 assert!(state.idle_seconds() >= 10);
125 }
126
127 #[test]
128 fn test_auth_context_management() {
129 let mut state = ClientState::new("client_1".to_string());
130 let ctx = AuthContext::new("tenant_1".to_string(), vec!["read".to_string()]);
131
132 state.add_auth_context("tenant_1".to_string(), ctx);
133 assert!(state.is_authenticated_to("tenant_1"));
134 assert_eq!(state.authenticated_tenants().len(), 1);
135 }
136}