[Fix] Omikron Con no longer breaks on iota disconnect :D

src/server/omikron_connection.rs

@@ -63,7 +63,11 @@ impl OmikronConnection {
     }
 
     pub async fn get_omikron_id(&self) -> i64 {
-        *self.omikron_id.read().await
+        let omikron_id = {
+            let guard = self.omikron_id.read().await;
+            guard.clone()
+        };
+        omikron_id
     }
     pub async fn is_identified(&self) -> bool {
         *self.identified.read().await && *self.challenged.read().await
@@ -267,18 +271,23 @@ impl OmikronConnection {
         }
         if cv.is_type(CommunicationType::iota_disconnected) {
             log_in!(PrintType::Omega, "IOTA disconnected");
-            if let Some(iota_id) = cv.get_data(DataTypes::iota_id).and_then(|v| v.as_i64()) {
+
-                let iota_offline =
+            if let Some(v) = cv.get_data(DataTypes::iota_id) {
-                    user_online_tracker::untrack_iota_connection(iota_id, omikron_id);
+                if let Some(iota_id) = v.as_i64() {
-                if iota_offline {
+                    let iota_offline =
-                    if let Ok(users) = sql::get_users_by_iota_id(iota_id).await {
+                        user_online_tracker::untrack_iota_connection(iota_id, omikron_id);
-                        let user_ids: Vec<i64> = users.iter().map(|u| u.0).collect();
+                    if iota_offline {
-                        user_online_tracker::untrack_many_users(&user_ids);
+                        if let Ok(users) = sql::get_users_by_iota_id(iota_id).await {
+                            let user_ids: Vec<i64> = users.iter().map(|u| u.0).collect();
+                            user_online_tracker::untrack_many_users(&user_ids);
+                        }
                     }
                 }
             }
+
             return;
         }
+
         if cv.is_type(CommunicationType::sync_client_iota_status) {
             if let Some(json::JsonValue::Array(user_ids)) =
                 cv.get_data(DataTypes::user_ids).cloned()

src/sql/sql.rs

@@ -65,9 +65,9 @@ pub async fn initialize_db() -> Result<(), sqlx::Error> {
         about VARCHAR(200) COLLATE utf8mb4_bin,
         avatar MEDIUMBLOB,
         sub_level INT(11) NOT NULL DEFAULT 0,
-        sub_end BIGINT(20) NOT NULL,
+        sub_end BIGINT(20) NOT NULL DEFAULT 0,
         public_key TEXT NOT NULL COLLATE utf8mb4_bin,
-        private_key_hash TEXT NOT NULL COLLATE utf8mb4_bin,
+        private_key_hash TEXT NOT NULL COLLATE utf8mb4_bin DEFAULT '',
         iota_id BIGINT UNSIGNED NOT NULL,
         token VARCHAR(255) NOT NULL UNIQUE COLLATE utf8mb4_bin
         )",

src/sql/user_online_tracker.rs

@@ -33,23 +33,30 @@ pub fn track_iota_connection(iota_id: i64, omikron_id: i64, primary: bool) {
 }
 
 pub fn untrack_iota_connection(iota_id: i64, omikron_id: i64) -> bool {
-    if let Some(mut connections) = IOTA_OMIKRON_CONNECTIONS.get_mut(&iota_id) {
+    let connections_empty = if let Some(r) = IOTA_OMIKRON_CONNECTIONS.get(&iota_id) {
-        connections.retain(|&id| id != omikron_id);
+        let mut vec = r.value().clone();
+        vec.retain(|&id| id != omikron_id);
+        let empty = vec.is_empty();
+        drop(r);
+        IOTA_OMIKRON_CONNECTIONS.insert(iota_id, vec);
+        empty
+    } else {
+        false
+    };
 
-        if IOTA_PRIMARY_OMIKRON_CONNECTION
+    if let Some(primary_ref) = IOTA_PRIMARY_OMIKRON_CONNECTION.get(&iota_id) {
-            .get(&iota_id)
+        let primary_id = *primary_ref.value();
-            .map(|p| *p == omikron_id)
+        drop(primary_ref);
-            .unwrap_or(false)
+        if primary_id == omikron_id {
-        {
             IOTA_PRIMARY_OMIKRON_CONNECTION.remove(&iota_id);
         }
-
-        if connections.is_empty() {
-            IOTA_OMIKRON_CONNECTIONS.remove(&iota_id);
-            return true;
-        }
     }
-    false
+
+    if connections_empty {
+        IOTA_OMIKRON_CONNECTIONS.remove(&iota_id);
+    }
+
+    connections_empty
 }
 
 pub fn get_iota_primary_omikron_connection(iota_id: i64) -> Option<i64> {
@@ -81,34 +88,48 @@ pub fn untrack_many_users(user_ids: &[i64]) {
 }
 
 pub async fn untrack_omikron(omikron_id: i64) {
+    // Step 1: Collect all iota_ids where this omikron_id is connected
     let mut offline_iotas = Vec::new();
+    let mut primary_to_remove = Vec::new();
 
-    IOTA_OMIKRON_CONNECTIONS.retain(|iota_id, connections| {
+    // Collect iotas and primary info first
-        connections.retain(|id| *id != omikron_id);
+    for r in IOTA_OMIKRON_CONNECTIONS.iter() {
+        let iota_id = *r.key();
+        let mut connections = r.value().clone();
+        connections.retain(|&id| id != omikron_id);
+
+        if connections.is_empty() {
+            offline_iotas.push(iota_id);
+        }
 
         if IOTA_PRIMARY_OMIKRON_CONNECTION
-            .get(iota_id)
+            .get(&iota_id)
             .map(|p| *p == omikron_id)
             .unwrap_or(false)
         {
-            IOTA_PRIMARY_OMIKRON_CONNECTION.remove(iota_id);
+            primary_to_remove.push(iota_id);
         }
 
-        if connections.is_empty() {
+        // Update the connections vector after filtering
-            offline_iotas.push(*iota_id);
+        IOTA_OMIKRON_CONNECTIONS.insert(iota_id, connections);
-            false
+    }
-        } else {
-            true
-        }
-    });
 
+    // Step 2: Remove primary connections safely
+    for iota_id in primary_to_remove {
+        IOTA_PRIMARY_OMIKRON_CONNECTION.remove(&iota_id);
+    }
+
+    // Step 3: Remove users that were on this omikron
     USER_STATUS_MAP.retain(|_, status| status.omikron_id != omikron_id);
 
+    // Step 4: For offline iotas, remove associated users from USER_STATUS_MAP
     for iota_id in offline_iotas {
         if let Ok(users) = sql::sql::get_users_by_iota_id(iota_id).await {
             for user in users {
                 USER_STATUS_MAP.remove(&user.0);
             }
         }
+        // Finally remove the empty connections vector
+        IOTA_OMIKRON_CONNECTIONS.remove(&iota_id);
     }
 }