use async_tungstenite::{WebSocketReceiver, WebSocketSender}; use async_tungstenite::{WebSocketStream, tungstenite::Message}; use futures::SinkExt; use std::sync::{Arc, Weak}; use tokio::sync::Mutex; use tokio::sync::RwLock; use tokio_util::compat::Compat; use tungstenite::Utf8Bytes; use uuid::Uuid; use super::{rho_connection::RhoConnection, rho_manager}; use crate::util::print::PrintType; use crate::util::print::line; use crate::util::print::line_err; use crate::{ auth::auth_connector, // calls::call_manager::CallManager, data::{ communication::{CommunicationType, CommunicationValue, DataTypes}, user::{User, UserStatus}, }, omega::omega_connection::OmegaConnection, }; /// ClientConnection represents a WebSocket connection from a client device pub struct ClientConnection { /// WebSocket session pub sender: Arc>>>, pub receiver: Arc>>>, /// User ID associated with this client pub user_id: Arc>>, /// Whether this connection has been identified/authenticated pub identified: Arc>, /// Ping latency tracking pub ping: Arc>, /// Weak reference to RhoConnection to avoid circular references pub rho_connection: Arc>>>, /// List of user IDs this client is interested in receiving updates about pub interested_users: Arc>>, } impl ClientConnection { /// Create a new ClientConnection pub fn new( sender: WebSocketSender>, receiver: WebSocketReceiver>, ) -> Arc { Arc::new(Self { sender: Arc::new(RwLock::new(sender)), receiver: Arc::new(RwLock::new(receiver)), user_id: Arc::new(RwLock::new(None)), identified: Arc::new(RwLock::new(false)), ping: Arc::new(RwLock::new(-1)), rho_connection: Arc::new(RwLock::new(None)), interested_users: Arc::new(RwLock::new(Vec::new())), }) } /// Get the user ID pub async fn get_user_id(&self) -> Option { *self.user_id.read().await } /// Check if connection is identified pub async fn is_identified(&self) -> bool { *self.identified.read().await } /// Get current ping pub async fn get_ping(&self) -> i64 { *self.ping.read().await } /// Set the RhoConnection reference pub async fn set_rho_connection(&self, rho_connection: Weak) { let mut rho_ref = self.rho_connection.write().await; *rho_ref = Some(rho_connection); } /// Get RhoConnection if available pub async fn get_rho_connection(&self) -> Option> { let rho_ref = self.rho_connection.read().await; if let Some(weak_ref) = rho_ref.as_ref() { weak_ref.upgrade() } else { None } } /// Send a string message to the client pub async fn send_message_str(&self, message: &str) { let mut session = self.sender.write().await; if let Err(e) = session .send(Message::Text(Utf8Bytes::from(message.to_string()))) .await { line_err( PrintType::ClientOut, &format!("Failed to send message to client: {}", e), ); } } /// Send a CommunicationValue to the client pub async fn send_message(&self, cv: &CommunicationValue) { if !cv.is_type(CommunicationType::pong) { line(PrintType::ClientOut, &cv.to_json().to_string()); } self.send_message_str(&cv.to_json().to_string()).await; } /// Handle incoming message from client pub async fn handle_message(self: Arc, message: Utf8Bytes) { let cv = CommunicationValue::from_json(&message); // Handle identification if cv.is_type(CommunicationType::identification) && !self.is_identified().await { self.handle_identification(Arc::clone(&self), cv).await; return; } if !self.is_identified().await { return; } // Handle ping if cv.is_type(CommunicationType::ping) { self.handle_ping(cv).await; return; } line(PrintType::ClientIn, &cv.to_json().to_string()); // Handle client status changes if cv.is_type(CommunicationType::client_changed) { self.handle_client_changed(cv).await; return; } // Handle call invites if cv.is_type(CommunicationType::call_invite) { self.handle_call_invite(cv).await; return; } // Handle get call requests if cv.is_type(CommunicationType::get_call) { self.handle_get_call(cv).await; return; } // Forward other messages to Iota tokio::spawn(async move { self.forward_to_iota(cv).await; }); } /// Handle identification message async fn handle_identification(&self, sarc: Arc, mut cv: CommunicationValue) { // Extract user ID let user_id = match cv.get_data(DataTypes::user_id) { Some(id_str) => match Uuid::parse_str(&id_str.to_string()) { Ok(id) => id, Err(_) => { self.send_error_response( &cv.get_id(), CommunicationType::error_invalid_user_id, ) .await; return; } }, None => { self.send_error_response(&cv.get_id(), CommunicationType::error_invalid_user_id) .await; return; } }; // Validate private key if let Some(private_key_hash) = cv.get_data(DataTypes::private_key_hash) { println!("private_key_hash: {}", private_key_hash); let is_valid = auth_connector::is_private_key_valid(user_id, &private_key_hash.to_string()).await; if !is_valid { println!("Invalid private key"); self.send_error_response( &cv.get_id(), CommunicationType::error_invalid_private_key, ) .await; return; } } else { line(PrintType::ClientIn, "Missing private key"); self.send_error_response(&cv.get_id(), CommunicationType::error_invalid_private_key) .await; return; } // Find RhoConnection for this user let rho_connection = match rho_manager::get_rho_con_for_user(user_id).await { Some(rho) => rho, None => { self.send_error_response(&cv.get_id(), CommunicationType::error_no_iota) .await; return; } }; // Set identification data { let mut user_id_guard = self.user_id.write().await; *user_id_guard = Some(user_id); } { let mut identified_guard = self.identified.write().await; *identified_guard = true; } self.set_rho_connection(Arc::downgrade(&rho_connection)) .await; rho_connection.add_client_connection(Arc::from(sarc)).await; let response = CommunicationValue::new(CommunicationType::identification_response) .with_id(cv.get_id()); self.send_message(&response).await; } /// Handle ping message async fn handle_ping(&self, mut cv: CommunicationValue) { // Update our ping if provided if let Some(last_ping) = cv.get_data(DataTypes::last_ping) { if let Ok(ping_val) = last_ping.to_string().parse::() { let mut ping_guard = self.ping.write().await; *ping_guard = ping_val; } } // Get Iota ping from RhoConnection let iota_ping = if let Some(rho_conn) = self.get_rho_connection().await { rho_conn.get_iota_connection().get_ping().await } else { -1 }; // Send pong response let response = CommunicationValue::new(CommunicationType::pong) .with_id(cv.get_id()) .add_data_str(DataTypes::ping_iota, iota_ping.to_string()); self.send_message(&response).await; } /// Handle client status change async fn handle_client_changed(&self, mut cv: CommunicationValue) { if let Some(user_id) = self.get_user_id().await { if let Some(status_str) = cv.get_data(DataTypes::user_state) { // Parse user status - this would need to be implemented properly let user_status = UserStatus::online; // placeholder if let Some(rho_conn) = self.get_rho_connection().await { OmegaConnection::client_changed( rho_conn.get_iota_id().await, user_id, user_status, ) .await; } } } } /// Handle call invite async fn handle_call_invite(&self, cv: CommunicationValue) { let receiver_id = match cv.get_data(DataTypes::receiver_id) { Some(id_str) => match Uuid::parse_str(&id_str.to_string()) { Ok(id) => id, Err(_) => { self.send_error_response(&cv.get_id(), CommunicationType::error) .await; return; } }, None => { self.send_error_response(&cv.get_id(), CommunicationType::error) .await; return; } }; let call_id = match cv.get_data(DataTypes::call_id) { Some(id_str) => match Uuid::parse_str(&id_str.to_string()) { Ok(id) => id, Err(_) => { self.send_error_response(&cv.get_id(), CommunicationType::error) .await; return; } }, None => { self.send_error_response(&cv.get_id(), CommunicationType::error) .await; return; } }; let call_secret_sha = cv .get_data(DataTypes::call_secret_sha) .map(|s| s.to_string()); let call_secret = cv.get_data(DataTypes::call_secret).map(|s| s.to_string()); if call_secret_sha.is_none() || call_secret.is_none() { self.send_error_response(&cv.get_id(), CommunicationType::error) .await; return; } // Get call group - placeholder implementation // let call_group = CallManager::get_call_group(call_id, &call_secret_sha.unwrap(), true).await; // if call_group.is_none() { // self.send_error_response(&cv.get_id(), CommunicationType::error) // .await; // return; // } // Find target RhoConnection let target_rho = match rho_manager::get_rho_con_for_user(receiver_id).await { Some(rho) => rho, None => { self.send_error_response(&cv.get_id(), CommunicationType::error) .await; return; } }; // Get sender user ID let sender_id = match self.get_user_id().await { Some(id) => id, None => return, }; // Create and send call distribution message let distribute = CommunicationValue::new(CommunicationType::new_call) .with_receiver(receiver_id) .with_sender(sender_id) .add_data_str(DataTypes::call_id, call_id.to_string()) .add_data_str(DataTypes::receiver_id, receiver_id.to_string()) .add_data_str(DataTypes::sender_id, sender_id.to_string()) .add_data_str(DataTypes::call_secret, call_secret.unwrap()); target_rho.message_to_client(distribute).await; // Handle call group invitation logic here // This would require implementing CallGroup::Caller and related functionality // Send success response let response = CommunicationValue::new(CommunicationType::call_invite).with_id(cv.get_id()); self.send_message(&response).await; } /// Handle get call request async fn handle_get_call(&self, mut cv: CommunicationValue) { let user_id = match self.get_user_id().await { Some(id) => id, None => return, }; let call_id = match cv.get_data(DataTypes::call_id) { Some(id_str) => match Uuid::parse_str(&id_str.to_string()) { Ok(id) => id, Err(_) => { self.send_error_response(&cv.get_id(), CommunicationType::error) .await; return; } }, None => { self.send_error_response(&cv.get_id(), CommunicationType::error) .await; return; } }; let call_secret_sha = cv .get_data(DataTypes::call_secret_sha) .map(|s| s.to_string()); if call_secret_sha.is_none() { self.send_error_response(&cv.get_id(), CommunicationType::error) .await; return; } // Get call group let mut response = CommunicationValue::new(CommunicationType::get_call) .with_id(cv.get_id()) .with_receiver(user_id); // Placeholder call group logic // if let Some(call_group) = CallManager::get_call_group(call_id, &call_secret_sha.unwrap(), true).await { // response = response // .add_data_str(DataTypes::call_state, call_group.call_state.to_string()) // .add_data_str(DataTypes::start_date, call_group.started_at.to_string()); // // if call_group.ended_at != 0 { // response = response.add_data_str(DataTypes::end_date, call_group.ended_at.to_string()); // } // } else { response = response.add_data_str(DataTypes::call_state, "DESTROYED".to_string()); // } self.send_message(&response).await; } /// Forward message to Iota async fn forward_to_iota(&self, mut cv: CommunicationValue) { if let Some(user_id) = self.get_user_id().await { if let Some(rho_conn) = self.get_rho_connection().await { let updated_cv = cv.with_sender(user_id); rho_conn.message_to_iota(updated_cv).await; } } } /// Send error response async fn send_error_response(&self, message_id: &Uuid, error_type: CommunicationType) { let error = CommunicationValue::new(error_type).with_id(*message_id); self.send_message(&error).await; } /// Close the connection pub async fn close(&self) { let mut session = self.sender.write().await; let _ = session.close(None).await; } /// Set interested users list pub async fn set_interested_users(&self, interested_ids: Vec) { let mut interested_guard = self.interested_users.write().await; *interested_guard = interested_ids; } /// Check if interested in a user and send notification pub async fn are_you_interested(&self, user: &User) { let interested_guard = self.interested_users.read().await; if interested_guard.contains(&user.user_id) { let notification = CommunicationValue::new(CommunicationType::client_changed) .add_data_str(DataTypes::user_id, user.user_id.to_string()) .add_data_str( DataTypes::user_state, format!("{:?}", user.status.to_string()), ); self.send_message(¬ification).await; } } /// Handle connection close pub async fn handle_close(&self) { if self.is_identified().await { if let Some(user_id) = self.get_user_id().await { if let Some(rho_conn) = rho_manager::get_rho_con_for_user(user_id).await { rho_conn .close_client_connection(Arc::new(self.clone())) .await; } } } } } // Implement Clone to make it easier to work with Arc impl Clone for ClientConnection { fn clone(&self) -> Self { Self { sender: Arc::clone(&self.sender), receiver: Arc::clone(&self.receiver), user_id: Arc::clone(&self.user_id), identified: Arc::clone(&self.identified), ping: Arc::clone(&self.ping), rho_connection: Arc::clone(&self.rho_connection), interested_users: Arc::clone(&self.interested_users), } } } impl std::fmt::Debug for ClientConnection { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("ClientConnection") .field("user_id", &"[async]") .field("identified", &"[async]") .field("ping", &"[async]") .finish() } }