diff --git a/src/anonymous_clients/anonymous_client_connection.rs b/src/anonymous_clients/anonymous_client_connection.rs index b91776b..c8bf33d 100644 --- a/src/anonymous_clients/anonymous_client_connection.rs +++ b/src/anonymous_clients/anonymous_client_connection.rs @@ -13,7 +13,7 @@ use uuid::Uuid; use crate::anonymous_clients::anonymous_manager::{self, generate_username}; use crate::calls::call_manager; use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes}; -use crate::omega::omega_connection::{WAITING_TASKS, get_omega_connection}; +use crate::omega::omega_connection::get_omega_connection; use crate::rho::rho_manager; use crate::util::logger::PrintType; use crate::{log_in, log_out}; @@ -328,19 +328,17 @@ impl AnonymousClientConnection { } async fn handle_omega_forward(self: Arc, cv: CommunicationValue) { let client_for_closure = self.clone(); - WAITING_TASKS.insert( - cv.get_id(), - Box::new(move |_, response_cv| { - let client = client_for_closure.clone(); - tokio::spawn(async move { - client.send_message(&response_cv).await; - }); - true - }), - ); - get_omega_connection() - .send_message(&cv.with_sender(*self.user_id.read().await)) - .await; + tokio::spawn(async move { + let response_cv = get_omega_connection() + .await_response( + &cv.with_sender(*self.user_id.read().await), + Some(Duration::from_secs(20)), + ) + .await; + if let Ok(response_cv) = response_cv { + client_for_closure.send_message(&response_cv).await; + } + }); } /// Handle ping message diff --git a/src/main.rs b/src/main.rs index e77d70e..f6731f9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -19,7 +19,6 @@ use crate::{ anonymous_client_connection::AnonymousClientConnection, anonymous_manager, }, calls::call_util::garbage_collect_calls, - omega::omega_connection::OmegaConnection, rho::{client_connection::ClientConnection, iota_connection::IotaConnection}, util::{ crypto_helper::{load_public_key, load_secret_key}, @@ -39,9 +38,6 @@ pub fn get_public_key() -> x448::PublicKey { #[tokio::main] async fn main() { dotenv().ok(); - tokio::spawn(async move { - Arc::new(OmegaConnection::new()).connect(); - }); startup(); let address = format!( "{}:{}", diff --git a/src/omega/omega_connection.rs b/src/omega/omega_connection.rs index 9845567..30b6f03 100755 --- a/src/omega/omega_connection.rs +++ b/src/omega/omega_connection.rs @@ -31,9 +31,12 @@ use crate::{ }, }; -pub static WAITING_TASKS: Lazy< - DashMap, CommunicationValue) -> bool + Send + Sync>>, -> = Lazy::new(DashMap::new); +pub struct WaitingTask { + pub task: Box, CommunicationValue) -> bool + Send + Sync>, + pub inserted_at: Instant, +} + +pub static WAITING_TASKS: Lazy> = Lazy::new(DashMap::new); static OMEGA_CONNECTION: Lazy> = Lazy::new(|| { let conn = Arc::new(OmegaConnection::new()); @@ -41,6 +44,14 @@ static OMEGA_CONNECTION: Lazy> = Lazy::new(|| { tokio::spawn(async move { conn_clone.connect_internal(0).await; }); + + tokio::spawn(async { + loop { + sleep(Duration::from_secs(60)).await; + WAITING_TASKS.retain(|_, v| v.inserted_at.elapsed() < Duration::from_secs(60)); + } + }); + conn }); @@ -52,7 +63,7 @@ pub fn get_omega_connection() -> Arc { #[derive(Clone)] pub struct OmegaConnection { write: Arc< - RwLock< + Mutex< Option< WebSocketSender< Stream, TokioAdapter>>, @@ -61,7 +72,7 @@ pub struct OmegaConnection { >, >, read: Arc< - RwLock< + Mutex< Option< WebSocketReceiver< Stream, TokioAdapter>>, @@ -72,27 +83,48 @@ pub struct OmegaConnection { pingpong: Arc>>>, pub last_ping: Arc>, pub message_send_times: Arc>>, - pub is_connected: Arc>, + state: Arc>, +} + +#[derive(Clone, PartialEq, Eq)] +enum ConnectionState { + Disconnected, + Connecting, + Connected, } impl OmegaConnection { pub fn new() -> Self { OmegaConnection { - read: Arc::new(RwLock::new(None)), - write: Arc::new(RwLock::new(None)), + read: Arc::new(Mutex::new(None)), + write: Arc::new(Mutex::new(None)), pingpong: Arc::new(Mutex::new(None)), last_ping: Arc::new(Mutex::new(-1)), message_send_times: Arc::new(Mutex::new(HashMap::new())), - is_connected: Arc::new(RwLock::new(false)), + state: Arc::new(RwLock::new(ConnectionState::Disconnected)), } } - pub fn connect(self: Arc) { - let cloned_self = self.clone(); - tokio::spawn(async move { - cloned_self.connect_internal(0).await; - }); - } async fn connect_internal(self: Arc, mut retry: usize) { + if self.state.read().await.eq(&ConnectionState::Connected) { + return; + } + if self.state.read().await.eq(&ConnectionState::Connecting) { + let start = Instant::now(); + let timeout = Duration::from_secs(10); + while self.state.read().await.eq(&ConnectionState::Connecting) + && start.elapsed() < timeout + { + sleep(Duration::from_millis(100)).await; + } + return; + } + + *self.state.write().await = ConnectionState::Connecting; + + if let Some(handle) = self.pingpong.lock().await.take() { + handle.abort(); + } + loop { if retry > 5 { log_err!( @@ -100,6 +132,7 @@ impl OmegaConnection { PrintType::Omega, "Max retry attempts reached, giving up." ); + *self.state.write().await = ConnectionState::Disconnected; return; } @@ -107,16 +140,17 @@ impl OmegaConnection { env::var("OMEGA_HOST").unwrap_or("wss://omega.tensamin.net/ws/omikron".to_string()); match connect_async(&url_str).await { Ok((ws_stream, _)) => { - *self.is_connected.write().await = true; log_in!(0, PrintType::Omega, "WebSocket connected to {}", url_str); retry = 0; let (write, read) = ws_stream.split(); - *self.read.write().await = Some(read); - *self.write.write().await = Some(write); + *self.read.lock().await = Some(read); + *self.write.lock().await = Some(write); - let cloned_self = self.clone(); - tokio::spawn(async move { - cloned_self.clone().read_loop().await; + *self.state.write().await = ConnectionState::Connected; + + let read_loop_self = self.clone(); + let read_loop_handle = tokio::spawn(async move { + read_loop_self.read_loop().await; }); let cloned_self = self.clone(); @@ -137,148 +171,157 @@ impl OmegaConnection { ); WAITING_TASKS.insert( id, - Box::new(|selfc, cv| { - if cv.is_type(CommunicationType::error_not_found) { - log_err!(0, - PrintType::Omega, - "Identification failed: Omikron ID not found on Omega.", - ); - return false; - } - if !cv.is_type(CommunicationType::challenge) { - return false; - } - tokio::spawn(async move { - let task = async move { - let challenge = cv - .get_data(DataTypes::challenge) - .and_then(|v| v.as_str()) - .ok_or_else(|| { - "Challenge not found or not a string".to_string() - })?; - - let server_pub_key = cv - .get_data(DataTypes::public_key) - .and_then(|v| v.as_str()) - .ok_or_else(|| { - "Public key from server not found or not a string" - .to_string() - })?; - - let decrypted_challenge = decrypt_b64( - &secret_key_to_base64(&get_private_key()), - server_pub_key, - challenge, - ) - .map_err(|e| { - format!("Failed to decrypt challenge: {:?}", e) - })?; - - let response_msg = CommunicationValue::new( - CommunicationType::challenge_response, - ) - .with_id(cv.get_id()) - .add_data( - DataTypes::challenge, - JsonValue::String(decrypted_challenge), + WaitingTask { + task: Box::new(|selfc, cv| { + if cv.is_type(CommunicationType::error_not_found) { + log_err!(0, + PrintType::Omega, + "Identification failed: Omikron ID not found on Omega.", ); - - let response_id = response_msg.get_id(); - WAITING_TASKS.insert( - response_id, - Box::new(|selfc, final_cv| { - if !final_cv - .is_type(CommunicationType::identification_response) - { - log_err!(0, - PrintType::Omega, - "Expected identification_response, got something else.", - ); - return false; - } - - if let Some(accepted) = final_cv.get_data(DataTypes::accepted).and_then(|v| v.as_bool()) { - if !accepted { - log_err!(0, PrintType::Omega, "Omega did not accept identification."); - return false; - } - } else { - log_err!(0, PrintType::Omega, "Omega response did not contain 'accepted' field."); - return false; - } - - tokio::spawn(async move { - let mut connected_iota_ids: Vec = Vec::new(); - let mut connected_user_ids: Vec = Vec::new(); - let rho_connections_reader = RHO_CONNECTIONS.read().await; - - for iota_id in rho_connections_reader.keys() { - connected_iota_ids.push(JsonValue::from(*iota_id)); - } - - for rho in rho_connections_reader.values() { - for client_conn in rho.get_client_connections().await { - connected_user_ids.push(JsonValue::from(client_conn.get_user_id().await)); - } - } - - drop(rho_connections_reader); - - let sync_msg = CommunicationValue::new(CommunicationType::sync_client_iota_status) - .add_data(DataTypes::iota_ids, JsonValue::Array(connected_iota_ids)) - .add_data(DataTypes::user_ids, JsonValue::Array(connected_user_ids)) - .add_data(DataTypes::rho_connections, JsonValue::from(connection_count().await)); - - selfc.send_message(&sync_msg).await; - }); - log!(0, - PrintType::Omega, - "Successfully identified with Omega.", - ); - true - }), - ); - - selfc.send_message(&response_msg).await; - - Ok::<(), String>(()) - }; - - if let Err(e) = task.await { - log_err!(0, PrintType::Omega, "{}", &e); + return false; } - }); + if !cv.is_type(CommunicationType::challenge) { + return false; + } + tokio::spawn(async move { + let task = async move { + let challenge = cv + .get_data(DataTypes::challenge) + .and_then(|v| v.as_str()) + .ok_or_else(|| { + "Challenge not found or not a string".to_string() + })?; - true - }), + let server_pub_key = cv + .get_data(DataTypes::public_key) + .and_then(|v| v.as_str()) + .ok_or_else(|| { + "Public key from server not found or not a string" + .to_string() + })?; + + let decrypted_challenge = decrypt_b64( + &secret_key_to_base64(&get_private_key()), + server_pub_key, + challenge, + ) + .map_err(|e| { + format!("Failed to decrypt challenge: {:?}", e) + })?; + + let response_msg = CommunicationValue::new( + CommunicationType::challenge_response, + ) + .with_id(cv.get_id()) + .add_data( + DataTypes::challenge, + JsonValue::String(decrypted_challenge), + ); + + let response_id = response_msg.get_id(); + WAITING_TASKS.insert( + response_id, + WaitingTask { + task: Box::new(|selfc, final_cv| { + if !final_cv + .is_type(CommunicationType::identification_response) + { + log_err!(0, + PrintType::Omega, + "Expected identification_response, got something else.", + ); + return false; + } + + if let Some(accepted) = final_cv.get_data(DataTypes::accepted).and_then(|v| v.as_bool()) { + if !accepted { + log_err!(0, PrintType::Omega, "Omega did not accept identification."); + return false; + } + } else { + log_err!(0, PrintType::Omega, "Omega response did not contain 'accepted' field."); + return false; + } + + tokio::spawn(async move { + let mut connected_iota_ids: Vec = Vec::new(); + let mut connected_user_ids: Vec = Vec::new(); + let rho_connections_reader = RHO_CONNECTIONS.read().await; + + for iota_id in rho_connections_reader.keys() { + connected_iota_ids.push(JsonValue::from(*iota_id)); + } + + for rho in rho_connections_reader.values() { + for client_conn in rho.get_client_connections().await { + connected_user_ids.push(JsonValue::from(client_conn.get_user_id().await)); + } + } + + drop(rho_connections_reader); + + let sync_msg = CommunicationValue::new(CommunicationType::sync_client_iota_status) + .add_data(DataTypes::iota_ids, JsonValue::Array(connected_iota_ids)) + .add_data(DataTypes::user_ids, JsonValue::Array(connected_user_ids)) + .add_data(DataTypes::rho_connections, JsonValue::from(connection_count().await)); + + selfc.send_message(&sync_msg).await; + }); + log!(0, + PrintType::Omega, + "Successfully identified with Omega.", + ); + true + }), + inserted_at: Instant::now(), + } + ); + + selfc.send_message(&response_msg).await; + + Ok::<(), String>(()) + }; + + if let Err(e) = task.await { + log_err!(0, PrintType::Omega, "{}", &e); + } + }); + + true + }), + inserted_at: Instant::now(), + } ); cloned_self.send_message(&identify_msg).await }); - let cloned_self = self.clone(); + let ping_pong_self = self.clone(); let handle = tokio::spawn(async move { loop { - if *cloned_self.is_connected.read().await == false { + if *ping_pong_self.state.read().await != ConnectionState::Connected { break; } - cloned_self.send_ping().await; + ping_pong_self.send_ping().await; sleep(Duration::from_secs(5)).await; } }); - *self.is_connected.write().await = true; *self.pingpong.lock().await = Some(handle); - while *self.is_connected.read().await { - sleep(Duration::from_secs(2)).await; - } - *self.read.write().await = None; - *self.write.write().await = None; + read_loop_handle.await.unwrap_or_else(|e| { + log_err!(0, PrintType::Omega, "Read loop task failed: {}", e) + }); + + *self.read.lock().await = None; + *self.write.lock().await = None; + *self.state.write().await = ConnectionState::Disconnected; + log_err!(0, PrintType::Omega, "Connection lost. Retrying..."); retry += 1; sleep(Duration::from_secs(2)).await; } Err(e) => { + *self.state.write().await = ConnectionState::Disconnected; log_err!( 0, PrintType::Omega, @@ -295,15 +338,13 @@ impl OmegaConnection { } async fn read_loop(self: Arc) { + let mut reader = match self.read.lock().await.take() { + Some(reader) => reader, + None => return, + }; + loop { - let msg = { - let mut guard = self.read.write().await; - let ws = match guard.as_mut() { - Some(ws) => ws, - _ => break, - }; - ws.next().await - }; + let msg = reader.next().await; match msg { Some(Ok(Message::Text(msg))) => { @@ -314,9 +355,8 @@ impl OmegaConnection { } let msg_id = cv.get_id(); log_in!(0, PrintType::Omega, "{}", &cv.to_json().to_string()); - // Handle waiting tasks if let Some(task) = WAITING_TASKS.remove(&msg_id) { - if (task.1)(self.clone(), cv.clone()) { + if (task.1.task)(self.clone(), cv.clone()) { continue; } } else { @@ -328,21 +368,17 @@ impl OmegaConnection { _ => {} } } - - *self.is_connected.write().await = false; - *self.read.write().await = None; - *self.write.write().await = None; } pub async fn send_message(&self, cv: &CommunicationValue) { - let mut guard = self.write.write().await; + if !cv.is_type(CommunicationType::ping) { + log_out!(0, PrintType::Omega, "{}", &cv.to_json().to_string()); + } + let msg = cv.to_json().to_string(); + + let mut guard = self.write.lock().await; if let Some(ws) = guard.as_mut() { - if !cv.is_type(CommunicationType::ping) { - log_out!(0, PrintType::Omega, "{}", &cv.to_json().to_string()); - } - let _ = ws - .send(Message::Text(cv.to_json().to_string().into())) - .await; + let _ = ws.send(Message::Text(msg.into())).await; } } @@ -378,20 +414,22 @@ impl OmegaConnection { WAITING_TASKS.insert( msg_id, - Box::new( - move |_: Arc, response: CommunicationValue| { - let _ = Box::pin(async move |_: CommunicationValue| { - let rho = rho_manager::get_rho_con_for_user(user_id).await; - if let Some(rho) = rho { - for client in rho.get_client_connections_for_user(user_id).await { - client.send_message(&response).await; + WaitingTask { + task: Box::new( + move |_: Arc, response: CommunicationValue| { + tokio::spawn(async move { + let rho = rho_manager::get_rho_con_for_user(user_id).await; + if let Some(rho) = rho { + for client in rho.get_client_connections_for_user(user_id).await { + client.send_message(&response).await; + } } - } + }); true - }); - true - }, - ), + }, + ), + inserted_at: Instant::now(), + }, ); OmegaConnection::send_global(cv).await; @@ -402,17 +440,15 @@ impl OmegaConnection { } pub async fn await_connection(&self, timeout_duration: Option) -> Result<(), String> { - if *self.is_connected.read().await { + if *self.state.read().await == ConnectionState::Connected { return Ok(()); } let timeout = timeout_duration.unwrap_or(Duration::from_secs(10)); let start = Instant::now(); loop { - { - if *self.is_connected.read().await { - return Ok(()); - } + if *self.state.read().await == ConnectionState::Connected { + return Ok(()); } if start.elapsed() >= timeout { @@ -431,27 +467,29 @@ impl OmegaConnection { cv: &CommunicationValue, timeout_duration: Option, ) -> Result { - let _ = self.await_connection(timeout_duration).await; + self.await_connection(timeout_duration).await?; let (tx, mut rx) = mpsc::channel(1); let msg_id = cv.get_id(); let task_tx = tx.clone(); WAITING_TASKS.insert( msg_id, - Box::new(move |_, response_cv| { - let inner_tx = task_tx.clone(); - tokio::spawn(async move { - if let Err(e) = inner_tx.send(response_cv).await { - log_err!( - 0, - PrintType::Omega, - "Failed to send response back to awaiter: {}", - e - ); - } - }); - true - }), + WaitingTask { + task: Box::new(move |_, response_cv| { + let inner_tx = task_tx.clone(); + tokio::spawn(async move { + if inner_tx.send(response_cv).await.is_err() { + log_err!( + 0, + PrintType::Omega, + "Failed to send response back to awaiter", + ); + } + }); + true + }), + inserted_at: Instant::now(), + }, ); self.send_message(cv).await; diff --git a/src/omega/ping_pong_task.rs b/src/omega/ping_pong_task.rs index 95b63f0..695ef75 100644 --- a/src/omega/ping_pong_task.rs +++ b/src/omega/ping_pong_task.rs @@ -1,15 +1,21 @@ use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes}; use crate::omega::omega_connection::OmegaConnection; use json::number::Number; +use std::time::Duration; use tokio::time::Instant; use uuid::Uuid; +const PING_TIMEOUT: Duration = Duration::from_secs(30); + impl OmegaConnection { pub async fn send_ping(&self) { let uuid = Uuid::new_v4(); let send_time = Instant::now(); - self.message_send_times.lock().await.insert(uuid, send_time); + let mut message_send_times = self.message_send_times.lock().await; + message_send_times.retain(|_uuid, time| time.elapsed() < PING_TIMEOUT); + message_send_times.insert(uuid, send_time); + self.send_ping_message(uuid).await; } @@ -27,16 +33,10 @@ impl OmegaConnection { /// Handles incoming pong and calculates latency pub async fn handle_pong(&self, cv: &CommunicationValue, _log: bool) { let id = cv.get_id(); - let send_time_opt = { - let queue = self.message_send_times.lock().await; - queue.get(&id).cloned() - }; - - if let Some(send_time) = send_time_opt { + let mut message_send_times = self.message_send_times.lock().await; + if let Some(send_time) = message_send_times.remove(&id) { let ping = Instant::now().duration_since(send_time).as_millis() as i64; - self.message_send_times.lock().await.remove(&id); - - *self.last_ping.lock().await = ping as i64; + *self.last_ping.lock().await = ping; } } } diff --git a/src/rho/client_connection.rs b/src/rho/client_connection.rs index 7112e5e..79417e9 100644 --- a/src/rho/client_connection.rs +++ b/src/rho/client_connection.rs @@ -15,7 +15,7 @@ use uuid::Uuid; use super::{rho_connection::RhoConnection, rho_manager}; use crate::anonymous_clients::anonymous_manager; use crate::calls::{call_manager, call_util}; -use crate::omega::omega_connection::{WAITING_TASKS, get_omega_connection}; +use crate::omega::omega_connection::get_omega_connection; use crate::util::crypto_helper::{load_public_key, public_key_to_base64}; use crate::util::crypto_util::{DataFormat, SecurePayload}; use crate::util::logger::PrintType; @@ -374,19 +374,17 @@ impl ClientConnection { } async fn handle_omega_forward(self: Arc, cv: CommunicationValue) { let client_for_closure = self.clone(); - WAITING_TASKS.insert( - cv.get_id(), - Box::new(move |_, response_cv| { - let client = client_for_closure.clone(); - tokio::spawn(async move { - client.send_message(&response_cv).await; - }); - true - }), - ); - get_omega_connection() - .send_message(&cv.with_sender(*self.user_id.read().await)) - .await; + tokio::spawn(async move { + let response_cv = get_omega_connection() + .await_response( + &cv.with_sender(*self.user_id.read().await), + Some(Duration::from_secs(20)), + ) + .await; + if let Ok(response_cv) = response_cv { + client_for_closure.send_message(&response_cv).await; + } + }); } /// Handle ping message diff --git a/src/rho/iota_connection.rs b/src/rho/iota_connection.rs index 36e9fe5..3f753a9 100755 --- a/src/rho/iota_connection.rs +++ b/src/rho/iota_connection.rs @@ -5,7 +5,6 @@ use crate::get_public_key; use crate::log_err; use crate::log_in; use crate::log_out; -use crate::omega::omega_connection::WAITING_TASKS; use crate::omega::omega_connection::get_omega_connection; use crate::util::crypto_helper::load_public_key; use crate::util::crypto_helper::public_key_to_base64; @@ -460,19 +459,17 @@ impl IotaConnection { async fn handle_omega_forward(self: Arc, cv: CommunicationValue) { let iota_for_closure = self.clone(); - WAITING_TASKS.insert( - cv.get_id(), - Box::new(move |_, response_cv| { - let iota = iota_for_closure.clone(); - tokio::spawn(async move { - iota.send_message(&response_cv).await; - }); - true - }), - ); - get_omega_connection() - .send_message(&cv.with_sender(*self.iota_id.read().await)) - .await; + tokio::spawn(async move { + let response_cv = get_omega_connection() + .await_response( + &cv.with_sender(*self.iota_id.read().await), + Some(Duration::from_secs(20)), + ) + .await; + if let Ok(response_cv) = response_cv { + iota_for_closure.send_message(&response_cv).await; + } + }); } /// Handle ping message async fn handle_ping(&self, cv: CommunicationValue) {