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
}
}