From 10b5c05de4833f549e0432991724f39904e68397 Mon Sep 17 00:00:00 2001
From: Alex Emmet <111742636+Alex-Emmet@users.noreply.github.com>
Date: Wed, 14 Jan 2026 22:20:43 +0100
Subject: [PATCH] Stable Connections

---
 src/omega/omega_connection.rs | 62 ++++++++++++++++++++++-------------
 src/rho/iota_connection.rs    |  7 ++++
 2 files changed, 46 insertions(+), 23 deletions(-)

diff --git a/src/omega/omega_connection.rs b/src/omega/omega_connection.rs
index ff4f6a8..d4cad9d 100755
--- a/src/omega/omega_connection.rs
+++ b/src/omega/omega_connection.rs
@@ -8,7 +8,15 @@ use dashmap::DashMap;
 use futures::prelude::*;
 use json::{JsonValue, number::Number};
 use once_cell::sync::Lazy;
-use std::{collections::HashMap, env, sync::Arc, time::Duration};
+use std::{
+    collections::HashMap,
+    env,
+    sync::{
+        Arc,
+        atomic::{AtomicBool, Ordering},
+    },
+    time::Duration,
+};
 use tokio::{
     net::TcpStream,
     sync::{Mutex, RwLock, mpsc},
@@ -17,7 +25,6 @@ use tokio::{
 use tokio_native_tls::TlsStream;
 use uuid::Uuid;
 
-use crate::log_err;
 use crate::{
     data::{
         communication::{CommunicationType, CommunicationValue, DataTypes},
@@ -28,17 +35,27 @@ use crate::{
     util::crypto_helper::{decrypt, load_public_key},
     util::logger::PrintType,
 };
+use crate::{log_err, util::crypto_helper::secret_key_to_base64};
 
 pub static WAITING_TASKS: Lazy<
     DashMap<Uuid, Box<dyn Fn(Arc<OmegaConnection>, CommunicationValue) -> bool + Send + Sync>>,
 > = Lazy::new(DashMap::new);
 
+/// Flag to ensure the connection loop is only started once.
+static CONNECTION_LOOP_STARTED: AtomicBool = AtomicBool::new(false);
+
+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;
-    });
+    if !CONNECTION_LOOP_STARTED.swap(true, Ordering::SeqCst) {
+        let conn_clone = conn.clone();
+        tokio::spawn(async move {
+            conn_clone.connect_internal(0).await;
+        });
+    }
     conn
 });
 
@@ -85,10 +102,12 @@ impl OmegaConnection {
         }
     }
     pub fn connect(self: Arc<OmegaConnection>) {
-        let cloned_self = self.clone();
-        tokio::spawn(async move {
-            cloned_self.connect_internal(0).await;
-        });
+        if !CONNECTION_LOOP_STARTED.swap(true, Ordering::SeqCst) {
+            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) {
         loop {
@@ -193,16 +212,6 @@ impl OmegaConnection {
                                                     return false;
                                                 }
 
-                                                if let Some(accepted) = final_cv.get_data(DataTypes::accepted).and_then(|v| v.as_bool()) {
-                                                    if !accepted {
-                                                        log_err!(PrintType::Omega, "Omega did not accept identification.");
-                                                        return false;
-                                                    }
-                                                } else {
-                                                    log_err!(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();
@@ -307,17 +316,24 @@ impl OmegaConnection {
                         continue;
                     }
                     let msg_id = cv.get_id();
-                    log_in!(PrintType::Omega, "{}", &cv.to_json().to_string());
+                    log_in!(PrintType::Omikron, "{}", &cv.to_json().to_string());
                     // Handle waiting tasks
                     if let Some(task) = WAITING_TASKS.remove(&msg_id) {
                         if (task.1)(self.clone(), cv.clone()) {
-                            continue;
+                            // continue in the read_loop
                         }
                     } 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
+                            }
+                        }
                     }
                 }
                 #[allow(non_snake_case)]
-                Some(Ok(Message::Close(_))) | None => continue,
+                Some(Ok(Message::Close(_))) | None => break,
                 Some(Err(_)) => break,
                 _ => {}
             }
diff --git a/src/rho/iota_connection.rs b/src/rho/iota_connection.rs
index dabe675..02b01e0 100644
--- a/src/rho/iota_connection.rs
+++ b/src/rho/iota_connection.rs
@@ -236,6 +236,8 @@ impl IotaConnection {
             } else {
                 self.send_error_response(&cv.get_id(), CommunicationType::error_internal)
                     .await;
+                self.close().await;
+                return;
             }
 
             return;
@@ -294,6 +296,11 @@ impl IotaConnection {
                     .add_data(DataTypes::iota_id, JsonValue::from(new_iota_id));
 
                 iota_conn_for_task.send_message(&success_msg).await;
+            } else {
+                iota_conn_for_task
+                    .send_error_response(&cv.get_id(), CommunicationType::error_internal)
+                    .await;
+                iota_conn_for_task.close().await;
             }
             return;
         }