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::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<Self>, 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

4
src/main.rs
View file

@ -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!(
"{}:{}",

426
src/omega/omega_connection.rs
View file

@ -31,9 +31,12 @@ use crate::{
},
};
pub static WAITING_TASKS: Lazy<
DashMap<Uuid, Box<dyn Fn(Arc<OmegaConnection>, CommunicationValue) -> bool + Send + Sync>>,
> = Lazy::new(DashMap::new);
pub struct WaitingTask {
pub task: Box<dyn Fn(Arc<OmegaConnection>, CommunicationValue) -> bool + Send + Sync>,
pub inserted_at: Instant,
}
pub static WAITING_TASKS: Lazy<DashMap<Uuid, WaitingTask>> = Lazy::new(DashMap::new);
static OMEGA_CONNECTION: Lazy<Arc<OmegaConnection>> = Lazy::new(|| {
let conn = Arc::new(OmegaConnection::new());
@ -41,6 +44,14 @@ static OMEGA_CONNECTION: Lazy<Arc<OmegaConnection>> = 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<OmegaConnection> {
#[derive(Clone)]
pub struct OmegaConnection {
write: Arc<
RwLock<
Mutex<
Option<
WebSocketSender<
Stream<TokioAdapter<TcpStream>, TokioAdapter<TlsStream<TcpStream>>>,
@ -61,7 +72,7 @@ pub struct OmegaConnection {
>,
>,
read: Arc<
RwLock<
Mutex<
Option<
WebSocketReceiver<
Stream<TokioAdapter<TcpStream>, TokioAdapter<TlsStream<TcpStream>>>,
@ -72,27 +83,48 @@ pub struct OmegaConnection {
pingpong: Arc<Mutex<Option<tokio::task::JoinHandle<()>>>>,
pub last_ping: Arc<Mutex<i64>>,
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 {
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<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) {
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<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);
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<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
});
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<Self>) {
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<OmegaConnection>, 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<OmegaConnection>, 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<Duration>) -> 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<Duration>,
) -> Result<CommunicationValue, String> {
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;

20
src/omega/ping_pong_task.rs
View file

@ -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;
}
}
}

26
src/rho/client_connection.rs
View file

@ -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<Self>, 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

25
src/rho/iota_connection.rs
View file

@ -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<Self>, 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) {