async-tungstenite

This commit is contained in:
Alex Emmet 2025-11-09 22:01:57 +00:00
commit 1baaad2f76
10 changed files with 423 additions and 272 deletions

32
Cargo.lock generated
View file

@ -7,6 +7,7 @@ name = "Omikron"
version = "0.1.0" version = "0.1.0"
dependencies = [ dependencies = [
"ansi_term", "ansi_term",
"async-tungstenite",
"axum", "axum",
"base64", "base64",
"bytes", "bytes",
@ -34,7 +35,6 @@ dependencies = [
"sysinfo", "sysinfo",
"tokio", "tokio",
"tokio-stream", "tokio-stream",
"tokio-tungstenite 0.28.0",
"tokio-util", "tokio-util",
"tokio_websocket_server", "tokio_websocket_server",
"tracing", "tracing",
@ -89,6 +89,22 @@ dependencies = [
"winapi", "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]] [[package]]
name = "atomic-waker" name = "atomic-waker"
version = "1.1.2" version = "1.1.2"
@ -2224,18 +2240,6 @@ dependencies = [
"tungstenite 0.26.2", "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]] [[package]]
name = "tokio-util" name = "tokio-util"
version = "0.7.16" version = "0.7.16"
@ -2265,7 +2269,7 @@ dependencies = [
"serde", "serde",
"tokio", "tokio",
"tokio-rustls", "tokio-rustls",
"tokio-tungstenite 0.26.2", "tokio-tungstenite",
"tracing", "tracing",
"uuid", "uuid",
] ]

View file

@ -4,42 +4,42 @@ version = "0.1.0"
edition = "2024" edition = "2024"
[dependencies] [dependencies]
tokio = { version = "*", features = ["full"] } ansi_term = "*"
tokio-tungstenite = { version = "*" } async-tungstenite = { version = "*" }
tokio-util = { version = "*", features = ["full"] } axum = "*"
hyper = { version = "*", features = ["full"] } base64 = "0.22.1"
bytes = "*"
chrono = "*"
cmake = "*"
crossterm = "*"
dashmap = "*"
der = "*"
futures = "*"
futures-util = "*" futures-util = "*"
hex = "*"
http = "*" 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" tokio_websocket_server = "0.1.0"
tracing = "*" tracing = "*"
tracing-subscriber = "*" tracing-subscriber = "*"
tungstenite = "*"
uuid = "*" uuid = "*"
axum = "*"
json = "*"
reqwest = "*"
hex = "*"
rustls = { version = "*", default-features = false, features = ["ring"] }
cmake = "*"
walkdir = "2.5.0" walkdir = "2.5.0"
sysinfo = "0.30" warp = "*"
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"
x448 = { version = "*" } x448 = { version = "*" }
x509 = "*" x509 = "*"
sha2 = "*"
der = "*"
pkcs8 = { version = "*", features = ["alloc"] }
futures = "*"
dashmap = "*"
warp = { version = "*" }
chrono = "*"
tokio-stream = "*"
loom = "0.7.2"
ansi_term = "*"

View file

@ -1,184 +1,249 @@
use std::sync::Arc; use async_tungstenite::{
WebSocketReceiver, WebSocketSender, WebSocketStream, tungstenite::Message,
use crate::{
calls::{call_manager::CallManagerState, caller::Caller},
data::communication::{CommunicationType, CommunicationValue, DataTypes},
}; };
use futures::{SinkExt, StreamExt, lock::Mutex}; use futures::SinkExt;
use json::JsonValue; use std::sync::Arc;
use tokio::sync::mpsc::UnboundedSender; use tokio::sync::{
use tokio::sync::mpsc::unbounded_channel; Mutex, RwLock,
mpsc::{UnboundedSender, unbounded_channel},
};
use tokio_util::compat::Compat;
use tungstenite::Utf8Bytes; use tungstenite::Utf8Bytes;
use uuid::Uuid; use uuid::Uuid;
pub async fn handle_connection( use crate::{
raw_stream: tokio::net::TcpStream, calls::call_manager,
state: Arc<Mutex<CallManagerState>>, data::communication::{CommunicationType, CommunicationValue, DataTypes},
) { util::print::{PrintType, line, line_err},
let ws_stream = tokio_tungstenite::accept_async(raw_stream) };
.await use json::JsonValue;
.expect("Error during the websocket handshake");
let (outgoing, mut incoming) = ws_stream.split(); pub struct CallConnection {
pub sender: Arc<RwLock<WebSocketSender<Compat<tokio::net::TcpStream>>>>,
pub receiver: Arc<RwLock<WebSocketReceiver<Compat<tokio::net::TcpStream>>>>,
pub tx: UnboundedSender<Utf8Bytes>,
pub user_id: Arc<RwLock<Option<Uuid>>>,
pub call_id: Arc<RwLock<Option<Uuid>>>,
}
// We create a channel so other parts can send to this client impl CallConnection {
pub type Tx = UnboundedSender<Utf8Bytes>; pub async fn new(
let (tx, mut rx): (UnboundedSender<_>, _) = unbounded_channel(); sender: WebSocketSender<Compat<tokio::net::TcpStream>>,
receiver: WebSocketReceiver<Compat<tokio::net::TcpStream>>,
) -> Arc<Self> {
let (tx, mut rx) = unbounded_channel::<Utf8Bytes>();
// Spawn a task to forward from rx → outgoing let conn = Arc::new(Self {
let mut outgoing_clone = outgoing; 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 { tokio::spawn(async move {
while let Some(msg) = rx.recv().await { while let Some(msg) = rx.recv().await {
let _ = outgoing_clone let mut sess = sender_clone.write().await;
.send(tokio_tungstenite::tungstenite::Message::Text(msg)) let _ = sess.send(Message::Text(msg)).await;
.await;
} }
}); });
// Each connection will track its user_id, call_group, etc. conn
let mut maybe_user_id: Option<Uuid> = None; }
let mut maybe_call_id: Option<Uuid> = None;
/// Handle a single incoming message
pub async fn handle_message(self: Arc<Self>, 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());
}
while let Some(msg) = incoming.next().await {
let msg = match msg {
Ok(m) => m,
Err(_) => break,
};
if let tokio_tungstenite::tungstenite::Message::Text(s) = msg {
// parse CommunicationValue
let cv = CommunicationValue::from_json(&s);
match cv.comm_type { match cv.comm_type {
CommunicationType::identification => { CommunicationType::identification => self.handle_identification(cv).await,
// extract fields CommunicationType::ping => self.handle_ping(cv).await,
let user_str = cv.get_data(DataTypes::user_id).unwrap().as_str().unwrap(); CommunicationType::client_changed => self.handle_client_changed(cv).await,
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()));
}
}
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()));
}
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);
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 => { CommunicationType::start_stream | CommunicationType::end_stream => {
if let Some(uid) = maybe_user_id { self.handle_stream_toggle(cv).await
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_sdp
| CommunicationType::webrtc_ice | CommunicationType::webrtc_ice
| CommunicationType::watch_stream => { | CommunicationType::watch_stream => self.handle_direct_relay(cv).await,
if let Some(uid) = maybe_user_id { _ => {}
if let Some(JsonValue::String(receiver_str)) = }
cv.get_data(DataTypes::receiver_id) }
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 Ok(receiver_id) = Uuid::parse_str(&receiver_str) { *self.user_id.write().await = Some(uid);
if let Some(cid) = maybe_call_id { *self.call_id.write().await = Some(cid);
let st = state.lock().await; }
if let Some(group) = st.call_groups.get(&cid) {
let group = call_manager::get_or_create_group(cid, secret_sha).await;
{
group.lock().await.add_member(uid, self.tx.clone());
}
// 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 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;
}
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(); let mut bc = cv.clone();
bc = bc.add_data( bc = bc.add_data(DataTypes::sender_id, JsonValue::String(uid.to_string()));
DataTypes::sender_id, group.broadcast(&bc.to_json().to_string());
JsonValue::String(uid.to_string()), }
); }
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 group
.lock() .lock()
.await .await
.send_to(&receiver_id, &bc.to_json().to_string()); .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) {
// other types you may handle 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 {
}
}
// 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); group.lock().await.remove_member(uid);
} }
st.remove_inactive().await; call_manager::remove_inactive().await;
} }
}
fn caller_state_change(_caller: &mut Caller, _state_str: &str) { let mut session = self.sender.write().await;
// parse and set your enum, e.g. match _state_str { "active" => ..., etc. } let _ = session.close(None).await;
}
pub async fn handle_close(&self) {}
} }

