use crate::calls::call_group::CallGroup; use crate::calls::call_manager; use crate::util::print::PrintType; use crate::util::print::line; use crate::util::print::line_err; use async_tungstenite::WebSocketReceiver; use async_tungstenite::WebSocketSender; use async_tungstenite::tungstenite::Message; use json::JsonValue; use std::{ collections::HashMap, sync::{Arc, Weak}, }; use tokio::sync::RwLock; use tokio_util::compat::Compat; use tungstenite::Utf8Bytes; use uuid::Uuid; use super::{rho_connection::RhoConnection, rho_manager}; use crate::{ auth::auth_connector, // calls::call_manager::CallManager, data::communication::{CommunicationType, CommunicationValue, DataTypes}, omega::omega_connection::OmegaConnection, }; pub struct IotaConnection { pub sender: Arc>>>, pub receiver: Arc>>>, pub iota_id: Arc>, pub user_ids: Arc>>, pub identified: Arc>, pub ping: Arc>, pub rho_connection: Arc>>>, } impl IotaConnection { /// Create a new IotaConnection pub fn new( sender: WebSocketSender>, receiver: WebSocketReceiver>, ) -> Arc { Arc::new(Self { sender: Arc::new(RwLock::new(sender)), receiver: Arc::new(RwLock::new(receiver)), iota_id: Arc::new(RwLock::new(Uuid::nil())), user_ids: Arc::new(RwLock::new(Vec::new())), identified: Arc::new(RwLock::new(false)), ping: Arc::new(RwLock::new(0)), rho_connection: Arc::new(RwLock::new(None)), }) } /// Get the Iota ID pub async fn get_iota_id(&self) -> Uuid { *self.iota_id.read().await } /// Get the user IDs pub async fn get_user_ids(&self) -> Vec { self.user_ids.read().await.clone() } /// 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 message to the Iota 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::IotaOut, &format!("Failed to send WebSocket message: {:?}", e), ); } } /// Send a CommunicationValue to the Iota pub async fn send_message(&self, cv: CommunicationValue) { if !cv.is_type(CommunicationType::pong) { line(PrintType::IotaOut, &cv.to_json().to_string()); } self.send_message_str(&cv.to_json().to_string()).await; } /// Handle incoming message from Iota 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(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::IotaIn, &cv.to_json().to_string()); // Handle forwarding to other Iotas or clients let receiver_id = cv.get_receiver(); if !self.get_user_ids().await.contains(&receiver_id) || cv.is_type(CommunicationType::message_other_iota) || cv.is_type(CommunicationType::send_chat) { self.handle_forward_message(cv).await; return; } // Handle GET_CHATS if cv.is_type(CommunicationType::get_chats) { self.handle_get_chats(cv).await; return; } // Forward to client self.forward_to_client(cv).await; } /// Handle identification message async fn handle_identification(self: Arc, cv: CommunicationValue) { // Parse Iota ID let iota_id = match cv.get_data(DataTypes::iota_id) { Some(id_str) => match Uuid::parse_str(&id_str.to_string()) { Ok(id) => id, Err(_) => { self.send_error_response(&cv.get_id()).await; return; } }, None => { self.send_error_response(&cv.get_id()).await; return; } }; // Parse user IDs let mut validated_user_ids: Vec = Vec::new(); if let Some(user_ids_str) = cv.get_data(DataTypes::user_ids) { for id_str in user_ids_str.to_string().split(',') { if id_str.is_empty() { continue; } if let Ok(user_id) = Uuid::parse_str(id_str.trim()) { if let Some(auth_iota_id) = auth_connector::get_iota_id(user_id).await { if auth_iota_id == iota_id { validated_user_ids.push(user_id); } } else { line(PrintType::IotaIn, "User ID not found"); } } } } // Set identification data { let mut iota_id_guard = self.iota_id.write().await; *iota_id_guard = iota_id; } { let mut user_ids_guard = self.user_ids.write().await; *user_ids_guard = validated_user_ids.clone(); } { let mut identified_guard = self.identified.write().await; *identified_guard = true; } // Check for existing connection and close it if rho_manager::contains_iota(iota_id).await { if let Some(existing_rho) = rho_manager::get_rho_by_iota(iota_id).await { existing_rho.close_iota_connection().await; } } // Create RhoConnection let rho_connection = Arc::new(RhoConnection::new(self.clone(), validated_user_ids.clone()).await); // Set up bidirectional reference self.set_rho_connection(Arc::downgrade(&rho_connection)) .await; // Add to manager rho_manager::add_rho(rho_connection).await; // Send response let mut str = String::new(); for id in &validated_user_ids { str.push_str(&format!(",{}", id)); } let response = CommunicationValue::new(CommunicationType::identification_response) .with_id(cv.get_id()) .add_data_str(DataTypes::accepted_ids, str) .add_data_str(DataTypes::accepted, validated_user_ids.len().to_string()); self.send_message(response).await; } /// Handle ping message async fn handle_ping(&self, cv: CommunicationValue) { 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; } } let client_pings = if let Some(rho_conn) = self.get_rho_connection().await { rho_conn.get_client_pings().await } else { HashMap::new() }; let pings = client_pings .into_iter() .map(|(k, v)| (k, JsonValue::String(v.to_string()))) .collect(); let response = CommunicationValue::new(CommunicationType::pong) .with_id(cv.get_id()) .add_data(DataTypes::ping_clients, JsonValue::Object(pings)); self.send_message(response).await; } /// Handle message forwarding to other Iotas async fn handle_forward_message(&self, cv: CommunicationValue) { let receiver_id = cv.get_receiver(); let sender_id = cv.get_sender(); if !self.get_user_ids().await.contains(&sender_id) { self.send_message(CommunicationValue::new( CommunicationType::error_invalid_user_id, )) .await; return; } if let Some(target_rho) = rho_manager::get_rho_con_for_user(receiver_id).await { target_rho.message_to_iota(cv).await; } else { let error = CommunicationValue::new(CommunicationType::error_no_iota) .with_id(cv.get_id()) .with_sender(cv.get_sender()); self.send_message(error).await; } } /// Handle GET_CHATS message async fn handle_get_chats(&self, cv: CommunicationValue) { let receiver_id = cv.get_receiver(); let mut interested_ids: Vec = Vec::new(); let calls: Vec> = call_manager::get_call_groups(receiver_id).await; let mut invites: HashMap> = HashMap::new(); for call in calls { for inviter in call.members.read().await.iter() { let inviter_id = inviter.user_id; if let Some(call_ids) = invites.get_mut(&inviter_id) { call_ids.push(JsonValue::String(call.call_id.to_string())); } else { invites.insert( inviter_id, vec![JsonValue::String(call.call_id.to_string())], ); } } } // Process contacts and add call information let mut enriched_contacts = JsonValue::new_array(); if let Some(contacts_data) = cv.get_data(DataTypes::user_ids) { if let JsonValue::Array(user_ids) = contacts_data { for user_json in user_ids { let user_id_str = user_json["user_id"].as_str().unwrap_or(""); if let Ok(user_id) = Uuid::parse_str(&user_id_str) { interested_ids.push(user_id); let mut enriched_contact = JsonValue::new_object(); let _ = enriched_contact.insert("user_id", user_id.to_string()); let _ = enriched_contact.insert( "calls", JsonValue::Array(invites.get(&user_id).unwrap_or(&vec![]).clone()), ); let _ = enriched_contacts.push(enriched_contact); } } } else { enriched_contacts = contacts_data.clone(); } } // Notify OmegaConnection about user states OmegaConnection::user_states(receiver_id, interested_ids.clone()).await; // Set interested users in RhoConnection if let Some(rho_conn) = self.get_rho_connection().await { rho_conn.set_interested(receiver_id, interested_ids).await; } // Forward to client self.forward_to_client(cv.add_data(DataTypes::user_ids, enriched_contacts)) .await; } /// Forward message to client async fn forward_to_client(&self, cv: CommunicationValue) { if let Some(rho_conn) = self.get_rho_connection().await { let updated_cv = cv.with_sender(self.get_iota_id().await); let _receiver_id = updated_cv.get_receiver(); rho_conn.message_to_client(updated_cv).await; } } async fn send_error_response(&self, message_id: &Uuid) { let error = CommunicationValue::new(CommunicationType::error).with_id(*message_id); self.send_message(error).await; } pub async fn handle_close(&self) { if self.is_identified().await { if let Some(rho_conn) = self.get_rho_connection().await { rho_conn.close_iota_connection().await; } } } } impl std::fmt::Debug for IotaConnection { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("IotaConnection") .field("iota_id", &"[async]") .field("identified", &"[async]") .field("ping", &"[async]") .finish() } }