iota/src/omikron/omikron_connection.rs
Alex Emmet 7bf90a3699 Reload Endpoint Reload Keybind Reload Logic
Remove TU Endpoint
2025-12-04 17:35:07 +00:00

606 lines
30 KiB
Rust

use crate::auth::local_auth;
use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes};
use crate::gui::log_panel::{log_cv, log_message, log_message_trans};
use crate::users::contact::Contact;
use crate::users::user_community_util::UserCommunityUtil;
use crate::util::chat_files;
use crate::util::chats_util::{get_user, get_users, mod_user};
use crate::util::file_util::{get_children, load_file, save_file};
use crate::{ACTIVE_TASKS, RECONNECT, SHUTDOWN};
use futures::Stream;
use futures::stream::{SplitSink, SplitStream};
use futures_util::sink::Sink;
use futures_util::{SinkExt, StreamExt};
use hyper::upgrade::Upgraded;
use hyper_util::rt::TokioIo;
use json::JsonValue;
use std::collections::HashMap;
use std::str::FromStr;
use std::sync::{Arc, LazyLock};
use std::time::{SystemTime, UNIX_EPOCH};
use tokio::sync::{Mutex, RwLock};
use tokio::time::{Duration, Instant, sleep};
use tokio_tungstenite::{connect_async, tungstenite::protocol::Message};
use tungstenite::Utf8Bytes;
use uuid::Uuid;
pub static OMIKRON_CONNECTION: LazyLock<Arc<RwLock<Option<Arc<OmikronConnection>>>>> =
LazyLock::new(|| Arc::new(RwLock::new(None)));
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ConnectionVariant {
Omikron,
ClientUnauthenticated,
ClientAuthenticated,
}
#[derive(Clone)]
pub struct OmikronConnection {
pub variant: Arc<RwLock<ConnectionVariant>>,
pub user_id: Arc<RwLock<Option<Uuid>>>,
pub(crate) writer:
Arc<Mutex<Option<Box<dyn Sink<Message, Error = tungstenite::Error> + Send + Unpin>>>>,
waiting: Arc<Mutex<HashMap<Uuid, Box<dyn Fn(CommunicationValue) + Send + Sync>>>>, // waiting for responses
pingpong: Arc<Mutex<Option<tokio::task::JoinHandle<()>>>>, // ping-pong handler
pub last_ping: Arc<Mutex<i64>>,
pub message_send_times: Arc<Mutex<HashMap<Uuid, Instant>>>,
pub is_connected: Arc<Mutex<bool>>,
}
impl OmikronConnection {
pub fn new() -> Self {
Self {
variant: Arc::new(RwLock::new(ConnectionVariant::Omikron)),
user_id: Arc::new(RwLock::new(None)),
writer: Arc::new(Mutex::new(None)),
waiting: Arc::new(Mutex::new(HashMap::new())),
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(Mutex::new(false)),
}
}
pub async fn client(
writer: SplitSink<tokio_tungstenite::WebSocketStream<TokioIo<Upgraded>>, Message>,
reader: SplitStream<tokio_tungstenite::WebSocketStream<TokioIo<Upgraded>>>,
) -> Arc<Self> {
let connection = Arc::new(Self {
variant: Arc::new(RwLock::new(ConnectionVariant::ClientUnauthenticated)),
user_id: Arc::new(RwLock::new(None)),
writer: Arc::new(Mutex::new(Some(Box::new(writer)
as Box<dyn Sink<Message, Error = tungstenite::Error> + Send + Unpin>))),
waiting: Arc::new(Mutex::new(HashMap::new())),
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(Mutex::new(false)),
});
let boxed_reader: Box<
dyn Stream<Item = Result<Message, tungstenite::Error>> + Send + Unpin,
> = Box::new(reader);
connection.spawn_listener(boxed_reader).await;
connection
}
pub async fn is_connected(&self) -> bool {
*self.is_connected.lock().await
}
/// Connect loop with retry
pub async fn connect(self: &Arc<Self>) {
loop {
if *SHUTDOWN.read().await {
break;
}
if *RECONNECT.read().await {
break;
}
match connect_async("wss://app.tensamin.net/ws/iota/").await {
Ok((ws_stream, _)) => {
let (write_half, read_half) = ws_stream.split();
*self.writer.lock().await = Some(Box::new(write_half));
let boxed_reader: Box<
dyn Stream<Item = Result<Message, tungstenite::Error>> + Send + Unpin,
> = Box::new(read_half);
self.clone().spawn_listener(boxed_reader).await;
let cloned_self = self.clone();
{
ACTIVE_TASKS.lock().unwrap().push("PingPong".to_string());
}
let handle = tokio::spawn(async move {
loop {
if *SHUTDOWN.read().await {
break;
}
if *RECONNECT.read().await {
break;
}
cloned_self.send_ping().await;
sleep(Duration::from_secs(1)).await;
}
});
{
ACTIVE_TASKS
.lock()
.unwrap()
.retain(|t| !t.eq(&"PingPong".to_string()));
}
*self.is_connected.lock().await = true;
*self.pingpong.lock().await = Some(handle);
break;
}
Err(_) => {
*self.is_connected.lock().await = false;
sleep(Duration::from_secs(2)).await;
}
}
}
}
pub async fn send_message(&self, msg: String) {
Self::send_message_static(&self.writer, msg).await
}
pub async fn set_variant(self: &Arc<Self>, variant: ConnectionVariant) {
*self.variant.write().await = variant;
}
pub async fn set_user_id(self: &Arc<Self>, user_id: Uuid) {
*self.user_id.write().await = Some(user_id);
}
/// Listener for all incoming messages
async fn spawn_listener(
self: &Arc<Self>,
mut read_half: Box<dyn Stream<Item = Result<Message, tungstenite::Error>> + Send + Unpin>,
) {
let waiting_out = self.waiting.clone();
let writer_out = self.writer.clone();
let is_connected_out = self.is_connected.clone();
let sel_out = self.clone();
let variant = self.variant.clone();
let sel_arc_out = self.clone();
{
ACTIVE_TASKS.lock().unwrap().push("Listener".to_string());
}
tokio::spawn(async move {
while let Some(msg) = read_half.next().await {
if *SHUTDOWN.read().await {
break;
}
if *RECONNECT.read().await {
break;
}
let waiting = waiting_out.clone();
let writer = writer_out.clone();
let is_connected = is_connected_out.clone();
let sel = sel_out.clone();
let variant = variant.clone();
let sel_arc = sel_arc_out.clone();
tokio::spawn(async move {
match msg {
Ok(Message::Close(Some(frame))) => {
log_message(format!("[Omikron] Closed: {:?}", frame));
*is_connected.lock().await = false;
return;
}
Ok(Message::Text(text)) => {
let mut cv = CommunicationValue::from_json(&text);
if cv.is_type(CommunicationType::pong) {
sel.handle_pong(&cv, true).await;
return;
}
let com = variant.read().await.clone();
if com == ConnectionVariant::ClientUnauthenticated {
if cv.is_type(CommunicationType::identification) {
// Extract user ID
let user_id = match cv.get_data(DataTypes::user_id) {
Some(id_str) => {
match Uuid::parse_str(&id_str.to_string()) {
Ok(id) => id,
Err(_) => {
sel_arc.send_message(
CommunicationValue::new(CommunicationType::error_invalid_user_id)
.with_id(cv.get_id())
.to_json()
.to_string()
)
.await;
return;
}
}
}
None => {
sel_arc
.send_message(
CommunicationValue::new(
CommunicationType::error_invalid_user_id,
)
.with_id(cv.get_id())
.to_json()
.to_string(),
)
.await;
return;
}
};
// Validate private key
if let Some(private_key_hash) =
cv.get_data(DataTypes::private_key_hash)
{
log_message(format!(
"private_key_hash: {}",
private_key_hash
));
let is_valid = local_auth::is_private_key_valid(
&user_id,
&private_key_hash.to_string(),
);
if !is_valid {
log_message("Invalid private key");
sel_arc.send_message(
CommunicationValue::new(
CommunicationType::error_invalid_private_key,
)
.with_id(cv.get_id())
.to_json()
.to_string(),
)
.await;
return;
}
} else {
log_message("Missing private key");
sel_arc
.send_message(
CommunicationValue::new(
CommunicationType::error_invalid_private_key,
)
.with_id(cv.get_id())
.to_json()
.to_string(),
)
.await;
return;
}
// Set identification data
sel_arc.set_user_id(user_id).await;
sel_arc
.set_variant(ConnectionVariant::ClientAuthenticated)
.await;
let response = CommunicationValue::new(
CommunicationType::identification_response,
)
.with_id(cv.get_id());
sel_arc.send_message(response.to_json().to_string()).await;
}
}
// ************************************************ //
// Direct messages //
// ************************************************ //
log_cv(&cv);
if let Some(x) = waiting.lock().await.remove(&cv.get_id()) {
x(cv);
return;
}
if cv.is_type(CommunicationType::message_other_iota) {
let sender_id = &cv.get_sender().unwrap();
let receiver_id = &cv.get_receiver().unwrap();
chat_files::add_message(
cv.get_data(DataTypes::send_time)
.unwrap_or(&JsonValue::new_object())
.as_i64()
.unwrap_or(
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_millis()
as i64,
) as u128,
false,
*receiver_id,
*sender_id,
cv.get_data(DataTypes::content).unwrap().as_str().unwrap(),
);
let user_forward =
CommunicationValue::new(CommunicationType::message_live)
.with_id(cv.get_id())
.with_receiver(*receiver_id)
.add_data(
DataTypes::send_time,
cv.get_data(DataTypes::send_time).unwrap().clone(),
)
.add_data(
DataTypes::message,
cv.get_data(DataTypes::content).unwrap().clone(),
)
.add_data(
DataTypes::sender_id,
JsonValue::String(cv.get_sender().unwrap().to_string()),
);
Self::send_message_static(
&writer.clone(),
user_forward.to_json().to_string(),
)
.await;
return;
}
if cv.is_type(CommunicationType::message_send) {
let my_id = cv.get_sender().unwrap();
let other_id = Uuid::from_str(
&*cv.get_data(DataTypes::receiver_id).unwrap().to_string(),
)
.unwrap();
chat_files::add_message(
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_millis() as u128,
true,
my_id,
other_id,
&*cv.get_data(DataTypes::content).unwrap().to_string(),
);
let ack = CommunicationValue::new(CommunicationType::message)
.with_id(cv.get_id())
.with_receiver(my_id);
Self::send_message_static(
&writer.clone(),
ack.to_json().to_string(),
)
.await;
let forward = CommunicationValue::forward_to_other_iota(&mut cv);
Self::send_message_static(
&writer.clone(),
forward.to_json().to_string(),
)
.await;
return;
}
if cv.is_type(CommunicationType::messages_get) {
let my_id = cv.get_sender().unwrap();
let partner_id = Uuid::from_str(
&*cv.get_data(DataTypes::user_id).unwrap().to_string(),
)
.unwrap();
let offset = cv
.get_data(DataTypes::offset)
.unwrap_or(&JsonValue::Null)
.to_string()
.parse::<i64>()
.unwrap_or(0);
let amount = cv
.get_data(DataTypes::amount)
.unwrap_or(&JsonValue::Null)
.to_string()
.parse::<i64>()
.unwrap_or(0);
let messages =
chat_files::get_messages(my_id, partner_id, offset, amount);
let resp = CommunicationValue::new(CommunicationType::messages_get)
.with_id(cv.get_id())
.with_receiver(my_id)
.add_data(DataTypes::messages, messages);
Self::send_message_static(
&writer.clone(),
resp.to_json().to_string(),
)
.await;
return;
}
if cv.is_type(CommunicationType::get_chats) {
let user_id = cv.get_sender().unwrap();
let users = get_users(user_id);
let resp = CommunicationValue::new(CommunicationType::get_chats)
.with_id(cv.get_id())
.with_receiver(user_id)
.add_data(DataTypes::user_ids, users);
Self::send_message_static(
&writer.clone(),
resp.to_json().to_string(),
)
.await;
return;
}
if cv.is_type(CommunicationType::add_chat) {
let user_id = cv.get_sender().unwrap();
let other_id = Uuid::from_str(
&*cv.get_data(DataTypes::user_id).unwrap().to_string(),
)
.unwrap();
let mut contact =
get_user(user_id, other_id).unwrap_or(Contact::new(other_id)); // needs ChatsUtil + Contact
contact.set_last_message_at(
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_millis() as i64,
);
mod_user(user_id, &contact);
let resp = CommunicationValue::new(CommunicationType::add_chat)
.with_id(cv.get_id())
.with_receiver(user_id);
Self::send_message_static(
&writer.clone(),
resp.to_json().to_string(),
)
.await;
return;
}
if cv.is_type(CommunicationType::add_community) {
UserCommunityUtil::add_community(
cv.get_sender().unwrap(),
cv.get_data(DataTypes::community_address)
.unwrap()
.to_string(),
cv.get_data(DataTypes::community_title).unwrap().to_string(),
cv.get_data(DataTypes::position).unwrap().to_string(),
);
let resp =
CommunicationValue::new(CommunicationType::add_community)
.with_id(cv.get_id())
.with_receiver(cv.get_sender().unwrap());
Self::send_message_static(
&writer.clone(),
resp.to_json().to_string(),
)
.await;
return;
}
if cv.is_type(CommunicationType::get_communities) {
let resp =
CommunicationValue::new(CommunicationType::get_communities)
.with_id(cv.get_id())
.with_receiver(cv.get_sender().unwrap())
.add_array(
DataTypes::communities,
UserCommunityUtil::get_communities(
cv.get_sender().unwrap(),
),
);
Self::send_message_static(
&writer.clone(),
resp.to_json().to_string(),
)
.await;
return;
}
if cv.is_type(CommunicationType::remove_community) {
UserCommunityUtil::remove_community(
cv.get_sender().unwrap(),
cv.get_data(DataTypes::community_address)
.unwrap()
.to_string(),
); // needs UserCommunityUtil
let resp =
CommunicationValue::new(CommunicationType::remove_community)
.with_id(cv.get_id())
.with_receiver(cv.get_sender().unwrap());
Self::send_message_static(
&writer.clone(),
resp.to_json().to_string(),
)
.await;
return;
}
if cv.is_type(CommunicationType::settings_save) {
let my_id = cv.get_sender().unwrap();
let settings_name =
cv.get_data(DataTypes::settings_name).unwrap().to_string();
let settings_value =
cv.get_data(DataTypes::payload).unwrap().to_string();
save_file(
&format!("users/{}/settings/", my_id),
&format!("{}.settings", settings_name),
&settings_value,
);
let response =
CommunicationValue::new(CommunicationType::settings_save)
.with_receiver(my_id)
.with_id(cv.get_id());
Self::send_message_static(
&writer.clone(),
response.to_json().to_string(),
)
.await;
return;
}
if cv.is_type(CommunicationType::settings_load) {
let my_id = cv.get_sender().unwrap();
let settings_name =
cv.get_data(DataTypes::settings_name).unwrap().to_string();
let settings_value_str = load_file(
&format!("users/{}/settings/", my_id),
&format!("{}.settings", settings_name),
);
let settings_value_json = JsonValue::from(settings_value_str);
let response =
CommunicationValue::new(CommunicationType::settings_load)
.with_id(cv.get_id())
.with_receiver(my_id)
.add_data(DataTypes::payload, settings_value_json)
.add_data_str(DataTypes::settings_name, settings_name);
Self::send_message_static(
&writer.clone(),
response.to_json().to_string(),
)
.await;
return;
}
if cv.is_type(CommunicationType::settings_list) {
let my_id = cv.get_sender().unwrap();
let settings = get_children(&format!("users/{}/settings/", my_id));
let mut settings_json = JsonValue::new_array();
for s in settings {
let s = s.replace(".settings", "");
if s.is_empty() {
continue;
}
let _ = settings_json.push(JsonValue::String(s));
}
let response =
CommunicationValue::new(CommunicationType::settings_list)
.with_id(cv.get_id())
.with_receiver(my_id)
.add_data(DataTypes::settings, settings_json);
Self::send_message_static(
&writer.clone(),
response.to_json().to_string(),
)
.await;
return;
}
}
Err(e) => {
log_message(format!("[Omikron] Error: {}", e));
*is_connected.lock().await = false;
return;
}
_ => {}
}
});
}
});
{
ACTIVE_TASKS
.lock()
.unwrap()
.retain(|t| !t.eq(&"Listener".to_string()));
}
}
pub async fn send_message_static(
writer: &Arc<
Mutex<Option<Box<dyn Sink<Message, Error = tungstenite::Error> + Send + Unpin>>>,
>,
msg: String,
) {
let mut guard = writer.lock().await;
if let Some(writer) = guard.as_mut() {
if let Ok(_) = writer.send(Message::Text(Utf8Bytes::from(msg))).await {
if let Ok(_) = writer.flush().await {
return;
}
}
}
log_message_trans("send_message_failed");
}
}