tensamin/omega
Watch
6
0
Fork
You've already forked omega
 0
Code Issues Pull requests 1 Projects Releases Packages Wiki Activity Actions

[Add] Connection Tracking

Browse source
This commit is contained in:
Alex Emmet 2026-01-29 10:58:43 +01:00
commit d3f2f09b7b
7 changed files with 143 additions and 138 deletions
Download patch file Download diff file Expand all files Collapse all files

49
src/server/api.rs
View file

@ -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/<id> -> 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::<i64>().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
{

43
src/server/omikron_connection.rs
View file

@ -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<i64> = 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;
}
}

34
src/server/omikron_manager.rs
View file

@ -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<DashMap<i64, Arc<OmikronConnection>>> =
Lazy::new(|| DashMap::new());
pub async fn add_omikron(omikron_conn: Arc<OmikronConnection>) {
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<Arc<OmikronConnection>, ()> {
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<Arc<OmikronConnection>> {
if let Some(omikron) = OMIKRON_CONNECTIONS.get(&omikron_id) {
Some(omikron.clone())
} else {
None
}
}

27
src/sql/iota_omikron_tracker.rs
View file

@ -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<Arc<RwLock<HashMap<i64, i64>>>> =
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<i64> {
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);
}

1
src/sql/mod.rs
View file

@ -1,4 +1,3 @@
pub mod connection_status;
pub mod iota_omikron_tracker;
pub mod sql;
pub mod user_online_tracker;

33
src/sql/sql.rs
View file

@ -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<u8>, Vec<u8>)>("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(())
}

94
src/sql/user_online_tracker.rs
View file

@ -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<DashMap<i64, i64>> = Lazy::new(DashMap::new);
// IotaID -> Vec<OmikronID>
static IOTA_OMIKRON_CONNECTIONS: Lazy<Arc<RwLock<HashMap<i64, Vec<i64>>>>> =
Lazy::new(|| Arc::new(RwLock::new(HashMap::new())));
static IOTA_OMIKRON_CONNECTIONS: Lazy<DashMap<i64, Vec<i64>>> = Lazy::new(DashMap::new);
// UserID -> UserStatus
static USER_STATUS_MAP: Lazy<Arc<RwLock<HashMap<i64, UserStatus>>>> =
Lazy::new(|| Arc::new(RwLock::new(HashMap::new())));
static USER_STATUS_MAP: Lazy<DashMap<i64, UserStatus>> = 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<Vec<i64>> {
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<i64> {
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<Vec<i64>> {
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<UserStatus> {
let user_map = USER_STATUS_MAP.read().await;
user_map.get(&user_id).cloned()
pub fn get_user_status(user_id: i64) -> Option<UserStatus> {
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);
}
}
}