View file

@ -1,45 +1,57 @@
use crate::calls::call_group::CallGroup; use crate::calls::call_group::CallGroup;
use futures::lock::Mutex; use futures::lock::Mutex;
use once_cell::sync::Lazy;
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::Arc; use std::sync::Arc;
use tokio::sync::RwLock;
use tokio::sync::mpsc::UnboundedSender; use tokio::sync::mpsc::UnboundedSender;
use tungstenite::Utf8Bytes; use tungstenite::Utf8Bytes;
use uuid::Uuid; use uuid::Uuid;
pub type Tx = UnboundedSender<Utf8Bytes>; pub type Tx = UnboundedSender<Utf8Bytes>;
pub struct CallManagerState { pub static CALL_GROUPS: Lazy<RwLock<Mutex<HashMap<Uuid, Arc<Mutex<CallGroup>>>>>> =
pub call_groups: HashMap<Uuid, Arc<Mutex<CallGroup>>>, Lazy::new(|| RwLock::new(Mutex::new(HashMap::new())));
pub async fn get_group(call_id: Uuid) -> Option<Arc<Mutex<CallGroup>>> {
CALL_GROUPS
.read()
.await
.lock()
.await
.get_mut(&call_id)
.cloned()
} }
pub async fn get_or_create_group(call_id: Uuid, secret: &str) -> Arc<Mutex<CallGroup>> {
impl CallManagerState { let g = CALL_GROUPS.write().await;
pub fn new() -> Self { if let Some(group) = g.lock().await.get_mut(&call_id) {
CallManagerState {
call_groups: HashMap::new(),
}
}
pub fn get_or_create_group(&mut self, call_id: Uuid, secret: &str) -> Arc<Mutex<CallGroup>> {
let g = self.call_groups.get_mut(&call_id);
if let Some(group) = g {
group.clone() group.clone()
} else { } else {
let cg = Arc::new(Mutex::new(CallGroup::new(call_id))); let cg = Arc::new(Mutex::new(CallGroup::new(call_id)));
self.call_groups.insert(call_id, cg); g.lock().await.insert(call_id, cg.clone());
self.call_groups.get(&call_id).unwrap().clone() cg
}
} }
}
pub async fn remove_inactive(&mut self) { pub async fn remove_inactive() {
let mut rem = Vec::new(); let mut rem = Vec::new();
for cg in self.call_groups.keys() { for cg in CALL_GROUPS.read().await.lock().await.keys() {
let group = self.call_groups.get(cg).unwrap().lock().await; if CALL_GROUPS
if group.callers.is_empty() { .write()
.await
.lock()
.await
.get(cg)
.unwrap()
.lock()
.await
.callers
.is_empty()
{
rem.push(cg.clone()); rem.push(cg.clone());
} }
} }
for cg in rem { for cg in rem {
self.call_groups.remove(&cg); CALL_GROUPS.write().await.lock().await.remove(&cg);
}
} }
} }

