From d2a4d0c7ae68150512f815b13568ec202d6ab6f0 Mon Sep 17 00:00:00 2001 From: Alex Emmet <111742636+Alex-Emmet@users.noreply.github.com> Date: Thu, 29 Jan 2026 10:58:43 +0100 Subject: [PATCH] [Add] Connection Tracking --- src/server/api.rs | 49 +++++++++++------ src/server/omikron_connection.rs | 43 +++++++-------- src/server/omikron_manager.rs | 34 ++++++++++++ src/sql/iota_omikron_tracker.rs | 27 --------- src/sql/mod.rs | 1 - src/sql/sql.rs | 33 ----------- src/sql/user_online_tracker.rs | 94 +++++++++++++++++++------------- 7 files changed, 143 insertions(+), 138 deletions(-) delete mode 100644 src/sql/iota_omikron_tracker.rs diff --git a/src/server/api.rs b/src/server/api.rs index 9d77b5b..a2e40ec 100644 --- a/src/server/api.rs +++ b/src/server/api.rs @@ -1,11 +1,12 @@ use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes}; use crate::get_public_key; +use crate::server::omikron_manager::get_random_omikron; use crate::sql::sql; +use crate::sql::user_online_tracker::{ + get_iota_omikron_connections, get_iota_primary_omikron_connection, +}; use crate::{ - sql::{ - iota_omikron_tracker::get_omikron_for_iota, - sql::{get_by_user_id, get_omikron_by_id, get_random_omikron}, - }, + sql::sql::{get_by_user_id, get_omikron_by_id}, util::crypto_helper::public_key_to_base64, }; use axum::http::HeaderValue; @@ -43,17 +44,33 @@ pub async fn handle( // api/get/omikron/ -> omikron for id (user / iota / omikron) "omikron" => { if path_parts.len() == 3 { - if let Ok((id, public_key, ip_address)) = get_random_omikron().await { - ( - StatusCode::OK, - "application/json", - format!( - "{{\"id\": {}, \"public_key\": \"{}\", \"ip_address\": \"{}\"}}", - id, public_key, ip_address - ), - ) + if let Ok(omikron_conn) = get_random_omikron().await { + if let Ok((public_key, ip_address)) = + sql::get_omikron_by_id(omikron_conn.get_omikron_id().await).await + { + ( + StatusCode::OK, + "application/json", + format!( + "{{\"id\": {}, \"public_key\": \"{}\", \"ip_address\": \"{}\"}}", + omikron_conn.get_omikron_id().await, + public_key, + ip_address + ), + ) + } else { + ( + StatusCode::INTERNAL_SERVER_ERROR, + "text/plain", + "selected an invalid omikron".to_string(), + ) + } } else { - not_found() + ( + StatusCode::NOT_FOUND, + "text/plain", + "couldn't find online omikron".to_string(), + ) } } else if path_parts.len() == 4 { let id = path_parts[3].parse::().unwrap_or(0); @@ -68,7 +85,7 @@ pub async fn handle( id, public_key, ip_address ), ) - } else if let Some(omikron_id) = get_omikron_for_iota(id).await { + } else if let Some(omikron_id) = get_iota_primary_omikron_connection(id) { if let Ok((public_key, ip_address)) = get_omikron_by_id(omikron_id).await { @@ -86,7 +103,7 @@ pub async fn handle( } else if let Ok((_, iota_id, _, _, _, _, _, _, _, _, _, _)) = get_by_user_id(id).await { - if let Some(omikron_id) = get_omikron_for_iota(iota_id).await { + if let Some(omikron_id) = get_iota_primary_omikron_connection(iota_id) { if let Ok((public_key, ip_address)) = get_omikron_by_id(omikron_id).await { diff --git a/src/server/omikron_connection.rs b/src/server/omikron_connection.rs index 011389a..5c3d29c 100644 --- a/src/server/omikron_connection.rs +++ b/src/server/omikron_connection.rs @@ -1,4 +1,5 @@ use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes}; +use crate::server::omikron_manager; use crate::server::short_link::add_short_link; use crate::sql::connection_status::ConnectionType; use crate::sql::sql::{self, get_by_user_id, get_by_username, get_iota_by_id, get_omikron_by_id}; @@ -11,6 +12,7 @@ use dashmap::DashMap; use futures::SinkExt; use futures::stream::SplitSink; use futures::stream::SplitStream; +use futures::task::UnsafeFutureObj; use hyper::upgrade::Upgraded; use hyper_util::rt::TokioIo; use json::JsonValue; @@ -183,8 +185,7 @@ impl OmikronConnection { if client_response == *self.challenge.read().await { *self.challenged.write().await = true; - let _ = sql::set_omikron_active(self.get_omikron_id().await, true); - + omikron_manager::add_omikron(self.clone()).await; self.send_message( &CommunicationValue::new(CommunicationType::identification_response) .with_id(cv.get_id()) @@ -225,20 +226,18 @@ impl OmikronConnection { // ONLINE STATUS TRACKING if cv.is_type(CommunicationType::user_connected) { if let Some(user_id) = cv.get_data(DataTypes::user_id).and_then(|v| v.as_i64()) { - user_online_tracker::track_user_status(user_id, ConnectionType::Online, omikron_id) - .await; + user_online_tracker::track_user_status(user_id, ConnectionType::Online, omikron_id); } return; } if cv.is_type(CommunicationType::user_disconnected) { if let Some(user_id) = cv.get_data(DataTypes::user_id).and_then(|v| v.as_i64()) { - if let Some(status) = user_online_tracker::get_user_status(user_id).await { + if let Some(status) = user_online_tracker::get_user_status(user_id) { user_online_tracker::track_user_status( user_id, ConnectionType::UserOffline, status.omikron_id, - ) - .await; + ); } } return; @@ -247,7 +246,7 @@ impl OmikronConnection { if cv.is_type(CommunicationType::iota_connected) { log_in!(PrintType::Omega, "IOTA connected"); if let Some(iota_id) = cv.get_data(DataTypes::iota_id).and_then(|v| v.as_i64()) { - user_online_tracker::track_iota_connection(iota_id, omikron_id).await; + user_online_tracker::track_iota_connection(iota_id, omikron_id, true); let mut user_ids = JsonValue::Array(Vec::new()); if let Ok(users) = sql::get_users_by_iota_id(iota_id).await { for user in users { @@ -256,8 +255,7 @@ impl OmikronConnection { user.0, ConnectionType::UserOffline, omikron_id, - ) - .await; + ); } } else { log_in!(PrintType::General, "SQL error, when loading users for IOTA"); @@ -276,11 +274,11 @@ impl OmikronConnection { if cv.is_type(CommunicationType::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).await; + 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).await; + user_online_tracker::untrack_many_users(&user_ids); } } } @@ -296,8 +294,7 @@ impl OmikronConnection { user_id, ConnectionType::Online, omikron_id, - ) - .await; + ); } } } @@ -306,7 +303,7 @@ impl OmikronConnection { { for iota_id_json in iota_ids { if let Some(iota_id) = iota_id_json.as_i64() { - user_online_tracker::track_iota_connection(iota_id, omikron_id).await; + user_online_tracker::track_iota_connection(iota_id, omikron_id, true); } } } @@ -365,10 +362,9 @@ impl OmikronConnection { response.add_data_str(DataTypes::avatar, STANDARD.encode(avatar)); } - let user_status = user_online_tracker::get_user_status(id).await; + let user_status = user_online_tracker::get_user_status(id); let iota_connections = user_online_tracker::get_iota_omikron_connections(iota_id) - .await .unwrap_or_default(); response = response.add_data( DataTypes::omikron_connections, @@ -451,10 +447,9 @@ impl OmikronConnection { response.add_data_str(DataTypes::avatar, STANDARD.encode(avatar)); } - let user_status = user_online_tracker::get_user_status(id).await; + let user_status = user_online_tracker::get_user_status(id); let iota_connections = user_online_tracker::get_iota_omikron_connections(iota_id) - .await .unwrap_or_default(); if let Some(user_status) = user_status { @@ -506,7 +501,6 @@ impl OmikronConnection { let iota_connections = user_online_tracker::get_iota_omikron_connections(iota_id) - .await .unwrap_or_default(); response = response.add_data( @@ -544,7 +538,6 @@ impl OmikronConnection { let iota_connections = user_online_tracker::get_iota_omikron_connections(iota_id) - .await .unwrap_or_default(); response = response.add_data( @@ -585,7 +578,6 @@ impl OmikronConnection { let iota_connections = user_online_tracker::get_iota_omikron_connections(iota_id) - .await .unwrap_or_default(); response = response.add_data( @@ -904,7 +896,11 @@ impl OmikronConnection { pub async fn close(&self) { let mut sender = self.sender.write().await; if self.is_identified().await { - let _ = sql::set_omikron_active(self.get_omikron_id().await, false); + let omikron_id = self.get_omikron_id().await; + if omikron_id != 0 { + omikron_manager::remove_omikron(omikron_id).await; + user_online_tracker::untrack_omikron(omikron_id).await; + } } let _ = sender.close().await; } @@ -912,6 +908,7 @@ impl OmikronConnection { if self.is_identified().await { let omikron_id = self.get_omikron_id().await; if omikron_id != 0 { + omikron_manager::remove_omikron(omikron_id).await; user_online_tracker::untrack_omikron(omikron_id).await; } } diff --git a/src/server/omikron_manager.rs b/src/server/omikron_manager.rs index e69de29..8440835 100644 --- a/src/server/omikron_manager.rs +++ b/src/server/omikron_manager.rs @@ -0,0 +1,34 @@ +use crate::{log_in, server::omikron_connection::OmikronConnection, util::logger::PrintType}; +use dashmap::DashMap; +use once_cell::sync::Lazy; +use rand::prelude::IteratorRandom; +use std::sync::Arc; + +pub static OMIKRON_CONNECTIONS: Lazy>> = + Lazy::new(|| DashMap::new()); + +pub async fn add_omikron(omikron_conn: Arc) { + OMIKRON_CONNECTIONS.insert(omikron_conn.get_omikron_id().await, omikron_conn); +} + +pub async fn remove_omikron(omikron_id: i64) { + OMIKRON_CONNECTIONS.remove(&omikron_id); +} + +pub async fn get_random_omikron() -> Result, ()> { + log_in!(0, PrintType::Iota, "{}", OMIKRON_CONNECTIONS.len()); + let mut rng = rand::thread_rng(); + if let Some((_, val)) = OMIKRON_CONNECTIONS.clone().into_iter().choose(&mut rng) { + return Ok(val); + } else { + return Err(()); + } +} + +pub async fn get_omikron(omikron_id: i64) -> Option> { + if let Some(omikron) = OMIKRON_CONNECTIONS.get(&omikron_id) { + Some(omikron.clone()) + } else { + None + } +} diff --git a/src/sql/iota_omikron_tracker.rs b/src/sql/iota_omikron_tracker.rs deleted file mode 100644 index cf11310..0000000 --- a/src/sql/iota_omikron_tracker.rs +++ /dev/null @@ -1,27 +0,0 @@ -use once_cell::sync::Lazy; -use std::collections::HashMap; -use std::sync::Arc; -use tokio::sync::RwLock; - -static IOTA_OMIKRON_MAP: Lazy>>> = - Lazy::new(|| Arc::new(RwLock::new(HashMap::new()))); - -pub async fn track_iota_omikron(iota: i64, omikron: i64) { - let mut c = IOTA_OMIKRON_MAP.write().await; - c.insert(iota, omikron); -} - -pub async fn get_omikron_for_iota(iota: i64) -> Option { - let c = IOTA_OMIKRON_MAP.read().await; - c.get(&iota).cloned() -} - -pub async fn untrack_iota(iota: i64) { - let mut c = IOTA_OMIKRON_MAP.write().await; - c.remove(&iota); -} - -pub async fn untrack_by_omikron(omikron: i64) { - let mut c = IOTA_OMIKRON_MAP.write().await; - c.retain(|_, v| *v != omikron); -} diff --git a/src/sql/mod.rs b/src/sql/mod.rs index 033fd01..91ba841 100644 --- a/src/sql/mod.rs +++ b/src/sql/mod.rs @@ -1,4 +1,3 @@ pub mod connection_status; -pub mod iota_omikron_tracker; pub mod sql; pub mod user_online_tracker; diff --git a/src/sql/sql.rs b/src/sql/sql.rs index ec0f1d2..20dd688 100644 --- a/src/sql/sql.rs +++ b/src/sql/sql.rs @@ -87,7 +87,6 @@ pub async fn initialize_db() -> Result<(), sqlx::Error> { "CREATE TABLE IF NOT EXISTS omikrons ( id BIGINT UNSIGNED NOT NULL PRIMARY KEY, - is_active INT(1) NOT NULL DEFAULT 0, public_key VARCHAR(255) NOT NULL COLLATE utf8mb4_bin, location VARCHAR(255) NOT NULL COLLATE utf8mb4_bin, ip_address VARCHAR(255) NOT NULL COLLATE utf8mb4_bin @@ -589,23 +588,6 @@ pub async fn delete_iota(id: i64) -> Result<(), sqlx::Error> { // OMIKRONS // ========================================================================================== -pub async fn get_random_omikron() -> Result<(i64, String, String), sqlx::Error> { - let db_lock = SQL_DB.read().await; - let pool = db_lock.as_ref().expect("Database pool is not initialized"); - let row = sqlx::query_as::<_, (i64, Vec, Vec)>("SELECT id, public_key, ip_address FROM omikrons WHERE is_active = 1 ORDER BY RAND() LIMIT 1") - .fetch_optional(pool) - .await?; - - match row { - Some((id, public_key, ip_address)) => Ok(( - id, - String::from_utf8_lossy(&public_key).to_string(), - String::from_utf8_lossy(&ip_address).to_string(), - )), - _ => Err(sqlx::Error::RowNotFound), - } -} - pub async fn get_omikron_by_id(id: i64) -> Result<(String, String), sqlx::Error> { let db_lock = SQL_DB.read().await; let pool = db_lock.as_ref().expect("Database pool is not initialized"); @@ -625,18 +607,3 @@ pub async fn get_omikron_by_id(id: i64) -> Result<(String, String), sqlx::Error> _ => Err(sqlx::Error::RowNotFound), } } - -pub async fn set_omikron_active(id: i64, active: bool) -> Result<(), sqlx::Error> { - let db_lock = SQL_DB.read().await; - let pool = db_lock.as_ref().expect("Database pool is not initialized"); - - let active = if active { 1 } else { 0 }; - - sqlx::query("UPDATE omikrons SET active = ? WHERE id = CAST(? AS UNSIGNED)") - .bind(active) - .bind(id) - .execute(pool) - .await?; - - Ok(()) -} diff --git a/src/sql/user_online_tracker.rs b/src/sql/user_online_tracker.rs index 95e94c6..3a1ab42 100644 --- a/src/sql/user_online_tracker.rs +++ b/src/sql/user_online_tracker.rs @@ -1,9 +1,7 @@ use crate::sql; use crate::sql::connection_status::ConnectionType; +use dashmap::DashMap; use once_cell::sync::Lazy; -use std::collections::HashMap; -use std::sync::Arc; -use tokio::sync::RwLock; #[derive(Debug, Clone)] pub struct UserStatus { @@ -11,42 +9,59 @@ pub struct UserStatus { pub omikron_id: i64, } +// IotaID -> Primary OmikronID +static IOTA_PRIMARY_OMIKRON_CONNECTION: Lazy> = Lazy::new(DashMap::new); + // IotaID -> Vec -static IOTA_OMIKRON_CONNECTIONS: Lazy>>>> = - Lazy::new(|| Arc::new(RwLock::new(HashMap::new()))); +static IOTA_OMIKRON_CONNECTIONS: Lazy>> = Lazy::new(DashMap::new); // UserID -> UserStatus -static USER_STATUS_MAP: Lazy>>> = - Lazy::new(|| Arc::new(RwLock::new(HashMap::new()))); +static USER_STATUS_MAP: Lazy> = Lazy::new(DashMap::new); -pub async fn track_iota_connection(iota_id: i64, omikron_id: i64) { - let mut iota_map = IOTA_OMIKRON_CONNECTIONS.write().await; - let connections = iota_map.entry(iota_id).or_default(); - if !connections.contains(&omikron_id) { - connections.push(omikron_id); +pub fn track_iota_connection(iota_id: i64, omikron_id: i64, primary: bool) { + let mut entry = IOTA_OMIKRON_CONNECTIONS + .entry(iota_id) + .or_insert_with(Vec::new); + + if !entry.contains(&omikron_id) { + entry.push(omikron_id); + } + + if primary { + IOTA_PRIMARY_OMIKRON_CONNECTION.insert(iota_id, omikron_id); } } -pub async fn untrack_iota_connection(iota_id: i64, omikron_id: i64) -> bool { - let mut iota_map = IOTA_OMIKRON_CONNECTIONS.write().await; - if let Some(connections) = iota_map.get_mut(&iota_id) { +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); + + if IOTA_PRIMARY_OMIKRON_CONNECTION + .get(&iota_id) + .map(|p| *p == omikron_id) + .unwrap_or(false) + { + IOTA_PRIMARY_OMIKRON_CONNECTION.remove(&iota_id); + } + if connections.is_empty() { - iota_map.remove(&iota_id); - return true; // Iota is now offline + IOTA_OMIKRON_CONNECTIONS.remove(&iota_id); + return true; } } false } -pub async fn get_iota_omikron_connections(iota_id: i64) -> Option> { - let iota_map = IOTA_OMIKRON_CONNECTIONS.read().await; - iota_map.get(&iota_id).cloned() +pub fn get_iota_primary_omikron_connection(iota_id: i64) -> Option { + IOTA_PRIMARY_OMIKRON_CONNECTION.get(&iota_id).map(|v| *v) } -pub async fn track_user_status(user_id: i64, status: ConnectionType, omikron_id: i64) { - let mut user_map = USER_STATUS_MAP.write().await; - user_map.insert( +pub fn get_iota_omikron_connections(iota_id: i64) -> Option> { + IOTA_OMIKRON_CONNECTIONS.get(&iota_id).map(|v| v.clone()) +} + +pub fn track_user_status(user_id: i64, status: ConnectionType, omikron_id: i64) { + USER_STATUS_MAP.insert( user_id, UserStatus { connection_type: status, @@ -55,31 +70,34 @@ pub async fn track_user_status(user_id: i64, status: ConnectionType, omikron_id: ); } -pub async fn get_user_status(user_id: i64) -> Option { - let user_map = USER_STATUS_MAP.read().await; - user_map.get(&user_id).cloned() +pub fn get_user_status(user_id: i64) -> Option { + USER_STATUS_MAP.get(&user_id).map(|v| v.clone()) } -pub async fn untrack_user(user_id: i64) { - let mut user_map = USER_STATUS_MAP.write().await; - user_map.remove(&user_id); +pub fn untrack_user(user_id: i64) { + USER_STATUS_MAP.remove(&user_id); } -pub async fn untrack_many_users(user_ids: &[i64]) { - let mut user_map = USER_STATUS_MAP.write().await; +pub fn untrack_many_users(user_ids: &[i64]) { for user_id in user_ids { - user_map.remove(user_id); + USER_STATUS_MAP.remove(user_id); } } pub async fn untrack_omikron(omikron_id: i64) { - let mut iota_map = IOTA_OMIKRON_CONNECTIONS.write().await; - let mut user_map = USER_STATUS_MAP.write().await; - let mut offline_iotas = Vec::new(); - iota_map.retain(|iota_id, connections| { + IOTA_OMIKRON_CONNECTIONS.retain(|iota_id, connections| { connections.retain(|id| *id != omikron_id); + + if IOTA_PRIMARY_OMIKRON_CONNECTION + .get(iota_id) + .map(|p| *p == omikron_id) + .unwrap_or(false) + { + IOTA_PRIMARY_OMIKRON_CONNECTION.remove(iota_id); + } + if connections.is_empty() { offline_iotas.push(*iota_id); false @@ -88,12 +106,12 @@ pub async fn untrack_omikron(omikron_id: i64) { } }); - user_map.retain(|_, status| status.omikron_id != omikron_id); + USER_STATUS_MAP.retain(|_, status| status.omikron_id != omikron_id); 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_map.remove(&user.0); + USER_STATUS_MAP.remove(&user.0); } } }