register logic (bricked) user management !! omikron con wait for

response method !!

src/omikron/omikron_connection.rs

@@ -1,28 +1,35 @@
 use crate::auth::local_auth;
-use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes};
-use crate::gui::log_panel::{log_cv, log_message, log_message_format};
+use crate::gui::log_panel::{log_cv, log_message_format};
+use crate::{
+    data::communication::{CommunicationType, CommunicationValue, DataTypes},
+    gui::log_panel::{log_message, log_message_trans},
+    util::{config_util::CONFIG, crypto_helper},
+};
+use json::JsonValue;
+
 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::chats_util::{get_user, mod_user};
 use crate::util::file_util::{get_children, load_file, save_file};
+use crate::util::{chat_files, chats_util};
 use crate::{ACTIVE_TASKS, SHUTDOWN};
+use dashmap::DashMap;
 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 json::number::Number;
 use std::collections::HashMap;
 use std::sync::{Arc, LazyLock};
 use std::time::{SystemTime, UNIX_EPOCH};
-use tokio::sync::{Mutex, RwLock};
+use tokio::sync::{Mutex, RwLock, mpsc};
 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)));
 
@@ -39,7 +46,7 @@ pub struct OmikronConnection {
     pub user_id: Arc<RwLock<i64>>,
     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: Arc<DashMap<Uuid, Box<dyn Fn(CommunicationValue) + Send + Sync>>>,
     pingpong: Arc<Mutex<Option<tokio::task::JoinHandle<()>>>>,
     pub last_ping: Arc<Mutex<i64>>,
     pub message_send_times: Arc<Mutex<HashMap<Uuid, Instant>>>,
@@ -52,7 +59,7 @@ impl OmikronConnection {
             variant: Arc::new(RwLock::new(ConnectionVariant::Omikron)),
             user_id: Arc::new(RwLock::new(0)),
             writer: Arc::new(Mutex::new(None)),
-            waiting: Arc::new(Mutex::new(HashMap::new())),
+            waiting: Arc::new(DashMap::new()),
             pingpong: Arc::new(Mutex::new(None)),
             last_ping: Arc::new(Mutex::new(-1)),
             message_send_times: Arc::new(Mutex::new(HashMap::new())),
@@ -68,7 +75,7 @@ impl OmikronConnection {
             user_id: Arc::new(RwLock::new(0)),
             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())),
+            waiting: Arc::new(DashMap::new()),
             pingpong: Arc::new(Mutex::new(None)),
             last_ping: Arc::new(Mutex::new(-1)),
             message_send_times: Arc::new(Mutex::new(HashMap::new())),
@@ -85,54 +92,103 @@ impl OmikronConnection {
         *self.is_connected.lock().await
     }
     /// Connect loop with retry
-    pub async fn connect(self: &Arc<Self>) {
-        loop {
-            if *SHUTDOWN.read().await {
-                break;
+    pub async fn connect(self: &Arc<Self>, user_ids: String) {
+        if self.is_connected().await {
+            return;
+        }
+
+        let conf = CONFIG.read().await;
+        let iota_id = conf.get_iota_id();
+        let public_key = conf.get_public_key();
+        let private_key = conf.get_private_key();
+
+        if iota_id == 0 || public_key.is_none() || private_key.is_none() {
+            // Registration flow
+            drop(conf); // release read lock
+            log_message_trans("iota_register_new");
+            let key_pair = crypto_helper::generate_keypair();
+            let public_key_base64 = crypto_helper::public_key_to_base64(&key_pair.public);
+            let private_key_base64 = crypto_helper::secret_key_to_base64(&key_pair.secret);
+
+            let mut conf_write = CONFIG.write().await;
+            conf_write.change("public_key", JsonValue::String(public_key_base64.clone()));
+            conf_write.change("private_key", JsonValue::String(private_key_base64));
+            conf_write.update();
+            drop(conf_write);
+
+            if self.connect_internal().await {
+                // a new helper function to just connect
+                self.send_message(
+                    CommunicationValue::new(CommunicationType::register_iota)
+                        .add_data(DataTypes::public_key, JsonValue::String(public_key_base64))
+                        .to_json()
+                        .to_string(),
+                )
+                .await;
             }
-            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 *cloned_self.is_connected.lock().await == false {
-                                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;
-                }
+        } else {
+            // Login flow
+            if self.connect_internal().await {
+                self.send_message(
+                    CommunicationValue::new(CommunicationType::identification)
+                        .add_data(
+                            DataTypes::iota_id,
+                            JsonValue::Number(json::number::Number::from(iota_id)),
+                        )
+                        .to_json()
+                        .to_string(),
+                )
+                .await;
             }
         }
     }
+
+    async fn connect_internal(self: &Arc<Self>) -> bool {
+        if self.is_connected().await {
+            return true;
+        }
+        log_message_trans("omikron_connecting");
+
+        // connect to omikron
+        let conf = CONFIG.read().await;
+        let addr = conf
+            .get("omikron_addr")
+            .as_str()
+            .unwrap_or("wss://app.tensamin.net/ws/iota/");
+        let stream_res = connect_async(addr).await;
+        if let Err(e) = stream_res {
+            log_message(format!("omikron_connection_error {}", e.to_string()));
+            return false;
+        }
+        let (stream, _) = stream_res.unwrap();
+        log_message_trans("omikron_connection_success");
+
+        let (write_half, read_half) = 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.spawn_listener(boxed_reader).await;
+
+        let mut is_connected = self.is_connected.lock().await;
+        *is_connected = true;
+        drop(is_connected);
+
+        let sel_arc_clone = self.clone();
+        tokio::spawn(async move {
+            loop {
+                if !sel_arc_clone.is_connected().await {
+                    break;
+                }
+                sel_arc_clone.send_ping().await;
+                sleep(Duration::from_secs(10)).await;
+            }
+        });
+
+        true
+    }
+
     pub async fn send_message(&self, msg: String) {
         Self::send_message_static(&self.writer, Arc::clone(&self.is_connected), msg).await;
     }
@@ -186,7 +242,7 @@ impl OmikronConnection {
     }
     pub fn handle_message(
         msg: Result<Message, tungstenite::Error>,
-        waiting: Arc<Mutex<HashMap<Uuid, Box<dyn Fn(CommunicationValue) + Send + Sync + 'static>>>>,
+        waiting: Arc<DashMap<Uuid, Box<dyn Fn(CommunicationValue) + Send + Sync + 'static>>>,
         writer: Arc<
             Mutex<
                 Option<Box<dyn Sink<Message, Error = tungstenite::Error> + Send + Unpin + 'static>>,
@@ -205,11 +261,86 @@ impl OmikronConnection {
                     return;
                 }
                 Ok(Message::Text(text)) => {
-                    let mut cv = CommunicationValue::from_json(&text);
+                    let cv = CommunicationValue::from_json(&text);
                     if cv.is_type(CommunicationType::pong) {
                         sel.handle_pong(&cv, true).await;
                         return;
                     }
+                    if cv.is_type(CommunicationType::challenge) {
+                        let conf = CONFIG.read().await;
+                        let private_key = conf.get_private_key().unwrap();
+                        drop(conf);
+
+                        let omikron_public_key = cv
+                            .get_data(DataTypes::public_key)
+                            .unwrap()
+                            .as_str()
+                            .unwrap();
+                        let encrypted_challenge =
+                            cv.get_data(DataTypes::challenge).unwrap().as_str().unwrap();
+
+                        let solved_challenge = crypto_helper::decrypt(
+                            &private_key,
+                            omikron_public_key,
+                            encrypted_challenge,
+                        );
+
+                        if let Ok(decrypted) = solved_challenge {
+                            let response =
+                                CommunicationValue::new(CommunicationType::challenge_response)
+                                    .with_id(cv.get_id())
+                                    .add_data(DataTypes::challenge, JsonValue::String(decrypted));
+
+                            sel_arc.send_message(response.to_json().to_string()).await;
+                        } else {
+                            log_message("Failed to decrypt challenge");
+                        }
+
+                        return;
+                    }
+                    if cv.is_type(CommunicationType::register_iota_success) {
+                        let iota_id = cv
+                            .get_data(DataTypes::iota_id)
+                            .unwrap()
+                            .as_i64()
+                            .unwrap_or(0);
+                        if iota_id != 0 {
+                            let mut conf = CONFIG.write().await;
+                            conf.change("iota_id", JsonValue::Number(iota_id.into()));
+                            conf.update();
+                            log_message(format!("Iota registered with ID: {}", iota_id));
+
+                            // Now, proceed to login
+                            let login_message =
+                                CommunicationValue::new(CommunicationType::identification)
+                                    .add_data(
+                                        DataTypes::iota_id,
+                                        JsonValue::Number(json::number::Number::from(iota_id)),
+                                    )
+                                    .to_json()
+                                    .to_string();
+
+                            let sel_arc_clone = sel_arc.clone();
+                            tokio::spawn(async move {
+                                sel_arc_clone.send_message(login_message).await;
+                            });
+                        } else {
+                            log_message("Iota registration failed.");
+                        }
+                        return;
+                    }
+                    if cv.is_type(CommunicationType::identification_response) {
+                        if let Some(accepted) = cv.get_data(DataTypes::accepted) {
+                            if accepted.as_str().unwrap_or("0") != "0" {
+                                log_message(format!("Omikron connected: {}", accepted.to_string()));
+                            } else {
+                                log_message("omikron_connection_failed");
+                                // Maybe add a retry logic here or close the app
+                            }
+                        }
+                        return;
+                    }
+
                     let com = variant.read().await.clone();
                     if com == ConnectionVariant::ClientUnauthenticated {
                         if cv.is_type(CommunicationType::identification) {
@@ -288,8 +419,8 @@ impl OmikronConnection {
                     // Direct messages                                  //
                     // ************************************************ //
                     log_cv(&cv);
-                    if let Some(x) = waiting.lock().await.remove(&cv.get_id()) {
-                        x(cv);
+                    if let Some((_, y)) = waiting.remove(&cv.get_id()) {
+                        y(cv);
                         return;
                     }
                     if cv.is_type(CommunicationType::message_other_iota) {
@@ -433,7 +564,7 @@ impl OmikronConnection {
 
                     if cv.is_type(CommunicationType::get_chats) {
                         let user_id = cv.get_sender();
-                        let users = get_users(user_id);
+                        let users = chats_util::get_users(user_id);
                         let resp = CommunicationValue::new(CommunicationType::get_chats)
                             .with_id(cv.get_id())
                             .with_receiver(user_id)
@@ -641,4 +772,41 @@ impl OmikronConnection {
             *connected.lock().await = false;
         }
     }
+    pub async fn await_response(
+        self: Arc<Self>,
+        cv: &CommunicationValue,
+        timeout_duration: Option<Duration>,
+    ) -> Result<CommunicationValue, String> {
+        let (tx, mut rx) = mpsc::channel(1);
+        let msg_id = cv.get_id();
+
+        let task_tx = tx.clone();
+        self.waiting.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_message(format!("Failed to send response back to awaiter: {}", &e));
+                    }
+                });
+            }),
+        );
+
+        self.send_message(cv.to_json().to_string()).await;
+
+        let timeout = timeout_duration.unwrap_or(Duration::from_secs(10));
+
+        match tokio::time::timeout(timeout, rx.recv()).await {
+            Ok(Some(response_cv)) => Ok(response_cv),
+            Ok(None) => Err("Failed to receive response, channel was closed.".to_string()),
+            Err(_) => {
+                self.waiting.remove(&msg_id);
+                Err(format!(
+                    "Request timed out after {} seconds.",
+                    timeout.as_secs()
+                ))
+            }
+        }
+    }
 }