diff --git a/Cargo.lock b/Cargo.lock index 3d6c4f1..616ca06 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7,6 +7,7 @@ name = "Omikron" version = "0.1.0" dependencies = [ "ansi_term", + "async-tungstenite", "axum", "base64", "bytes", @@ -34,7 +35,6 @@ dependencies = [ "sysinfo", "tokio", "tokio-stream", - "tokio-tungstenite 0.28.0", "tokio-util", "tokio_websocket_server", "tracing", @@ -89,6 +89,22 @@ dependencies = [ "winapi", ] +[[package]] +name = "async-tungstenite" +version = "0.32.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f89c129ab749940f95509d84950c62092c8b4bc6e386ddb162229037a6ec91" +dependencies = [ + "atomic-waker", + "futures-core", + "futures-io", + "futures-task", + "futures-util", + "log", + "pin-project-lite", + "tungstenite 0.28.0", +] + [[package]] name = "atomic-waker" version = "1.1.2" @@ -2224,18 +2240,6 @@ dependencies = [ "tungstenite 0.26.2", ] -[[package]] -name = "tokio-tungstenite" -version = "0.28.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d25a406cddcc431a75d3d9afc6a7c0f7428d4891dd973e4d54c56b46127bf857" -dependencies = [ - "futures-util", - "log", - "tokio", - "tungstenite 0.28.0", -] - [[package]] name = "tokio-util" version = "0.7.16" @@ -2265,7 +2269,7 @@ dependencies = [ "serde", "tokio", "tokio-rustls", - "tokio-tungstenite 0.26.2", + "tokio-tungstenite", "tracing", "uuid", ] diff --git a/Cargo.toml b/Cargo.toml index 7c97bdb..69a4c56 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,42 +4,42 @@ version = "0.1.0" edition = "2024" [dependencies] -tokio = { version = "*", features = ["full"] } -tokio-tungstenite = { version = "*" } -tokio-util = { version = "*", features = ["full"] } -hyper = { version = "*", features = ["full"] } +ansi_term = "*" +async-tungstenite = { version = "*" } +axum = "*" +base64 = "0.22.1" +bytes = "*" +chrono = "*" +cmake = "*" +crossterm = "*" +dashmap = "*" +der = "*" +futures = "*" futures-util = "*" +hex = "*" http = "*" -tungstenite = "*" +hyper = { version = "*", features = ["full"] } +json = "*" +loom = "0.7.2" +once_cell = "1.21.3" +pkcs8 = { version = "*", features = ["alloc"] } +rand = "0.8" +rand_core = { version = "0.6", features = ["getrandom", "std"] } +reqwest = "*" +rustls = { version = "*", default-features = false, features = ["ring"] } +serde = { version = "1.0.219", features = ["derive"] } +sha2 = "*" +sys-info = "*" +sysinfo = "0.30" +tokio = { version = "*", features = ["full"] } +tokio-stream = "*" +tokio-util = { version = "*", features = ["full"] } tokio_websocket_server = "0.1.0" tracing = "*" tracing-subscriber = "*" +tungstenite = "*" uuid = "*" -axum = "*" -json = "*" -reqwest = "*" -hex = "*" -rustls = { version = "*", default-features = false, features = ["ring"] } -cmake = "*" walkdir = "2.5.0" -sysinfo = "0.30" -serde = { version = "1.0.219", features = ["derive"] } -bytes = "*" -crossterm = "*" -sys-info = "*" -rand = "0.8" -rand_core = { version = "0.6", features = ["getrandom", "std"] } -base64 = "0.22.1" -once_cell = "1.21.3" +warp = "*" x448 = { version = "*" } x509 = "*" -sha2 = "*" -der = "*" -pkcs8 = { version = "*", features = ["alloc"] } -futures = "*" -dashmap = "*" -warp = { version = "*" } -chrono = "*" -tokio-stream = "*" -loom = "0.7.2" -ansi_term = "*" diff --git a/src/calls/call_connection.rs b/src/calls/call_connection.rs index d147da5..83b1b02 100644 --- a/src/calls/call_connection.rs +++ b/src/calls/call_connection.rs @@ -1,184 +1,249 @@ -use std::sync::Arc; - -use crate::{ - calls::{call_manager::CallManagerState, caller::Caller}, - data::communication::{CommunicationType, CommunicationValue, DataTypes}, +use async_tungstenite::{ + WebSocketReceiver, WebSocketSender, WebSocketStream, tungstenite::Message, }; -use futures::{SinkExt, StreamExt, lock::Mutex}; -use json::JsonValue; -use tokio::sync::mpsc::UnboundedSender; -use tokio::sync::mpsc::unbounded_channel; +use futures::SinkExt; +use std::sync::Arc; +use tokio::sync::{ + Mutex, RwLock, + mpsc::{UnboundedSender, unbounded_channel}, +}; +use tokio_util::compat::Compat; use tungstenite::Utf8Bytes; use uuid::Uuid; -pub async fn handle_connection( - raw_stream: tokio::net::TcpStream, - state: Arc>, -) { - let ws_stream = tokio_tungstenite::accept_async(raw_stream) - .await - .expect("Error during the websocket handshake"); +use crate::{ + calls::call_manager, + data::communication::{CommunicationType, CommunicationValue, DataTypes}, + util::print::{PrintType, line, line_err}, +}; +use json::JsonValue; - let (outgoing, mut incoming) = ws_stream.split(); +pub struct CallConnection { + pub sender: Arc>>>, + pub receiver: Arc>>>, + pub tx: UnboundedSender, + pub user_id: Arc>>, + pub call_id: Arc>>, +} - // We create a channel so other parts can send to this client - pub type Tx = UnboundedSender; - let (tx, mut rx): (UnboundedSender<_>, _) = unbounded_channel(); +impl CallConnection { + pub async fn new( + sender: WebSocketSender>, + receiver: WebSocketReceiver>, + ) -> Arc { + let (tx, mut rx) = unbounded_channel::(); - // Spawn a task to forward from rx → outgoing - let mut outgoing_clone = outgoing; - tokio::spawn(async move { - while let Some(msg) = rx.recv().await { - let _ = outgoing_clone - .send(tokio_tungstenite::tungstenite::Message::Text(msg)) - .await; + let conn = Arc::new(Self { + sender: Arc::new(RwLock::new(sender)), + receiver: Arc::new(RwLock::new(receiver)), + tx, + user_id: Arc::new(RwLock::new(None)), + call_id: Arc::new(RwLock::new(None)), + }); + + // Spawn sender task to forward tx → session + let sender_clone = Arc::clone(&conn.sender); + tokio::spawn(async move { + while let Some(msg) = rx.recv().await { + let mut sess = sender_clone.write().await; + let _ = sess.send(Message::Text(msg)).await; + } + }); + + conn + } + + /// Handle a single incoming message + pub async fn handle_message(self: Arc, msg: Utf8Bytes) { + let mut cv = CommunicationValue::from_json(&msg); + if let Some(sender) = *self.user_id.read().await { + cv = cv.with_sender(sender); + } + if !cv.is_type(CommunicationType::identification) && !cv.is_type(CommunicationType::ping) { + line(PrintType::CallIn, &cv.to_json().to_string()); } - }); - // Each connection will track its user_id, call_group, etc. - let mut maybe_user_id: Option = None; - let mut maybe_call_id: Option = None; + match cv.comm_type { + CommunicationType::identification => self.handle_identification(cv).await, + CommunicationType::ping => self.handle_ping(cv).await, + CommunicationType::client_changed => self.handle_client_changed(cv).await, + CommunicationType::start_stream | CommunicationType::end_stream => { + self.handle_stream_toggle(cv).await + } + CommunicationType::webrtc_sdp + | CommunicationType::webrtc_ice + | CommunicationType::watch_stream => self.handle_direct_relay(cv).await, + _ => {} + } + } - while let Some(msg) = incoming.next().await { - let msg = match msg { - Ok(m) => m, - Err(_) => break, + async fn handle_identification(&self, cv: CommunicationValue) { + let user_str = cv + .get_data(DataTypes::user_id) + .and_then(|v| v.as_str()) + .unwrap_or(""); + let call_str = cv + .get_data(DataTypes::call_id) + .and_then(|v| v.as_str()) + .unwrap_or(""); + let secret_sha = cv + .get_data(DataTypes::call_secret_sha) + .and_then(|v| v.as_str()) + .unwrap_or(""); + + let (Ok(uid), Ok(cid)) = (Uuid::parse_str(user_str), Uuid::parse_str(call_str)) else { + return; }; - if let tokio_tungstenite::tungstenite::Message::Text(s) = msg { - // parse CommunicationValue - let cv = CommunicationValue::from_json(&s); - match cv.comm_type { - CommunicationType::identification => { - // extract fields - let user_str = cv.get_data(DataTypes::user_id).unwrap().as_str().unwrap(); - let call_str = cv.get_data(DataTypes::call_id).unwrap().as_str().unwrap(); - let _secret_sha = cv - .get_data(DataTypes::call_secret_sha) - .unwrap() - .as_str() - .unwrap(); - let user_id = Uuid::parse_str(user_str).ok(); - let call_id = Uuid::parse_str(call_str).ok(); - if let (Some(uid), Some(cid)) = (user_id, call_id) { - maybe_user_id = Some(uid); - maybe_call_id = Some(cid); - // Add to call group - let group = state.lock().await.get_or_create_group(cid, &_secret_sha); - group.lock().await.add_member(uid, tx.clone()); - // Response: identification_response plus states - let mut response = - CommunicationValue::new(CommunicationType::identification_response) - .with_id(cv.get_id().clone()); - // build user states - let mut users = JsonValue::new_object(); - for c in group.lock().await.callers.values() { - let mut user_info = JsonValue::new_object(); - let _ = user_info.insert("state", JsonValue::from("muted".to_string())); - let _ = user_info.insert("streaming", JsonValue::from(false)); - let _ = - users.insert(&c.user_id.to_string(), JsonValue::from(user_info)); - } - response = response.add_data(DataTypes::about, JsonValue::from(users)); - // send back - let _ = tx.send(Utf8Bytes::from(response.to_json().to_string())); - } - } + { + *self.user_id.write().await = Some(uid); + *self.call_id.write().await = Some(cid); + } - CommunicationType::ping => { - // optionally parse LAST_PING - // reply with PONG with same message_id - let resp = CommunicationValue::new(CommunicationType::pong) - .with_id(cv.get_id().clone()); - let _ = tx.send(Utf8Bytes::from(resp.to_json().to_string())); - } + let group = call_manager::get_or_create_group(cid, secret_sha).await; + { + group.lock().await.add_member(uid, self.tx.clone()); + } - CommunicationType::client_changed => { - match (maybe_user_id, cv.get_data(DataTypes::call_state)) { - (Some(uid), Some(JsonValue::String(state_str))) => { - if let Some(cid) = maybe_call_id { - { - let mut st = state.lock().await; - let mut group = - st.call_groups.get_mut(&cid).unwrap().lock().await; - let caller = group.caller_state_mut(&uid).unwrap(); - caller_state_change(caller, &state_str); + // Build broadcast + let broadcast = CommunicationValue::new(CommunicationType::client_connected) + .with_id(cv.get_id().clone()) + .add_data_str(DataTypes::user_id, uid.to_string()) + .add_data_str(DataTypes::call_state, "muted".to_string()); + { + group + .lock() + .await + .broadcast(&broadcast.to_json().to_string()); + } + // Build response + let mut response = CommunicationValue::new(CommunicationType::identification_response) + .with_id(cv.get_id().clone()); - let mut bc = cv.clone(); - bc = bc.add_data( - DataTypes::sender_id, - JsonValue::String(uid.to_string()), - ); - group.broadcast(&bc.to_json().to_string()); - } - } - } - _ => (), - } - } - - CommunicationType::start_stream | CommunicationType::end_stream => { - if let Some(uid) = maybe_user_id { - if let Some(cid) = maybe_call_id { - let mut st = state.lock().await; - let mut group = st.call_groups.get_mut(&cid).unwrap().lock().await; - if let Some(caller) = group.caller_state_mut(&uid) { - let streaming = cv.comm_type == CommunicationType::start_stream; - // set streaming - // caller.streaming = streaming; // if you store streaming - } - let mut bc = cv.clone(); - bc = bc - .add_data(DataTypes::sender_id, JsonValue::String(uid.to_string())); - group.broadcast(&bc.to_json().to_string()); - } - } - } - - CommunicationType::webrtc_sdp - | CommunicationType::webrtc_ice - | CommunicationType::watch_stream => { - if let Some(uid) = maybe_user_id { - if let Some(JsonValue::String(receiver_str)) = - cv.get_data(DataTypes::receiver_id) - { - if let Ok(receiver_id) = Uuid::parse_str(&receiver_str) { - if let Some(cid) = maybe_call_id { - let st = state.lock().await; - if let Some(group) = st.call_groups.get(&cid) { - let mut bc = cv.clone(); - bc = bc.add_data( - DataTypes::sender_id, - JsonValue::String(uid.to_string()), - ); - group - .lock() - .await - .send_to(&receiver_id, &bc.to_json().to_string()); - } - } - } - } - } - } - - _ => { - // other types you may handle - } + let mut users = JsonValue::new_object(); + { + for caller in group.lock().await.callers.values() { + let mut user_info = JsonValue::new_object(); + let _ = user_info.insert("state", JsonValue::from("muted")); + let _ = user_info.insert("streaming", JsonValue::from(false)); + let _ = users.insert(&caller.user_id.to_string(), user_info); } } + + response = response.add_data(DataTypes::about, users); + self.send_message(&response).await; } - // on close / disconnection: - if let (Some(uid), Some(cid)) = (maybe_user_id, maybe_call_id) { - let mut st = state.lock().await; - if let Some(group) = st.call_groups.get_mut(&cid) { - group.lock().await.remove_member(uid); + async fn handle_ping(&self, cv: CommunicationValue) { + let resp = CommunicationValue::new(CommunicationType::pong).with_id(cv.get_id().clone()); + self.send_message(&resp).await; + } + + async fn handle_client_changed(&self, cv: CommunicationValue) { + let uid = match *self.user_id.read().await { + Some(id) => id, + None => return, + }; + + let cid = match *self.call_id.read().await { + Some(id) => id, + None => return, + }; + + if let Some(JsonValue::String(_state_str)) = cv.get_data(DataTypes::call_state) { + let Some(group_lock) = call_manager::get_group(cid).await else { + return; + }; + let mut group = group_lock.lock().await; + if let Some(_caller) = group.caller_state_mut(&uid) { + // caller_state_change(caller, &state_str); + } + + let mut bc = cv.clone(); + bc = bc.add_data(DataTypes::sender_id, JsonValue::String(uid.to_string())); + group.broadcast(&bc.to_json().to_string()); } - st.remove_inactive().await; } -} -fn caller_state_change(_caller: &mut Caller, _state_str: &str) { - // parse and set your enum, e.g. match _state_str { "active" => ..., etc. } + async fn handle_stream_toggle(&self, cv: CommunicationValue) { + let uid = match *self.user_id.read().await { + Some(id) => id, + None => return, + }; + + let cid = match *self.call_id.read().await { + Some(id) => id, + None => return, + }; + + if let Some(group_lock) = call_manager::get_group(cid).await { + let group = group_lock.lock().await; + let _streaming = cv.comm_type == CommunicationType::start_stream; + let mut bc = cv.clone(); + bc = bc.add_data(DataTypes::sender_id, JsonValue::String(uid.to_string())); + group.broadcast(&bc.to_json().to_string()); + } + } + + async fn handle_direct_relay(&self, cv: CommunicationValue) { + let uid = match *self.user_id.read().await { + Some(id) => id, + None => return, + }; + let cid = match *self.call_id.read().await { + Some(id) => id, + None => return, + }; + if let Some(JsonValue::String(receiver_str)) = cv.get_data(DataTypes::receiver_id) { + if let Ok(receiver_id) = Uuid::parse_str(&receiver_str) { + if let Some(group) = call_manager::get_group(cid).await { + let mut bc = cv.clone(); + bc = bc.add_data(DataTypes::sender_id, JsonValue::String(uid.to_string())); + group + .lock() + .await + .send_to(&receiver_id, &bc.to_json().to_string()); + line( + PrintType::CallOut, + &format!("Forwarded WebRTC message to receiver: {}", receiver_str), + ); + } else { + line_err(PrintType::CallOut, "Failed to find the group for the call."); + } + } else { + line_err(PrintType::CallOut, "Invalid receiver_id in WebRTC message."); + } + } else { + line_err(PrintType::CallOut, "Missing receiver_id in WebRTC message."); + } + } + pub async fn send_message(&self, cv: &CommunicationValue) { + if !cv.is_type(CommunicationType::pong) { + line(PrintType::CallOut, &cv.to_json().to_string()); + } + let mut session = self.sender.write().await; + if let Err(e) = session + .send(Message::Text(Utf8Bytes::from(cv.to_json().to_string()))) + .await + { + line_err( + PrintType::CallOut, + &format!("Failed to send message to client: {}", e), + ); + } + } + + pub async fn close(&self) { + let (uid, cid) = { (*self.user_id.read().await, *self.call_id.read().await) }; + if let (Some(uid), Some(cid)) = (uid, cid) { + if let Some(group) = call_manager::get_group(cid).await { + group.lock().await.remove_member(uid); + } + call_manager::remove_inactive().await; + } + + let mut session = self.sender.write().await; + let _ = session.close(None).await; + } + pub async fn handle_close(&self) {} } diff --git a/src/calls/call_manager.rs b/src/calls/call_manager.rs index ba96fa4..3b6eb5d 100644 --- a/src/calls/call_manager.rs +++ b/src/calls/call_manager.rs @@ -1,45 +1,57 @@ use crate::calls::call_group::CallGroup; use futures::lock::Mutex; +use once_cell::sync::Lazy; use std::collections::HashMap; use std::sync::Arc; +use tokio::sync::RwLock; use tokio::sync::mpsc::UnboundedSender; use tungstenite::Utf8Bytes; use uuid::Uuid; pub type Tx = UnboundedSender; -pub struct CallManagerState { - pub call_groups: HashMap>>, +pub static CALL_GROUPS: Lazy>>>>> = + Lazy::new(|| RwLock::new(Mutex::new(HashMap::new()))); + +pub async fn get_group(call_id: Uuid) -> Option>> { + CALL_GROUPS + .read() + .await + .lock() + .await + .get_mut(&call_id) + .cloned() +} +pub async fn get_or_create_group(call_id: Uuid, secret: &str) -> Arc> { + let g = CALL_GROUPS.write().await; + if let Some(group) = g.lock().await.get_mut(&call_id) { + group.clone() + } else { + let cg = Arc::new(Mutex::new(CallGroup::new(call_id))); + g.lock().await.insert(call_id, cg.clone()); + cg + } } -impl CallManagerState { - pub fn new() -> Self { - CallManagerState { - call_groups: HashMap::new(), +pub async fn remove_inactive() { + let mut rem = Vec::new(); + for cg in CALL_GROUPS.read().await.lock().await.keys() { + if CALL_GROUPS + .write() + .await + .lock() + .await + .get(cg) + .unwrap() + .lock() + .await + .callers + .is_empty() + { + rem.push(cg.clone()); } } - - pub fn get_or_create_group(&mut self, call_id: Uuid, secret: &str) -> Arc> { - let g = self.call_groups.get_mut(&call_id); - if let Some(group) = g { - group.clone() - } else { - let cg = Arc::new(Mutex::new(CallGroup::new(call_id))); - self.call_groups.insert(call_id, cg); - self.call_groups.get(&call_id).unwrap().clone() - } - } - - pub async fn remove_inactive(&mut self) { - let mut rem = Vec::new(); - for cg in self.call_groups.keys() { - let group = self.call_groups.get(cg).unwrap().lock().await; - if group.callers.is_empty() { - rem.push(cg.clone()); - } - } - for cg in rem { - self.call_groups.remove(&cg); - } + for cg in rem { + CALL_GROUPS.write().await.lock().await.remove(&cg); } } diff --git a/src/calls/caller.rs b/src/calls/caller.rs index 452ebf9..83a43d5 100644 --- a/src/calls/caller.rs +++ b/src/calls/caller.rs @@ -2,7 +2,7 @@ use tokio::sync::mpsc::UnboundedSender; use tungstenite::Utf8Bytes; use uuid::Uuid; -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct Caller { pub user_id: Uuid, pub tx: UnboundedSender, diff --git a/src/main.rs b/src/main.rs index 0e03ce6..eab4acd 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,29 +5,29 @@ mod omega; mod rho; mod util; -use ansi_term::Color; -use crossterm::style::PrintStyledContent; +use async_tungstenite::accept_hdr_async; use futures::StreamExt; -use std::sync::Arc; +use std::{sync::Arc, time::Duration}; use tokio::net::TcpListener; -use tokio_tungstenite::accept_hdr_async; -use tungstenite::handshake::server::{Request, Response}; +use tokio_util::compat::TokioAsyncReadCompatExt; +use tungstenite::{ + Message, Utf8Bytes, + handshake::server::{Request, Response}, +}; use crate::{ + calls::call_connection::CallConnection, omega::omega_connection::OmegaConnection, rho::{client_connection::ClientConnection, iota_connection::IotaConnection}, - util::print::PrintType, - util::print::line, - util::print::line_err, - util::print::print_start_message, + util::print::{PrintType, line, line_err, print_start_message}, }; #[tokio::main] async fn main() { print_start_message(); - - OmegaConnection::new().connect().await; - + tokio::spawn(async move { + OmegaConnection::new().connect().await; + }); let listener = TcpListener::bind("0.0.0.0:959").await.unwrap(); line( PrintType::OmegaIn, @@ -37,11 +37,12 @@ async fn main() { while let Ok((stream, _)) = listener.accept().await { tokio::spawn(async move { let mut path: String = "/".to_string(); + let callback = |req: &Request, response: Response| { - path = format!("{}", &req.uri().path()); + path = req.uri().path().to_string(); // Extract URI path Ok(response) }; - let ws_stream = match accept_hdr_async(stream, callback).await { + let ws_stream = match accept_hdr_async(stream.compat(), callback).await { Ok(ws) => ws, Err(e) => { line_err( @@ -51,13 +52,14 @@ async fn main() { return; } }; + let (sender, receiver) = ws_stream.split(); if path == "/ws/client/" { line(PrintType::ClientIn, "New Client connection"); let client_conn: Arc = - Arc::from(ClientConnection::new(ws_stream)); + Arc::from(ClientConnection::new(sender, receiver)); loop { let msg_result = { - let mut session_lock = client_conn.session.lock().await; + let mut session_lock = client_conn.receiver.write().await; session_lock.next().await }; @@ -86,10 +88,11 @@ async fn main() { } } else if path == "/ws/iota/" { line(PrintType::IotaIn, "New Iota connection"); - let iota_conn: Arc = Arc::from(IotaConnection::new(ws_stream)); + let iota_conn: Arc = + Arc::from(IotaConnection::new(sender, receiver)); loop { let msg_result = { - let mut session_lock = iota_conn.session.lock().await; + let mut session_lock = iota_conn.receiver.write().await; session_lock.next().await }; @@ -117,6 +120,39 @@ async fn main() { } } } + } else if path == "/ws/call/" { + line(PrintType::CallIn, "New Call connection"); + let call_conn: Arc = + Arc::from(CallConnection::new(sender, receiver).await); + loop { + let msg_result = { + let mut session_lock = call_conn.receiver.write().await; + session_lock.next().await + }; + + match msg_result { + Some(Ok(msg)) => { + if msg.is_text() { + let text = msg.into_text().unwrap(); + call_conn.clone().handle_message(text).await; + } else if msg.is_close() { + line(PrintType::CallIn, "Call disconnected"); + call_conn.handle_close().await; + return; + } + } + Some(Err(e)) => { + line_err(PrintType::CallIn, &format!("WebSocket error: {}", e)); + call_conn.handle_close().await; + return; + } + None => { + line(PrintType::CallIn, "Call stream ended"); + call_conn.handle_close().await; + return; + } + } + } } }); } diff --git a/src/omega/omega_connection.rs b/src/omega/omega_connection.rs index 923b865..9d4617f 100644 --- a/src/omega/omega_connection.rs +++ b/src/omega/omega_connection.rs @@ -3,7 +3,6 @@ use std::time::Duration; use crate::data::communication::{CommunicationType, CommunicationValue}; use crate::util::print::PrintType; -use crate::util::print::line; use crate::util::print::line_err; use crate::{ data::{ @@ -13,21 +12,23 @@ use crate::{ rho::rho_manager, util::config_util::CONFIG, }; +use async_tungstenite::tungstenite::protocol::Message; use dashmap::DashMap; use futures::{SinkExt, StreamExt}; use json::JsonValue; use once_cell::sync::Lazy; use tokio::sync::Mutex; use tokio::time::sleep; -use tokio_tungstenite::{connect_async, tungstenite::protocol::Message}; -use tungstenite::Utf8Bytes; +use tokio_util::compat::Compat; +use tungstenite::{Utf8Bytes, connect}; use uuid::Uuid; static WAITING_TASKS: Lazy bool + Send + Sync>>> = Lazy::new(DashMap::new); pub struct OmegaConnection { - ws_stream: Arc>>>, + ws_stream: + Arc>>>>, } impl OmegaConnection { @@ -50,7 +51,7 @@ impl OmegaConnection { return; } - match connect_async("wss://tensamin.methanium.net/ws/omega").await { + match connect("wss://tensamin.methanium.net/ws/omega") { Ok((_, _)) => { retry = 0; let identify_msg = CommunicationValue::new(CommunicationType::identification) @@ -79,7 +80,9 @@ impl OmegaConnection { } async fn read_loop( - ws_stream: Arc>>>, + ws_stream: Arc< + Mutex>>>, + >, ) { loop { let mut lock = ws_stream.lock().await; diff --git a/src/rho/client_connection.rs b/src/rho/client_connection.rs index ce39182..2b1e07e 100644 --- a/src/rho/client_connection.rs +++ b/src/rho/client_connection.rs @@ -1,8 +1,10 @@ +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_tungstenite::{WebSocketStream, tungstenite::Message}; +use tokio_util::compat::Compat; use tungstenite::Utf8Bytes; use uuid::Uuid; @@ -23,7 +25,8 @@ use crate::{ /// ClientConnection represents a WebSocket connection from a client device pub struct ClientConnection { /// WebSocket session - pub session: Arc>>, + pub sender: Arc>>>, + pub receiver: Arc>>>, /// User ID associated with this client pub user_id: Arc>>, /// Whether this connection has been identified/authenticated @@ -38,9 +41,13 @@ pub struct ClientConnection { impl ClientConnection { /// Create a new ClientConnection - pub fn new(session: WebSocketStream) -> Arc { + pub fn new( + sender: WebSocketSender>, + receiver: WebSocketReceiver>, + ) -> Arc { Arc::new(Self { - session: Arc::new(Mutex::new(session)), + 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)), @@ -82,7 +89,7 @@ impl ClientConnection { /// Send a string message to the client pub async fn send_message_str(&self, message: &str) { - let mut session = self.session.lock().await; + let mut session = self.sender.write().await; if let Err(e) = session .send(Message::Text(Utf8Bytes::from(message.to_string()))) .await @@ -423,7 +430,7 @@ impl ClientConnection { /// Close the connection pub async fn close(&self) { - let mut session = self.session.lock().await; + let mut session = self.sender.write().await; let _ = session.close(None).await; } @@ -466,7 +473,8 @@ impl ClientConnection { impl Clone for ClientConnection { fn clone(&self) -> Self { Self { - session: Arc::clone(&self.session), + 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), diff --git a/src/rho/iota_connection.rs b/src/rho/iota_connection.rs index abf981b..e7189a8 100644 --- a/src/rho/iota_connection.rs +++ b/src/rho/iota_connection.rs @@ -2,6 +2,10 @@ use crate::util::print::PrintType; use crate::util::print::line; use crate::util::print::line_err; use ansi_term::Color; +use async_tungstenite::WebSocketReceiver; +use async_tungstenite::WebSocketSender; +use async_tungstenite::tungstenite::protocol::WebSocketConfig; +use async_tungstenite::{WebSocketStream, tungstenite::Message}; use futures::FutureExt; use futures::SinkExt; use json::JsonValue; @@ -10,7 +14,7 @@ use std::{ sync::{Arc, Weak}, }; use tokio::sync::{Mutex, RwLock}; -use tokio_tungstenite::{WebSocketStream, tungstenite::Message}; +use tokio_util::compat::Compat; use tungstenite::Utf8Bytes; use uuid::Uuid; @@ -22,9 +26,9 @@ use crate::{ omega::omega_connection::OmegaConnection, }; -/// IotaConnection represents a WebSocket connection from an Iota device pub struct IotaConnection { - pub session: Arc>>, + pub sender: Arc>>>, + pub receiver: Arc>>>, pub iota_id: Arc>, pub user_ids: Arc>>, pub identified: Arc>, @@ -34,9 +38,13 @@ pub struct IotaConnection { impl IotaConnection { /// Create a new IotaConnection - pub fn new(session: WebSocketStream) -> Arc { + pub fn new( + sender: WebSocketSender>, + receiver: WebSocketReceiver>, + ) -> Arc { Arc::new(Self { - session: Arc::new(Mutex::new(session)), + 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)), @@ -83,7 +91,7 @@ impl IotaConnection { /// Send a message to the Iota pub async fn send_message_str(&self, message: &str) { - let mut session = self.session.lock().await; + let mut session = self.sender.write().await; if let Err(e) = session .send(Message::Text(Utf8Bytes::from(message.to_string()))) .await @@ -256,7 +264,8 @@ impl IotaConnection { if !self.get_user_ids().await.contains(&sender_id) { self.send_message(CommunicationValue::new( CommunicationType::error_invalid_user_id, - )); + )) + .await; return; } diff --git a/src/util/print.rs b/src/util/print.rs index 4ca77c1..72e8630 100644 --- a/src/util/print.rs +++ b/src/util/print.rs @@ -4,7 +4,9 @@ pub fn print_start_message() { println!("{}", Color::Yellow.paint("> Iota inbound")); println!("{}", Color::Purple.paint("< Iota outbound")); println!("{}", Color::Green.paint("> Client inbound")); - println!("{}", Color::Blue.paint("> Client outbound")); + println!("{}", Color::Blue.paint("< Client outbound")); + println!("{}", Color::Red.paint("> Call inbound")); + println!("{}", Color::Red.paint("< Call outbound")); println!("{}", Color::Cyan.paint("> Omega inbound")); println!("{}", Color::Cyan.paint("< Omega outbound")); println!("{}", Color::White.paint("General info")); @@ -17,6 +19,8 @@ pub enum PrintType { OmegaOut, ClientIn, ClientOut, + CallIn, + CallOut, General, } pub fn line(key: PrintType, message: &str) { @@ -39,6 +43,12 @@ pub fn line(key: PrintType, message: &str) { PrintType::ClientOut => { println!("{}{}", Color::Blue.paint("<"), Color::Blue.paint(message)) } + PrintType::CallIn => { + println!("{}{}", Color::Red.paint(">"), Color::Red.paint(message)) + } + PrintType::CallOut => { + println!("{}{}", Color::Red.paint("<"), Color::Red.paint(message)) + } PrintType::General => println!("{}", Color::White.paint(message)), } } @@ -66,6 +76,10 @@ pub fn line_err(key: PrintType, message: &str) { PrintType::ClientOut => { println!("{}{}", Color::Blue.paint("<<"), Color::Blue.paint(message)) } + PrintType::CallIn => println!("{}{}", Color::Red.paint(">>"), Color::Red.paint(message)), + PrintType::CallOut => { + println!("{}{}", Color::Red.paint("<<"), Color::Red.paint(message)) + } PrintType::General => println!("{}", Color::Red.paint(message)), } }