diff --git a/src/server/omikron_connection.rs b/src/server/omikron_connection.rs index f1634fc..11dc310 100644 --- a/src/server/omikron_connection.rs +++ b/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 = - user_online_tracker::untrack_iota_connection(iota_id, omikron_id); - if iota_offline { - if let Ok(users) = sql::get_users_by_iota_id(iota_id).await { - let user_ids: Vec = users.iter().map(|u| u.0).collect(); - user_online_tracker::untrack_many_users(&user_ids); + + if let Some(v) = cv.get_data(DataTypes::iota_id) { + if let Some(iota_id) = v.as_i64() { + let iota_offline = + user_online_tracker::untrack_iota_connection(iota_id, omikron_id); + if iota_offline { + if let Ok(users) = sql::get_users_by_iota_id(iota_id).await { + let user_ids: Vec = 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() diff --git a/src/sql/sql.rs b/src/sql/sql.rs index 1994551..8445e26 100644 --- a/src/sql/sql.rs +++ b/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 )", diff --git a/src/sql/user_online_tracker.rs b/src/sql/user_online_tracker.rs index 8cde17b..055c730 100644 --- a/src/sql/user_online_tracker.rs +++ b/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) { - connections.retain(|&id| id != omikron_id); + let connections_empty = if let Some(r) = IOTA_OMIKRON_CONNECTIONS.get(&iota_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 - .get(&iota_id) - .map(|p| *p == omikron_id) - .unwrap_or(false) - { + if let Some(primary_ref) = IOTA_PRIMARY_OMIKRON_CONNECTION.get(&iota_id) { + let primary_id = *primary_ref.value(); + drop(primary_ref); + 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 { @@ -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| { - connections.retain(|id| *id != omikron_id); + // Collect iotas and primary info first + 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() { - offline_iotas.push(*iota_id); - false - } else { - true - } - }); + // Update the connections vector after filtering + IOTA_OMIKRON_CONNECTIONS.insert(iota_id, connections); + } + // 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); } }