View file

@ -2,7 +2,7 @@ use tokio::sync::mpsc::UnboundedSender;
use tungstenite::Utf8Bytes; use tungstenite::Utf8Bytes;
use uuid::Uuid; use uuid::Uuid;
#[derive(Debug)] #[derive(Debug, Clone)]
pub struct Caller { pub struct Caller {
pub user_id: Uuid, pub user_id: Uuid,
pub tx: UnboundedSender<Utf8Bytes>, pub tx: UnboundedSender<Utf8Bytes>,

View file

@ -5,29 +5,29 @@ mod omega;
mod rho; mod rho;
mod util; mod util;
use ansi_term::Color; use async_tungstenite::accept_hdr_async;
use crossterm::style::PrintStyledContent;
use futures::StreamExt; use futures::StreamExt;
use std::sync::Arc; use std::{sync::Arc, time::Duration};
use tokio::net::TcpListener; use tokio::net::TcpListener;
use tokio_tungstenite::accept_hdr_async; use tokio_util::compat::TokioAsyncReadCompatExt;
use tungstenite::handshake::server::{Request, Response}; use tungstenite::{
Message, Utf8Bytes,
handshake::server::{Request, Response},
};
use crate::{ use crate::{
calls::call_connection::CallConnection,
omega::omega_connection::OmegaConnection, omega::omega_connection::OmegaConnection,
rho::{client_connection::ClientConnection, iota_connection::IotaConnection}, rho::{client_connection::ClientConnection, iota_connection::IotaConnection},
util::print::PrintType, util::print::{PrintType, line, line_err, print_start_message},
util::print::line,
util::print::line_err,
util::print::print_start_message,
}; };
#[tokio::main] #[tokio::main]
async fn main() { async fn main() {
print_start_message(); print_start_message();
tokio::spawn(async move {
OmegaConnection::new().connect().await; OmegaConnection::new().connect().await;
});
let listener = TcpListener::bind("0.0.0.0:959").await.unwrap(); let listener = TcpListener::bind("0.0.0.0:959").await.unwrap();
line( line(
PrintType::OmegaIn, PrintType::OmegaIn,
@ -37,11 +37,12 @@ async fn main() {
while let Ok((stream, _)) = listener.accept().await { while let Ok((stream, _)) = listener.accept().await {
tokio::spawn(async move { tokio::spawn(async move {
let mut path: String = "/".to_string(); let mut path: String = "/".to_string();
let callback = |req: &Request, response: Response| { let callback = |req: &Request, response: Response| {
path = format!("{}", &req.uri().path()); path = req.uri().path().to_string(); // Extract URI path
Ok(response) 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, Ok(ws) => ws,
Err(e) => { Err(e) => {
line_err( line_err(
@ -51,13 +52,14 @@ async fn main() {
return; return;
} }
}; };
let (sender, receiver) = ws_stream.split();
if path == "/ws/client/" { if path == "/ws/client/" {
line(PrintType::ClientIn, "New Client connection"); line(PrintType::ClientIn, "New Client connection");
let client_conn: Arc<ClientConnection> = let client_conn: Arc<ClientConnection> =
Arc::from(ClientConnection::new(ws_stream)); Arc::from(ClientConnection::new(sender, receiver));
loop { loop {
let msg_result = { 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 session_lock.next().await
}; };
@ -86,10 +88,11 @@ async fn main() {
} }
} else if path == "/ws/iota/" { } else if path == "/ws/iota/" {
line(PrintType::IotaIn, "New Iota connection"); line(PrintType::IotaIn, "New Iota connection");
let iota_conn: Arc<IotaConnection> = Arc::from(IotaConnection::new(ws_stream)); let iota_conn: Arc<IotaConnection> =
Arc::from(IotaConnection::new(sender, receiver));
loop { loop {
let msg_result = { 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 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<CallConnection> =
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;
}
}
}
} }
}); });
} }

View file

@ -3,7 +3,6 @@ use std::time::Duration;
use crate::data::communication::{CommunicationType, CommunicationValue}; use crate::data::communication::{CommunicationType, CommunicationValue};
use crate::util::print::PrintType; use crate::util::print::PrintType;
use crate::util::print::line;
use crate::util::print::line_err; use crate::util::print::line_err;
use crate::{ use crate::{
data::{ data::{
@ -13,21 +12,23 @@ use crate::{
rho::rho_manager, rho::rho_manager,
util::config_util::CONFIG, util::config_util::CONFIG,
}; };
use async_tungstenite::tungstenite::protocol::Message;
use dashmap::DashMap; use dashmap::DashMap;
use futures::{SinkExt, StreamExt}; use futures::{SinkExt, StreamExt};
use json::JsonValue; use json::JsonValue;
use once_cell::sync::Lazy; use once_cell::sync::Lazy;
use tokio::sync::Mutex; use tokio::sync::Mutex;
use tokio::time::sleep; use tokio::time::sleep;
use tokio_tungstenite::{connect_async, tungstenite::protocol::Message}; use tokio_util::compat::Compat;
use tungstenite::Utf8Bytes; use tungstenite::{Utf8Bytes, connect};
use uuid::Uuid; use uuid::Uuid;
static WAITING_TASKS: Lazy<DashMap<Uuid, Box<dyn Fn(CommunicationValue) -> bool + Send + Sync>>> = static WAITING_TASKS: Lazy<DashMap<Uuid, Box<dyn Fn(CommunicationValue) -> bool + Send + Sync>>> =
Lazy::new(DashMap::new); Lazy::new(DashMap::new);
pub struct OmegaConnection { pub struct OmegaConnection {
ws_stream: Arc<Mutex<Option<tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>>>>, ws_stream:
Arc<Mutex<Option<async_tungstenite::WebSocketStream<Compat<tokio::net::TcpStream>>>>>,
} }
impl OmegaConnection { impl OmegaConnection {
@ -50,7 +51,7 @@ impl OmegaConnection {
return; return;
} }
match connect_async("wss://tensamin.methanium.net/ws/omega").await { match connect("wss://tensamin.methanium.net/ws/omega") {
Ok((_, _)) => { Ok((_, _)) => {
retry = 0; retry = 0;
let identify_msg = CommunicationValue::new(CommunicationType::identification) let identify_msg = CommunicationValue::new(CommunicationType::identification)
@ -79,7 +80,9 @@ impl OmegaConnection {
} }
async fn read_loop( async fn read_loop(
ws_stream: Arc<Mutex<Option<tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>>>>, ws_stream: Arc<
Mutex<Option<async_tungstenite::WebSocketStream<Compat<tokio::net::TcpStream>>>>,
>,
) { ) {
loop { loop {
let mut lock = ws_stream.lock().await; let mut lock = ws_stream.lock().await;

View file

@ -1,8 +1,10 @@
use async_tungstenite::{WebSocketReceiver, WebSocketSender};
use async_tungstenite::{WebSocketStream, tungstenite::Message};
use futures::SinkExt; use futures::SinkExt;
use std::sync::{Arc, Weak}; use std::sync::{Arc, Weak};
use tokio::sync::Mutex; use tokio::sync::Mutex;
use tokio::sync::RwLock; use tokio::sync::RwLock;
use tokio_tungstenite::{WebSocketStream, tungstenite::Message}; use tokio_util::compat::Compat;
use tungstenite::Utf8Bytes; use tungstenite::Utf8Bytes;
use uuid::Uuid; use uuid::Uuid;
@ -23,7 +25,8 @@ use crate::{
/// ClientConnection represents a WebSocket connection from a client device /// ClientConnection represents a WebSocket connection from a client device
pub struct ClientConnection { pub struct ClientConnection {
/// WebSocket session /// WebSocket session
pub session: Arc<Mutex<WebSocketStream<tokio::net::TcpStream>>>, pub sender: Arc<RwLock<WebSocketSender<Compat<tokio::net::TcpStream>>>>,
pub receiver: Arc<RwLock<WebSocketReceiver<Compat<tokio::net::TcpStream>>>>,
/// User ID associated with this client /// User ID associated with this client
pub user_id: Arc<RwLock<Option<Uuid>>>, pub user_id: Arc<RwLock<Option<Uuid>>>,
/// Whether this connection has been identified/authenticated /// Whether this connection has been identified/authenticated
@ -38,9 +41,13 @@ pub struct ClientConnection {
impl ClientConnection { impl ClientConnection {
/// Create a new ClientConnection /// Create a new ClientConnection
pub fn new(session: WebSocketStream<tokio::net::TcpStream>) -> Arc<Self> { pub fn new(
sender: WebSocketSender<Compat<tokio::net::TcpStream>>,
receiver: WebSocketReceiver<Compat<tokio::net::TcpStream>>,
) -> Arc<Self> {
Arc::new(Self { 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)), user_id: Arc::new(RwLock::new(None)),
identified: Arc::new(RwLock::new(false)), identified: Arc::new(RwLock::new(false)),
ping: Arc::new(RwLock::new(-1)), ping: Arc::new(RwLock::new(-1)),
@ -82,7 +89,7 @@ impl ClientConnection {
/// Send a string message to the client /// Send a string message to the client
pub async fn send_message_str(&self, message: &str) { 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 if let Err(e) = session
.send(Message::Text(Utf8Bytes::from(message.to_string()))) .send(Message::Text(Utf8Bytes::from(message.to_string())))
.await .await
@ -423,7 +430,7 @@ impl ClientConnection {
/// Close the connection /// Close the connection
pub async fn close(&self) { pub async fn close(&self) {
let mut session = self.session.lock().await; let mut session = self.sender.write().await;
let _ = session.close(None).await; let _ = session.close(None).await;
} }
@ -466,7 +473,8 @@ impl ClientConnection {
impl Clone for ClientConnection { impl Clone for ClientConnection {
fn clone(&self) -> Self { fn clone(&self) -> 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), user_id: Arc::clone(&self.user_id),
identified: Arc::clone(&self.identified), identified: Arc::clone(&self.identified),
ping: Arc::clone(&self.ping), ping: Arc::clone(&self.ping),

View file

@ -2,6 +2,10 @@ use crate::util::print::PrintType;
use crate::util::print::line; use crate::util::print::line;
use crate::util::print::line_err; use crate::util::print::line_err;
use ansi_term::Color; 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::FutureExt;
use futures::SinkExt; use futures::SinkExt;
use json::JsonValue; use json::JsonValue;
@ -10,7 +14,7 @@ use std::{
sync::{Arc, Weak}, sync::{Arc, Weak},
}; };
use tokio::sync::{Mutex, RwLock}; use tokio::sync::{Mutex, RwLock};
use tokio_tungstenite::{WebSocketStream, tungstenite::Message}; use tokio_util::compat::Compat;
use tungstenite::Utf8Bytes; use tungstenite::Utf8Bytes;
use uuid::Uuid; use uuid::Uuid;
@ -22,9 +26,9 @@ use crate::{
omega::omega_connection::OmegaConnection, omega::omega_connection::OmegaConnection,
}; };
/// IotaConnection represents a WebSocket connection from an Iota device
pub struct IotaConnection { pub struct IotaConnection {
pub session: Arc<Mutex<WebSocketStream<tokio::net::TcpStream>>>, pub sender: Arc<RwLock<WebSocketSender<Compat<tokio::net::TcpStream>>>>,
pub receiver: Arc<RwLock<WebSocketReceiver<Compat<tokio::net::TcpStream>>>>,
pub iota_id: Arc<RwLock<Uuid>>, pub iota_id: Arc<RwLock<Uuid>>,
pub user_ids: Arc<RwLock<Vec<Uuid>>>, pub user_ids: Arc<RwLock<Vec<Uuid>>>,
pub identified: Arc<RwLock<bool>>, pub identified: Arc<RwLock<bool>>,
@ -34,9 +38,13 @@ pub struct IotaConnection {
impl IotaConnection { impl IotaConnection {
/// Create a new IotaConnection /// Create a new IotaConnection
pub fn new(session: WebSocketStream<tokio::net::TcpStream>) -> Arc<Self> { pub fn new(
sender: WebSocketSender<Compat<tokio::net::TcpStream>>,
receiver: WebSocketReceiver<Compat<tokio::net::TcpStream>>,
) -> Arc<Self> {
Arc::new(Self { 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())), iota_id: Arc::new(RwLock::new(Uuid::nil())),
user_ids: Arc::new(RwLock::new(Vec::new())), user_ids: Arc::new(RwLock::new(Vec::new())),
identified: Arc::new(RwLock::new(false)), identified: Arc::new(RwLock::new(false)),
@ -83,7 +91,7 @@ impl IotaConnection {
/// Send a message to the Iota /// Send a message to the Iota
pub async fn send_message_str(&self, message: &str) { 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 if let Err(e) = session
.send(Message::Text(Utf8Bytes::from(message.to_string()))) .send(Message::Text(Utf8Bytes::from(message.to_string())))
.await .await
@ -256,7 +264,8 @@ impl IotaConnection {
if !self.get_user_ids().await.contains(&sender_id) { if !self.get_user_ids().await.contains(&sender_id) {
self.send_message(CommunicationValue::new( self.send_message(CommunicationValue::new(
CommunicationType::error_invalid_user_id, CommunicationType::error_invalid_user_id,
)); ))
.await;
return; return;
} }

View file

@ -4,7 +4,9 @@ pub fn print_start_message() {
println!("{}", Color::Yellow.paint("> Iota inbound")); println!("{}", Color::Yellow.paint("> Iota inbound"));
println!("{}", Color::Purple.paint("< Iota outbound")); println!("{}", Color::Purple.paint("< Iota outbound"));
println!("{}", Color::Green.paint("> Client inbound")); 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 inbound"));
println!("{}", Color::Cyan.paint("< Omega outbound")); println!("{}", Color::Cyan.paint("< Omega outbound"));
println!("{}", Color::White.paint("General info")); println!("{}", Color::White.paint("General info"));
@ -17,6 +19,8 @@ pub enum PrintType {
OmegaOut, OmegaOut,
ClientIn, ClientIn,
ClientOut, ClientOut,
CallIn,
CallOut,
General, General,
} }
pub fn line(key: PrintType, message: &str) { pub fn line(key: PrintType, message: &str) {
@ -39,6 +43,12 @@ pub fn line(key: PrintType, message: &str) {
PrintType::ClientOut => { PrintType::ClientOut => {
println!("{}{}", Color::Blue.paint("<"), Color::Blue.paint(message)) 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)), PrintType::General => println!("{}", Color::White.paint(message)),
} }
} }
@ -66,6 +76,10 @@ pub fn line_err(key: PrintType, message: &str) {
PrintType::ClientOut => { PrintType::ClientOut => {
println!("{}{}", Color::Blue.paint("<<"), Color::Blue.paint(message)) 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)), PrintType::General => println!("{}", Color::Red.paint(message)),
} }
} }