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

Connection With Omega

Browse source
This commit is contained in:
Alex Emmet 2026-01-05 02:17:12 +01:00
commit 504e5ff015
20 changed files with 880 additions and 320 deletions
Download patch file Download diff file Expand all files Collapse all files

1
src/omega/mod.rs
View file

@ -1 +1,2 @@
pub mod omega_connection;
pub mod ping_pong_task;

373
src/omega/omega_connection.rs
View file

@ -1,75 +1,251 @@
use std::sync::Arc;
use std::time::Duration;
use crate::data::communication::{CommunicationType, CommunicationValue};
use crate::util::print::PrintType;
use crate::util::print::{line, line_err};
use crate::{
data::{
communication::DataTypes,
user::{User, UserStatus},
},
rho::rho_manager,
util::config_util::CONFIG,
use async_tungstenite::{
WebSocketReceiver, WebSocketSender, WebSocketStream,
stream::Stream,
tokio::{TokioAdapter, connect_async},
tungstenite::protocol::Message,
};
use async_tungstenite::tungstenite::protocol::Message;
use dashmap::DashMap;
use futures::StreamExt;
use json::JsonValue;
use futures::prelude::*;
use json::{JsonValue, number::Number};
use once_cell::sync::Lazy;
use tokio::sync::Mutex;
use tokio::time::sleep;
use tokio_util::compat::Compat;
use tungstenite::{Utf8Bytes, connect};
use std::{collections::HashMap, env, sync::Arc, time::Duration};
use tokio::{
net::TcpStream,
sync::{Mutex, RwLock},
time::{Instant, sleep},
};
use tokio_native_tls::TlsStream;
use uuid::Uuid;
static WAITING_TASKS: Lazy<DashMap<Uuid, Box<dyn Fn(CommunicationValue) -> bool + Send + Sync>>> =
Lazy::new(DashMap::new);
use crate::{
auth::crypto_helper::decrypt,
data::{
communication::{CommunicationType, CommunicationValue, DataTypes},
user::UserStatus,
},
get_private_key, log_in, log_out,
rho::rho_manager,
util::logger::PrintType,
};
use crate::{auth::crypto_helper::secret_key_to_base64, log_err};
static WAITING_TASKS: Lazy<
DashMap<Uuid, Box<dyn Fn(Arc<OmegaConnection>, CommunicationValue) -> bool + Send + Sync>>,
> = Lazy::new(DashMap::new);
static GENERIC_TASK: Lazy<
Mutex<Option<Box<dyn Fn(Arc<OmegaConnection>, CommunicationValue) -> bool + Send + Sync>>>,
> = Lazy::new(|| Mutex::new(None));
static OMEGA_CONNECTION: Lazy<Arc<OmegaConnection>> = Lazy::new(|| {
let conn = Arc::new(OmegaConnection::new());
let conn_clone = conn.clone();
tokio::spawn(async move {
conn_clone.connect_internal(0).await;
});
conn
});
pub fn get_omega_connection() -> Arc<OmegaConnection> {
OMEGA_CONNECTION.clone()
}
#[derive(Clone)]
pub struct OmegaConnection {
ws_stream:
Arc<Mutex<Option<async_tungstenite::WebSocketStream<Compat<tokio::net::TcpStream>>>>>,
write: Arc<
RwLock<
Option<
WebSocketSender<
Stream<TokioAdapter<TcpStream>, TokioAdapter<TlsStream<TcpStream>>>,
>,
>,
>,
>,
read: Arc<
RwLock<
Option<
WebSocketReceiver<
Stream<TokioAdapter<TcpStream>, TokioAdapter<TlsStream<TcpStream>>>,
>,
>,
>,
>,
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>>,
}
impl OmegaConnection {
pub fn new() -> Self {
OmegaConnection {
ws_stream: Arc::new(Mutex::new(None)),
read: Arc::new(RwLock::new(None)),
write: Arc::new(RwLock::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)),
}
}
pub async fn connect(&self) {
self.connect_internal(0).await;
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, mut retry: usize) {
async fn connect_internal(self: Arc<OmegaConnection>, mut retry: usize) {
loop {
if retry > 5 {
line_err(
PrintType::OmegaIn,
&"Max retry attempts reached, giving up.",
);
log_err!(PrintType::Omega, "Max retry attempts reached, giving up.");
return;
}
match connect("wss://tensamin.methanium.net/ws/omega") {
Ok((_, _)) => {
let url_str =
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!(PrintType::Omega, "WebSocket connected to {}", url_str);
retry = 0;
let identify_msg = CommunicationValue::new(CommunicationType::identification)
.add_data(
DataTypes::uuid,
JsonValue::String(CONFIG.read().await.omikron_id.to_string()),
);
self.send_message(&identify_msg).await;
let (write, read) = ws_stream.split();
*self.read.write().await = Some(read);
*self.write.write().await = Some(write);
let ws_stream_clone = self.ws_stream.clone();
let cloned_self = self.clone();
tokio::spawn(async move {
OmegaConnection::read_loop(ws_stream_clone).await;
cloned_self.clone().read_loop().await;
});
let cloned_self = self.clone();
tokio::spawn(async move {
let id = Uuid::new_v4();
let identify_msg =
CommunicationValue::new(CommunicationType::identification)
.with_id(id)
.add_data(
DataTypes::omikron,
JsonValue::Number(Number::from(
env::var("ID")
.unwrap_or("0".to_string())
.parse::<i64>()
.unwrap_or(0),
)),
);
WAITING_TASKS.insert(
id,
Box::new(|selfc, cv| {
if cv.is_type(CommunicationType::error_not_found) {
log_err!(
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(
&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,
Box::new(|_self, final_cv| {
if !final_cv
.is_type(CommunicationType::identification_response)
{
log_err!(
PrintType::Omega,
"Expected identification_response, got something else.",
);
return false;
}
log_err!(
PrintType::Omega,
"Successfully identified with Omega.",
);
true
}),
);
selfc.send_message(&response_msg).await;
Ok::<(), String>(())
};
if let Err(e) = task.await {
log_err!(PrintType::Omega, "{}", &e);
}
});
true
}),
);
cloned_self.send_message(&identify_msg).await
});
let cloned_self = self.clone();
let handle = tokio::spawn(async move {
loop {
if *cloned_self.is_connected.read().await == false {
break;
}
cloned_self.send_ping().await;
sleep(Duration::from_secs(1)).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;
log_err!(PrintType::Omega, "Connection lost. Retrying...");
retry += 1;
sleep(Duration::from_secs(2)).await;
}
Err(e) => {
line_err(
PrintType::OmegaIn,
&format!("WebSocket connection failed (attempt {}): {}", retry, e),
log_err!(
PrintType::Omega,
"WebSocket connection failed (attempt {}): {}",
retry + 1,
e,
);
retry += 1;
sleep(Duration::from_secs(2)).await;
@ -79,72 +255,60 @@ impl OmegaConnection {
}
}
async fn read_loop(
ws_stream: Arc<
Mutex<Option<async_tungstenite::WebSocketStream<Compat<tokio::net::TcpStream>>>>,
>,
) {
async fn read_loop(self: Arc<Self>) {
loop {
let mut lock = ws_stream.lock().await;
let Some(ws) = lock.as_mut() else {
break;
let msg = {
let mut guard = self.read.write().await;
let ws = match guard.as_mut() {
Some(ws) => ws,
None => break,
};
ws.next().await
};
match ws.next().await {
match msg {
Some(Ok(Message::Text(msg))) => {
let cv = CommunicationValue::from_json(&msg);
if cv.is_type(CommunicationType::pong) {
self.handle_pong(&cv, true).await;
continue;
}
let msg_id = cv.get_id();
log_in!(PrintType::Omikron, "{}", &cv.to_json().to_string());
// Handle waiting tasks
if let Some(task) = WAITING_TASKS.remove(&msg_id) {
if (task.1)(cv.clone()) {
continue;
if (task.1)(self.clone(), cv.clone()) {
// continue in the read_loop
}
}
// Handle CLIENT_CHANGED
if cv.is_type(CommunicationType::client_changed) {
let iota_id = cv
.get_data(DataTypes::iota_id)
.unwrap()
.as_i64()
.unwrap_or(0);
let user_id = cv
.get_data(DataTypes::user_id)
.unwrap()
.as_i64()
.unwrap_or(0);
let status_str = cv
.get_data(DataTypes::user_state)
.unwrap()
.as_str()
.unwrap();
let status = UserStatus::from_string(&status_str)
.unwrap_or(UserStatus::iota_offline);
let user = User::new(iota_id, user_id, status);
for rho_con in rho_manager::get_all_connections().await {
rho_con.are_they_interested(&user).await;
} else {
// Handle generic task
let generic_task_option = GENERIC_TASK.lock().await;
if let Some(generic_task) = generic_task_option.as_ref() {
if generic_task(self.clone(), cv.clone()) {
// continue in the read_loop
}
}
}
}
Some(Ok(Message::Close(_))) | None => {
break;
}
Some(Err(_)) => {
break;
}
Some(Ok(Message::Close(_))) | None => break,
Some(Err(_)) => break,
_ => {}
}
}
*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.ws_stream.lock().await;
let mut guard = self.write.write().await;
if let Some(ws) = guard.as_mut() {
line(PrintType::OmegaOut, &cv.to_json().to_string());
if !cv.is_type(CommunicationType::ping) {
log_out!(PrintType::Omega, "{}", &cv.to_json().to_string());
}
let _ = ws
.send(Message::Text(Utf8Bytes::from(cv.to_json().to_string())))
.send(Message::Text(cv.to_json().to_string().into()))
.await;
}
}
@ -187,25 +351,26 @@ impl OmegaConnection {
WAITING_TASKS.insert(
msg_id,
Box::new(move |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;
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;
}
}
}
true
});
true
});
true
}),
},
),
);
OmegaConnection::send_global(cv).await;
}
async fn send_global(cv: CommunicationValue) {
let conn = OmegaConnection::new();
conn.send_message(&cv).await;
OMEGA_CONNECTION.send_message(&cv).await;
}
}

42
src/omega/ping_pong_task.rs Normal file
View file

@ -0,0 +1,42 @@
use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes};
use crate::omega::omega_connection::OmegaConnection;
use json::number::Number;
use tokio::time::Instant;
use uuid::Uuid;
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);
self.send_ping_message(uuid).await;
}
pub async fn send_ping_message(&self, uuid: Uuid) {
let ping_message = CommunicationValue::new(CommunicationType::ping)
.with_id(uuid)
.add_data_num(
DataTypes::last_ping,
Number::from(*self.last_ping.lock().await),
);
self.send_message(&ping_message).await;
}
/// 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 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;
}
}
}