tensamin/omikron
Watch
6
0
Fork
You've already forked omikron
 0
Code Issues Pull requests 3 Projects Releases Packages Wiki Activity Actions

[Fix] More stable Connection

Browse source
This commit is contained in:
Alex-Emmet 2026-02-08 12:05:08 +01:00
commit 7f7b8a5cdd
6 changed files with 271 additions and 244 deletions
Download patch file Download diff file Expand all files Collapse all files

26
src/anonymous_clients/anonymous_client_connection.rs
View file

@ -13,7 +13,7 @@ use uuid::Uuid;
use crate::anonymous_clients::anonymous_manager::{self, generate_username}; use crate::anonymous_clients::anonymous_manager::{self, generate_username};
use crate::calls::call_manager; use crate::calls::call_manager;
use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes}; 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::rho::rho_manager;
use crate::util::logger::PrintType; use crate::util::logger::PrintType;
use crate::{log_in, log_out}; use crate::{log_in, log_out};
@ -328,19 +328,17 @@ impl AnonymousClientConnection {
} }
async fn handle_omega_forward(self: Arc<Self>, cv: CommunicationValue) { async fn handle_omega_forward(self: Arc<Self>, cv: CommunicationValue) {
let client_for_closure = self.clone(); let client_for_closure = self.clone();
WAITING_TASKS.insert( tokio::spawn(async move {
cv.get_id(), let response_cv = get_omega_connection()
Box::new(move |_, response_cv| { .await_response(
let client = client_for_closure.clone(); &cv.with_sender(*self.user_id.read().await),
tokio::spawn(async move { Some(Duration::from_secs(20)),
client.send_message(&response_cv).await; )
}); .await;
true if let Ok(response_cv) = response_cv {
}), client_for_closure.send_message(&response_cv).await;
); }
get_omega_connection() });
.send_message(&cv.with_sender(*self.user_id.read().await))
.await;
} }
/// Handle ping message /// Handle ping message

4
src/main.rs
View file

@ -19,7 +19,6 @@ use crate::{
anonymous_client_connection::AnonymousClientConnection, anonymous_manager, anonymous_client_connection::AnonymousClientConnection, anonymous_manager,
}, },
calls::call_util::garbage_collect_calls, calls::call_util::garbage_collect_calls,
omega::omega_connection::OmegaConnection,
rho::{client_connection::ClientConnection, iota_connection::IotaConnection}, rho::{client_connection::ClientConnection, iota_connection::IotaConnection},
util::{ util::{
crypto_helper::{load_public_key, load_secret_key}, crypto_helper::{load_public_key, load_secret_key},
@ -39,9 +38,6 @@ pub fn get_public_key() -> x448::PublicKey {
#[tokio::main] #[tokio::main]
async fn main() { async fn main() {
dotenv().ok(); dotenv().ok();
tokio::spawn(async move {
Arc::new(OmegaConnection::new()).connect();
});
startup(); startup();
let address = format!( let address = format!(
"{}:{}", "{}:{}",

426
src/omega/omega_connection.rs
View file

@ -31,9 +31,12 @@ use crate::{
}, },
}; };
pub static WAITING_TASKS: Lazy< pub struct WaitingTask {
DashMap<Uuid, Box<dyn Fn(Arc<OmegaConnection>, CommunicationValue) -> bool + Send + Sync>>, pub task: Box<dyn Fn(Arc<OmegaConnection>, CommunicationValue) -> bool + Send + Sync>,
> = Lazy::new(DashMap::new); pub inserted_at: Instant,
}
pub static WAITING_TASKS: Lazy<DashMap<Uuid, WaitingTask>> = Lazy::new(DashMap::new);
static OMEGA_CONNECTION: Lazy<Arc<OmegaConnection>> = Lazy::new(|| { static OMEGA_CONNECTION: Lazy<Arc<OmegaConnection>> = Lazy::new(|| {
let conn = Arc::new(OmegaConnection::new()); let conn = Arc::new(OmegaConnection::new());
@ -41,6 +44,14 @@ static OMEGA_CONNECTION: Lazy<Arc<OmegaConnection>> = Lazy::new(|| {
tokio::spawn(async move { tokio::spawn(async move {
conn_clone.connect_internal(0).await; 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 conn
}); });
@ -52,7 +63,7 @@ pub fn get_omega_connection() -> Arc<OmegaConnection> {
#[derive(Clone)] #[derive(Clone)]
pub struct OmegaConnection { pub struct OmegaConnection {
write: Arc< write: Arc<
RwLock< Mutex<
Option< Option<
WebSocketSender< WebSocketSender<
Stream<TokioAdapter<TcpStream>, TokioAdapter<TlsStream<TcpStream>>>, Stream<TokioAdapter<TcpStream>, TokioAdapter<TlsStream<TcpStream>>>,
@ -61,7 +72,7 @@ pub struct OmegaConnection {
>, >,
>, >,
read: Arc< read: Arc<
RwLock< Mutex<
Option< Option<
WebSocketReceiver< WebSocketReceiver<
Stream<TokioAdapter<TcpStream>, TokioAdapter<TlsStream<TcpStream>>>, Stream<TokioAdapter<TcpStream>, TokioAdapter<TlsStream<TcpStream>>>,
@ -72,27 +83,48 @@ pub struct OmegaConnection {
pingpong: Arc<Mutex<Option<tokio::task::JoinHandle<()>>>>, pingpong: Arc<Mutex<Option<tokio::task::JoinHandle<()>>>>,
pub last_ping: Arc<Mutex<i64>>, pub last_ping: Arc<Mutex<i64>>,
pub message_send_times: Arc<Mutex<HashMap<Uuid, Instant>>>, pub message_send_times: Arc<Mutex<HashMap<Uuid, Instant>>>,
pub is_connected: Arc<RwLock<bool>>, state: Arc<RwLock<ConnectionState>>,
}
#[derive(Clone, PartialEq, Eq)]
enum ConnectionState {
Disconnected,
Connecting,
Connected,
} }
impl OmegaConnection { impl OmegaConnection {
pub fn new() -> Self { pub fn new() -> Self {
OmegaConnection { OmegaConnection {
read: Arc::new(RwLock::new(None)), read: Arc::new(Mutex::new(None)),
write: Arc::new(RwLock::new(None)), write: Arc::new(Mutex::new(None)),
pingpong: Arc::new(Mutex::new(None)), pingpong: Arc::new(Mutex::new(None)),
last_ping: Arc::new(Mutex::new(-1)), last_ping: Arc::new(Mutex::new(-1)),
message_send_times: Arc::new(Mutex::new(HashMap::new())), 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<OmegaConnection>) {
let cloned_self = self.clone();
tokio::spawn(async move {
cloned_self.connect_internal(0).await;
});
}
async fn connect_internal(self: Arc<OmegaConnection>, mut retry: usize) { async fn connect_internal(self: Arc<OmegaConnection>, 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 { loop {
if retry > 5 { if retry > 5 {
log_err!( log_err!(
@ -100,6 +132,7 @@ impl OmegaConnection {
PrintType::Omega, PrintType::Omega,
"Max retry attempts reached, giving up." "Max retry attempts reached, giving up."
); );
*self.state.write().await = ConnectionState::Disconnected;
return; return;
} }
@ -107,16 +140,17 @@ impl OmegaConnection {
env::var("OMEGA_HOST").unwrap_or("wss://omega.tensamin.net/ws/omikron".to_string()); env::var("OMEGA_HOST").unwrap_or("wss://omega.tensamin.net/ws/omikron".to_string());
match connect_async(&url_str).await { match connect_async(&url_str).await {
Ok((ws_stream, _)) => { Ok((ws_stream, _)) => {
*self.is_connected.write().await = true;
log_in!(0, PrintType::Omega, "WebSocket connected to {}", url_str); log_in!(0, PrintType::Omega, "WebSocket connected to {}", url_str);
retry = 0; retry = 0;
let (write, read) = ws_stream.split(); let (write, read) = ws_stream.split();
*self.read.write().await = Some(read); *self.read.lock().await = Some(read);
*self.write.write().await = Some(write); *self.write.lock().await = Some(write);
let cloned_self = self.clone(); *self.state.write().await = ConnectionState::Connected;
tokio::spawn(async move {
cloned_self.clone().read_loop().await; 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(); let cloned_self = self.clone();
@ -137,148 +171,157 @@ impl OmegaConnection {
); );
WAITING_TASKS.insert( WAITING_TASKS.insert(
id, id,
Box::new(|selfc, cv| { WaitingTask {
if cv.is_type(CommunicationType::error_not_found) { task: Box::new(|selfc, cv| {
log_err!(0, if cv.is_type(CommunicationType::error_not_found) {
PrintType::Omega, log_err!(0,
"Identification failed: Omikron ID not found on Omega.", 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),
); );
return false;
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<JsonValue> = Vec::new();
let mut connected_user_ids: Vec<JsonValue> = 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);
} }
}); 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<JsonValue> = Vec::new();
let mut connected_user_ids: Vec<JsonValue> = 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 cloned_self.send_message(&identify_msg).await
}); });
let cloned_self = self.clone(); let ping_pong_self = self.clone();
let handle = tokio::spawn(async move { let handle = tokio::spawn(async move {
loop { loop {
if *cloned_self.is_connected.read().await == false { if *ping_pong_self.state.read().await != ConnectionState::Connected {
break; break;
} }
cloned_self.send_ping().await; ping_pong_self.send_ping().await;
sleep(Duration::from_secs(5)).await; sleep(Duration::from_secs(5)).await;
} }
}); });
*self.is_connected.write().await = true;
*self.pingpong.lock().await = Some(handle); *self.pingpong.lock().await = Some(handle);
while *self.is_connected.read().await { read_loop_handle.await.unwrap_or_else(|e| {
sleep(Duration::from_secs(2)).await; log_err!(0, PrintType::Omega, "Read loop task failed: {}", e)
} });
*self.read.write().await = None;
*self.write.write().await = None; *self.read.lock().await = None;
*self.write.lock().await = None;
*self.state.write().await = ConnectionState::Disconnected;
log_err!(0, PrintType::Omega, "Connection lost. Retrying..."); log_err!(0, PrintType::Omega, "Connection lost. Retrying...");
retry += 1; retry += 1;
sleep(Duration::from_secs(2)).await; sleep(Duration::from_secs(2)).await;
} }
Err(e) => { Err(e) => {
*self.state.write().await = ConnectionState::Disconnected;
log_err!( log_err!(
0, 0,
PrintType::Omega, PrintType::Omega,
@ -295,15 +338,13 @@ impl OmegaConnection {
} }
async fn read_loop(self: Arc<Self>) { async fn read_loop(self: Arc<Self>) {
let mut reader = match self.read.lock().await.take() {
Some(reader) => reader,
None => return,
};
loop { loop {
let msg = { let msg = reader.next().await;
let mut guard = self.read.write().await;
let ws = match guard.as_mut() {
Some(ws) => ws,
_ => break,
};
ws.next().await
};
match msg { match msg {
Some(Ok(Message::Text(msg))) => { Some(Ok(Message::Text(msg))) => {
@ -314,9 +355,8 @@ impl OmegaConnection {
} }
let msg_id = cv.get_id(); let msg_id = cv.get_id();
log_in!(0, PrintType::Omega, "{}", &cv.to_json().to_string()); log_in!(0, PrintType::Omega, "{}", &cv.to_json().to_string());
// Handle waiting tasks
if let Some(task) = WAITING_TASKS.remove(&msg_id) { 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; continue;
} }
} else { } 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) { 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 let Some(ws) = guard.as_mut() {
if !cv.is_type(CommunicationType::ping) { let _ = ws.send(Message::Text(msg.into())).await;
log_out!(0, PrintType::Omega, "{}", &cv.to_json().to_string());
}
let _ = ws
.send(Message::Text(cv.to_json().to_string().into()))
.await;
} }
} }
@ -378,20 +414,22 @@ impl OmegaConnection {
WAITING_TASKS.insert( WAITING_TASKS.insert(
msg_id, msg_id,
Box::new( WaitingTask {
move |_: Arc<OmegaConnection>, response: CommunicationValue| { task: Box::new(
let _ = Box::pin(async move |_: CommunicationValue| { move |_: Arc<OmegaConnection>, response: CommunicationValue| {
let rho = rho_manager::get_rho_con_for_user(user_id).await; tokio::spawn(async move {
if let Some(rho) = rho { let rho = rho_manager::get_rho_con_for_user(user_id).await;
for client in rho.get_client_connections_for_user(user_id).await { if let Some(rho) = rho {
client.send_message(&response).await; for client in rho.get_client_connections_for_user(user_id).await {
client.send_message(&response).await;
}
} }
} });
true true
}); },
true ),
}, inserted_at: Instant::now(),
), },
); );
OmegaConnection::send_global(cv).await; OmegaConnection::send_global(cv).await;
@ -402,17 +440,15 @@ impl OmegaConnection {
} }
pub async fn await_connection(&self, timeout_duration: Option<Duration>) -> Result<(), String> { pub async fn await_connection(&self, timeout_duration: Option<Duration>) -> Result<(), String> {
if *self.is_connected.read().await { if *self.state.read().await == ConnectionState::Connected {
return Ok(()); return Ok(());
} }
let timeout = timeout_duration.unwrap_or(Duration::from_secs(10)); let timeout = timeout_duration.unwrap_or(Duration::from_secs(10));
let start = Instant::now(); let start = Instant::now();
loop { loop {
{ if *self.state.read().await == ConnectionState::Connected {
if *self.is_connected.read().await { return Ok(());
return Ok(());
}
} }
if start.elapsed() >= timeout { if start.elapsed() >= timeout {
@ -431,27 +467,29 @@ impl OmegaConnection {
cv: &CommunicationValue, cv: &CommunicationValue,
timeout_duration: Option<Duration>, timeout_duration: Option<Duration>,
) -> Result<CommunicationValue, String> { ) -> Result<CommunicationValue, String> {
let _ = self.await_connection(timeout_duration).await; self.await_connection(timeout_duration).await?;
let (tx, mut rx) = mpsc::channel(1); let (tx, mut rx) = mpsc::channel(1);
let msg_id = cv.get_id(); let msg_id = cv.get_id();
let task_tx = tx.clone(); let task_tx = tx.clone();
WAITING_TASKS.insert( WAITING_TASKS.insert(
msg_id, msg_id,
Box::new(move |_, response_cv| { WaitingTask {
let inner_tx = task_tx.clone(); task: Box::new(move |_, response_cv| {
tokio::spawn(async move { let inner_tx = task_tx.clone();
if let Err(e) = inner_tx.send(response_cv).await { tokio::spawn(async move {
log_err!( if inner_tx.send(response_cv).await.is_err() {
0, log_err!(
PrintType::Omega, 0,
"Failed to send response back to awaiter: {}", PrintType::Omega,
e "Failed to send response back to awaiter",
); );
} }
}); });
true true
}), }),
inserted_at: Instant::now(),
},
); );
self.send_message(cv).await; self.send_message(cv).await;

20
src/omega/ping_pong_task.rs
View file

@ -1,15 +1,21 @@
use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes}; use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes};
use crate::omega::omega_connection::OmegaConnection; use crate::omega::omega_connection::OmegaConnection;
use json::number::Number; use json::number::Number;
use std::time::Duration;
use tokio::time::Instant; use tokio::time::Instant;
use uuid::Uuid; use uuid::Uuid;
const PING_TIMEOUT: Duration = Duration::from_secs(30);
impl OmegaConnection { impl OmegaConnection {
pub async fn send_ping(&self) { pub async fn send_ping(&self) {
let uuid = Uuid::new_v4(); let uuid = Uuid::new_v4();
let send_time = Instant::now(); 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; self.send_ping_message(uuid).await;
} }
@ -27,16 +33,10 @@ impl OmegaConnection {
/// Handles incoming pong and calculates latency /// Handles incoming pong and calculates latency
pub async fn handle_pong(&self, cv: &CommunicationValue, _log: bool) { pub async fn handle_pong(&self, cv: &CommunicationValue, _log: bool) {
let id = cv.get_id(); let id = cv.get_id();
let send_time_opt = { let mut message_send_times = self.message_send_times.lock().await;
let queue = self.message_send_times.lock().await; if let Some(send_time) = message_send_times.remove(&id) {
queue.get(&id).cloned()
};
if let Some(send_time) = send_time_opt {
let ping = Instant::now().duration_since(send_time).as_millis() as i64; 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;
*self.last_ping.lock().await = ping as i64;
} }
} }
} }

26
src/rho/client_connection.rs
View file

@ -15,7 +15,7 @@ use uuid::Uuid;
use super::{rho_connection::RhoConnection, rho_manager}; use super::{rho_connection::RhoConnection, rho_manager};
use crate::anonymous_clients::anonymous_manager; use crate::anonymous_clients::anonymous_manager;
use crate::calls::{call_manager, call_util}; 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_helper::{load_public_key, public_key_to_base64};
use crate::util::crypto_util::{DataFormat, SecurePayload}; use crate::util::crypto_util::{DataFormat, SecurePayload};
use crate::util::logger::PrintType; use crate::util::logger::PrintType;
@ -374,19 +374,17 @@ impl ClientConnection {
} }
async fn handle_omega_forward(self: Arc<Self>, cv: CommunicationValue) { async fn handle_omega_forward(self: Arc<Self>, cv: CommunicationValue) {
let client_for_closure = self.clone(); let client_for_closure = self.clone();
WAITING_TASKS.insert( tokio::spawn(async move {
cv.get_id(), let response_cv = get_omega_connection()
Box::new(move |_, response_cv| { .await_response(
let client = client_for_closure.clone(); &cv.with_sender(*self.user_id.read().await),
tokio::spawn(async move { Some(Duration::from_secs(20)),
client.send_message(&response_cv).await; )
}); .await;
true if let Ok(response_cv) = response_cv {
}), client_for_closure.send_message(&response_cv).await;
); }
get_omega_connection() });
.send_message(&cv.with_sender(*self.user_id.read().await))
.await;
} }
/// Handle ping message /// Handle ping message

25
src/rho/iota_connection.rs
View file

@ -5,7 +5,6 @@ use crate::get_public_key;
use crate::log_err; use crate::log_err;
use crate::log_in; use crate::log_in;
use crate::log_out; use crate::log_out;
use crate::omega::omega_connection::WAITING_TASKS;
use crate::omega::omega_connection::get_omega_connection; use crate::omega::omega_connection::get_omega_connection;
use crate::util::crypto_helper::load_public_key; use crate::util::crypto_helper::load_public_key;
use crate::util::crypto_helper::public_key_to_base64; use crate::util::crypto_helper::public_key_to_base64;
@ -460,19 +459,17 @@ impl IotaConnection {
async fn handle_omega_forward(self: Arc<Self>, cv: CommunicationValue) { async fn handle_omega_forward(self: Arc<Self>, cv: CommunicationValue) {
let iota_for_closure = self.clone(); let iota_for_closure = self.clone();
WAITING_TASKS.insert( tokio::spawn(async move {
cv.get_id(), let response_cv = get_omega_connection()
Box::new(move |_, response_cv| { .await_response(
let iota = iota_for_closure.clone(); &cv.with_sender(*self.iota_id.read().await),
tokio::spawn(async move { Some(Duration::from_secs(20)),
iota.send_message(&response_cv).await; )
}); .await;
true if let Ok(response_cv) = response_cv {
}), iota_for_closure.send_message(&response_cv).await;
); }
get_omega_connection() });
.send_message(&cv.with_sender(*self.iota_id.read().await))
.await;
} }
/// Handle ping message /// Handle ping message
async fn handle_ping(&self, cv: CommunicationValue) { async fn handle_ping(&self, cv: CommunicationValue) {