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

Format, Status, User Stati & Endpoints

Browse source
This commit is contained in:
Alex Emmet 2026-01-07 15:54:05 +01:00
commit 97f0a76c6c
11 changed files with 963 additions and 396 deletions
Download patch file Download diff file Expand all files Collapse all files

24
src/data/communication.rs
View file

@ -78,6 +78,9 @@ pub enum DataTypes {
communities,
rho_connections,
user,
online_status,
omikron_id,
omikron_connections,
}
impl DataTypes {
@ -155,6 +158,9 @@ impl DataTypes {
"communities" => DataTypes::communities,
"rhoconnections" => DataTypes::rho_connections,
"user" => DataTypes::user,
"onlinestatus" => DataTypes::online_status,
"omikronid" => DataTypes::omikron_id,
"omikronconnections" => DataTypes::omikron_connections,
_ => DataTypes::error_type, // fallback if unknown
}
}
@ -231,6 +237,15 @@ pub enum CommunicationType {
iota_connected,
iota_disconnected,
sync_client_iota_status,
get_user_data,
get_iota_data,
change_user_data,
change_iota_data,
start_register,
complete_register,
}
impl CommunicationType {
pub fn parse(p0: String) -> CommunicationType {
@ -305,6 +320,15 @@ impl CommunicationType {
"userdisconnected" => CommunicationType::user_disconnected,
"syncclientiotastatus" => CommunicationType::sync_client_iota_status,
"getuserdata" => CommunicationType::get_user_data,
"getiotadata" => CommunicationType::get_iota_data,
"changeuserdata" => CommunicationType::change_user_data,
"changeiotadata" => CommunicationType::change_iota_data,
"startregister" => CommunicationType::start_register,
"completeregister" => CommunicationType::complete_register,
_ => CommunicationType::error,
}
}

142
src/server/api.rs
View file

@ -1,13 +1,10 @@
use crate::{get_public_key, log};
use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes};
use crate::get_public_key;
use crate::sql::sql;
use crate::{
sql::{
iota_omikron_tracker::{get_omikron_for_iota, track_iota_omikron, untrack_iota},
sql::{
change_about, change_avatar, change_display_name, change_iota_id, change_iota_key,
change_keys, change_status, change_username, get_by_id, get_by_username,
get_iota_by_id, get_omikron_by_id, get_random_omikron, get_register_id,
register_complete_iota, register_complete_user,
},
iota_omikron_tracker::get_omikron_for_iota,
sql::{get_by_user_id, get_omikron_by_id, get_random_omikron},
},
util::crypto_helper::public_key_to_base64,
};
@ -16,6 +13,7 @@ use http_body_util::Full;
use hyper::body::Bytes;
use hyper::{HeaderMap, Response as HttpResponse, StatusCode};
use json::JsonValue;
use json::number::Number;
pub async fn handle(
path: &str,
@ -37,11 +35,6 @@ pub async fn handle(
// get/
// omikron/
// id/
// register/
// innit/
// complete
log!("{}", path);
log!("{:?} .len = {}", path_parts, path_parts.len());
let (status, content, body_text) = if path_parts.len() >= 2 {
match path_parts[1] {
"get" => match path_parts[2] {
@ -90,7 +83,7 @@ pub async fn handle(
not_found()
}
} else if let Ok((_, iota_id, _, _, _, _, _, _, _, _, _, _)) =
get_by_id(id).await
get_by_user_id(id).await
{
if let Some(omikron_id) = get_omikron_for_iota(iota_id).await {
if let Ok((public_key, ip_address)) =
@ -119,14 +112,131 @@ pub async fn handle(
}
// get/id/<username>
"id" => {
let username = path_parts[3];
bad_request()
let username = path_parts[2];
if username.is_empty() {
not_found()
} else {
if let Ok((
id,
iota_id,
username,
display,
status,
about,
avatar,
sub_level,
sub_end,
public_key,
_,
_,
)) = sql::get_by_username(username).await
{
(
StatusCode::OK,
"application/json",
CommunicationValue::new(CommunicationType::success)
.add_data_str(DataTypes::username, username)
.add_data_str(DataTypes::public_key, public_key)
.add_data(
DataTypes::user_id,
JsonValue::Number(Number::from(id)),
)
.add_data(
DataTypes::iota_id,
JsonValue::Number(Number::from(iota_id)),
)
.add_data_str(DataTypes::display, display)
.add_data_str(DataTypes::status, status)
.add_data_str(DataTypes::about, about)
.add_data_str(DataTypes::avatar, avatar)
.add_data(
DataTypes::sub_level,
JsonValue::Number(Number::from(sub_level)),
)
.add_data(
DataTypes::sub_end,
JsonValue::Number(Number::from(sub_end)),
)
.to_json()
.to_string(),
)
} else {
(
StatusCode::OK,
"application/json",
CommunicationValue::new(CommunicationType::error_not_found)
.to_json()
.to_string(),
)
}
}
}
"public_key" => (
StatusCode::OK,
"application/json",
public_key_to_base64(&get_public_key()),
),
"user" => {
let id = path_parts[2];
let id: i64 = id.parse().unwrap_or(0);
if id == 0 {
bad_request()
} else {
if let Ok((
id,
iota_id,
username,
display,
status,
about,
avatar,
sub_level,
sub_end,
public_key,
_,
_,
)) = sql::get_by_user_id(id).await
{
(
StatusCode::OK,
"application/json",
CommunicationValue::new(CommunicationType::success)
.add_data_str(DataTypes::username, username)
.add_data_str(DataTypes::public_key, public_key)
.add_data(
DataTypes::user_id,
JsonValue::Number(Number::from(id)),
)
.add_data(
DataTypes::iota_id,
JsonValue::Number(Number::from(iota_id)),
)
.add_data_str(DataTypes::display, display)
.add_data_str(DataTypes::status, status)
.add_data_str(DataTypes::about, about)
.add_data_str(DataTypes::avatar, avatar)
.add_data(
DataTypes::sub_level,
JsonValue::Number(Number::from(sub_level)),
)
.add_data(
DataTypes::sub_end,
JsonValue::Number(Number::from(sub_end)),
)
.to_json()
.to_string(),
)
} else {
(
StatusCode::OK,
"application/json",
CommunicationValue::new(CommunicationType::error_not_found)
.to_json()
.to_string(),
)
}
}
}
_ => {
let id = path_parts[2];
let id: i64 = id.parse().unwrap_or(0);

511
src/server/omikron_connection.rs
View file

@ -1,11 +1,7 @@
use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes};
use crate::sql::iota_omikron_tracker::{
track_iota_omikron, untrack_by_omikron as untrack_iota_by_omikron, untrack_iota,
};
use crate::sql::sql::get_omikron_by_id;
use crate::sql::user_online_tracker::{
track_user_omikron, untrack_by_omikron as untrack_user_by_omikron, untrack_user,
};
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};
use crate::sql::user_online_tracker::{self};
use crate::util::crypto_helper::encrypt;
use crate::util::logger::PrintType;
use crate::{get_private_key, log_out};
@ -18,6 +14,7 @@ use futures::stream::SplitStream;
use hyper::upgrade::Upgraded;
use hyper_util::rt::TokioIo;
use json::JsonValue;
use json::number::Number;
use rand::Rng;
use rand::distributions::Alphanumeric;
use std::sync::Arc;
@ -73,7 +70,7 @@ impl OmikronConnection {
}
sender.send(message_text).await.unwrap();
}
pub async fn get_user_id(&self) -> i64 {
pub async fn get_omikron_id(&self) -> i64 {
*self.omikron_id.read().await
}
pub async fn is_identified(&self) -> bool {
@ -102,137 +99,168 @@ impl OmikronConnection {
return;
}
// Handle identification
if !*self.identified.read().await && cv.is_type(CommunicationType::identification) {
let omikron_id = cv
.get_data(DataTypes::omikron)
.unwrap_or(&JsonValue::Null)
.as_i64()
.unwrap_or(0);
match get_omikron_by_id(omikron_id).await {
Ok((public_key, _)) => {
// Generate Challenge, encrypt it and send it to the omikron
*self.omikron_id.write().await = omikron_id;
// If not yet identified
if !self.is_identified().await {
// handle identification
if !*self.identified.read().await && cv.is_type(CommunicationType::identification) {
let omikron_id = cv
.get_data(DataTypes::omikron)
.unwrap_or(&JsonValue::Null)
.as_i64()
.unwrap_or(0);
match get_omikron_by_id(omikron_id).await {
Ok((public_key, _)) => {
// Generate Challenge, encrypt it and send it to the omikron
*self.omikron_id.write().await = omikron_id;
let challenge_str: String = rand::thread_rng()
.sample_iter(&Alphanumeric)
.take(32)
.map(char::from)
.collect();
let challenge_str: String = rand::thread_rng()
.sample_iter(&Alphanumeric)
.take(32)
.map(char::from)
.collect();
*self.challenge.write().await = challenge_str.clone();
*self.challenge.write().await = challenge_str.clone();
let user_public_key_bytes = match STANDARD.decode(&public_key) {
Ok(bytes) => bytes,
Err(_) => {
self.send_error_response(
&cv.get_id(),
CommunicationType::error_invalid_omikron_id,
)
.await;
return;
}
};
*self.pub_key.write().await = Some(user_public_key_bytes.clone());
let omikron_pub_key: PublicKey =
match PublicKey::from_bytes(&user_public_key_bytes) {
Some(key) => key,
None => {
let user_public_key_bytes = match STANDARD.decode(&public_key) {
Ok(bytes) => bytes,
Err(_) => {
self.send_error_response(
&cv.get_id(),
CommunicationType::error_invalid_public_key,
CommunicationType::error_invalid_omikron_id,
)
.await;
return;
}
};
*self.pub_key.write().await = Some(user_public_key_bytes.clone());
let encrypted_challenge =
encrypt(get_private_key(), omikron_pub_key, &challenge_str)
.unwrap_or("".to_string());
let omikron_pub_key: PublicKey =
match PublicKey::from_bytes(&user_public_key_bytes) {
Some(key) => key,
_ => {
self.send_error_response(
&cv.get_id(),
CommunicationType::error_invalid_public_key,
)
.await;
return;
}
};
let response = CommunicationValue::new(CommunicationType::challenge)
.add_data_str(
DataTypes::public_key,
STANDARD.encode(get_public_key().as_bytes()),
let encrypted_challenge =
encrypt(get_private_key(), omikron_pub_key, &challenge_str)
.unwrap_or("".to_string());
let response = CommunicationValue::new(CommunicationType::challenge)
.add_data_str(
DataTypes::public_key,
STANDARD.encode(get_public_key().as_bytes()),
)
.add_data_str(DataTypes::challenge, encrypted_challenge)
.with_id(cv.get_id());
self.send_message(&response).await;
*self.identified.write().await = true;
return;
}
Err(e) => {
self.send_message(
&CommunicationValue::new(CommunicationType::error_not_authenticated)
.with_id(cv.get_id())
.add_data_str(DataTypes::error_type, e.to_string()),
)
.add_data_str(DataTypes::challenge, encrypted_challenge)
.with_id(cv.get_id());
.await;
self.send_message(&response).await;
*self.identified.write().await = true;
return;
return;
}
}
Err(e) => {
self.send_message(
&CommunicationValue::new(CommunicationType::error_not_authenticated)
.with_id(cv.get_id())
.add_data_str(DataTypes::error_type, e.to_string()),
}
// Handle challenge response
if *self.identified.read().await
&& !*self.challenged.read().await
&& cv.is_type(CommunicationType::challenge_response)
{
let client_response = cv
.get_data(DataTypes::challenge)
.unwrap_or(&JsonValue::Null)
.as_str()
.unwrap_or("");
let expected_challenge = self.challenge.read().await.clone();
if client_response == expected_challenge {
*self.challenged.write().await = true;
let response =
CommunicationValue::new(CommunicationType::identification_response)
.with_id(cv.get_id());
let _ = sql::set_omikron_active(self.get_omikron_id().await, true);
self.send_message(&response).await;
} else {
self.send_error_response(
&cv.get_id(),
CommunicationType::error_invalid_challenge,
)
.await;
return;
self.close().await;
}
return;
}
}
// Handle challenge response
if *self.identified.read().await
&& !*self.challenged.read().await
&& cv.is_type(CommunicationType::challenge_response)
{
let client_response = cv
.get_data(DataTypes::challenge)
.unwrap_or(&JsonValue::Null)
.as_str()
.unwrap_or("");
let expected_challenge = self.challenge.read().await.clone();
if client_response == expected_challenge {
*self.challenged.write().await = true;
let response = CommunicationValue::new(CommunicationType::identification_response)
.with_id(cv.get_id());
self.send_message(&response).await;
} else {
self.send_error_response(&cv.get_id(), CommunicationType::error_invalid_challenge)
.await;
self.close().await;
}
return;
}
if !self.is_identified().await {
// if not identified && not identifying
self.send_error_response(&cv.get_id(), CommunicationType::error_not_authenticated)
.await;
self.close().await;
return;
}
let omikron_id = self.get_user_id().await;
let omikron_id = self.get_omikron_id().await;
if cv.is_type(CommunicationType::user_connected) {
if let Some(user_id) = cv.get_data(DataTypes::user_id).and_then(|v| v.as_i64()) {
track_user_omikron(user_id, omikron_id).await;
user_online_tracker::track_user_status(user_id, ConnectionType::Online, omikron_id)
.await;
}
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()) {
untrack_user(user_id).await;
if let Some(status) = user_online_tracker::get_user_status(user_id).await {
user_online_tracker::track_user_status(
user_id,
ConnectionType::UserOffline,
status.omikron_id,
)
.await;
}
}
return;
}
if cv.is_type(CommunicationType::iota_connected) {
if let Some(iota_id) = cv.get_data(DataTypes::iota_id).and_then(|v| v.as_i64()) {
track_iota_omikron(iota_id, omikron_id).await;
user_online_tracker::track_iota_connection(iota_id, omikron_id).await;
if let Ok(users) = sql::get_users_by_iota_id(iota_id).await {
for user in users {
user_online_tracker::track_user_status(
user.0,
ConnectionType::UserOffline,
omikron_id,
)
.await;
}
}
}
return;
}
if cv.is_type(CommunicationType::iota_disconnected) {
if let Some(iota_id) = cv.get_data(DataTypes::iota_id).and_then(|v| v.as_i64()) {
untrack_iota(iota_id).await;
let iota_offline =
user_online_tracker::untrack_iota_connection(iota_id, omikron_id).await;
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;
}
}
}
return;
}
@ -242,7 +270,12 @@ impl OmikronConnection {
{
for user_id_json in user_ids {
if let Some(user_id) = user_id_json.as_i64() {
track_user_omikron(user_id, omikron_id).await;
user_online_tracker::track_user_status(
user_id,
ConnectionType::Online,
omikron_id,
)
.await;
}
}
}
@ -251,12 +284,286 @@ impl OmikronConnection {
{
for iota_id_json in iota_ids {
if let Some(iota_id) = iota_id_json.as_i64() {
track_iota_omikron(iota_id, omikron_id).await;
user_online_tracker::track_iota_connection(iota_id, omikron_id).await;
}
}
}
return;
}
if cv.is_type(CommunicationType::get_user_data) {
if let Some(user_id) = cv.get_data(DataTypes::user_id).cloned() {
if let Some(user_id) = user_id.as_i64() {
if let Ok((
id,
iota_id,
username,
display,
status,
about,
avatar,
sub_level,
sub_end,
public_key,
_,
_,
)) = get_by_user_id(user_id).await
{
let mut response =
CommunicationValue::new(CommunicationType::get_user_data)
.add_data_str(DataTypes::username, username)
.add_data_str(DataTypes::public_key, public_key)
.add_data(DataTypes::user_id, JsonValue::Number(Number::from(id)))
.add_data(
DataTypes::iota_id,
JsonValue::Number(Number::from(iota_id)),
)
.add_data_str(DataTypes::display, display)
.add_data_str(DataTypes::status, status)
.add_data_str(DataTypes::about, about)
.add_data_str(DataTypes::avatar, avatar)
.add_data(
DataTypes::sub_level,
JsonValue::Number(Number::from(sub_level)),
)
.add_data(
DataTypes::sub_end,
JsonValue::Number(Number::from(sub_end)),
);
let user_status = user_online_tracker::get_user_status(id).await;
let iota_connections =
user_online_tracker::get_iota_omikron_connections(iota_id)
.await
.unwrap_or_default();
response = response.add_data(
DataTypes::omikron_connections,
JsonValue::Array(
iota_connections
.into_iter()
.map(|id| JsonValue::Number(Number::from(id)))
.collect(),
),
);
if let Some(user_status) = user_status {
response = response.add_data(
DataTypes::online_status,
JsonValue::String(user_status.connection_type.to_string()),
);
response = response.add_data(
DataTypes::omikron_id,
JsonValue::Number(Number::from(user_status.omikron_id)),
);
} else {
response = response.add_data(
DataTypes::online_status,
JsonValue::String(ConnectionType::IotaOffline.to_string()),
);
}
self.send_message(&response).await;
return;
}
}
}
if let Some(username) = cv.get_data(DataTypes::username).cloned() {
if let Some(username) = username.as_str() {
if let Ok((
id,
iota_id,
username,
display,
status,
about,
avatar,
sub_level,
sub_end,
public_key,
_,
_,
)) = get_by_username(username).await
{
let mut response =
CommunicationValue::new(CommunicationType::get_user_data)
.add_data_str(DataTypes::username, username)
.add_data_str(DataTypes::public_key, public_key)
.add_data(DataTypes::user_id, JsonValue::Number(Number::from(id)))
.add_data(
DataTypes::iota_id,
JsonValue::Number(Number::from(iota_id)),
)
.add_data_str(DataTypes::display, display)
.add_data_str(DataTypes::status, status)
.add_data_str(DataTypes::about, about)
.add_data_str(DataTypes::avatar, avatar)
.add_data(
DataTypes::sub_level,
JsonValue::Number(Number::from(sub_level)),
)
.add_data(
DataTypes::sub_end,
JsonValue::Number(Number::from(sub_end)),
);
let user_status = user_online_tracker::get_user_status(id).await;
let iota_connections =
user_online_tracker::get_iota_omikron_connections(iota_id)
.await
.unwrap_or_default();
if let Some(user_status) = user_status {
response = response.add_data(
DataTypes::online_status,
JsonValue::String(user_status.connection_type.to_string()),
);
response = response.add_data(
DataTypes::omikron_id,
JsonValue::Number(Number::from(user_status.omikron_id)),
);
} else {
response = response.add_data(
DataTypes::online_status,
JsonValue::String(ConnectionType::IotaOffline.to_string()),
);
}
response = response.add_data(
DataTypes::omikron_connections,
JsonValue::Array(
iota_connections
.iter()
.map(|&id| JsonValue::Number(Number::from(id)))
.collect(),
),
);
self.send_message(&response).await;
return;
}
}
}
let response =
CommunicationValue::new(CommunicationType::error_not_found).with_id(cv.get_id());
self.send_message(&response).await;
return;
}
if cv.is_type(CommunicationType::get_iota_data) {
if let Some(iota_id) = cv.get_data(DataTypes::iota_id).cloned() {
if let Some(iota_id) = iota_id.as_i64() {
if let Ok((iota_id, public_key)) = get_iota_by_id(iota_id).await {
let mut response =
CommunicationValue::new(CommunicationType::get_iota_data)
.add_data_str(DataTypes::public_key, public_key)
.add_data(
DataTypes::iota_id,
JsonValue::Number(Number::from(iota_id)),
);
let iota_connections =
user_online_tracker::get_iota_omikron_connections(iota_id)
.await
.unwrap_or_default();
response = response.add_data(
DataTypes::omikron_connections,
JsonValue::Array(
iota_connections
.iter()
.map(|&id| JsonValue::Number(Number::from(id)))
.collect(),
),
);
self.send_message(&response).await;
return;
}
}
}
if let Some(user_id) = cv.get_data(DataTypes::user_id).cloned() {
if let Some(user_id) = user_id.as_i64() {
if let Ok((_, iota_id, _, _, _, _, _, _, _, _, _, _)) =
get_by_user_id(user_id).await
{
if let Ok((iota_id, public_key)) = get_iota_by_id(iota_id).await {
let mut response =
CommunicationValue::new(CommunicationType::get_iota_data)
.add_data_str(DataTypes::public_key, public_key)
.add_data(
DataTypes::user_id,
JsonValue::Number(Number::from(user_id)),
)
.add_data(
DataTypes::iota_id,
JsonValue::Number(Number::from(iota_id)),
);
let iota_connections =
user_online_tracker::get_iota_omikron_connections(iota_id)
.await
.unwrap_or_default();
response = response.add_data(
DataTypes::omikron_connections,
JsonValue::Array(
iota_connections
.iter()
.map(|&id| JsonValue::Number(Number::from(id)))
.collect(),
),
);
self.send_message(&response).await;
return;
}
}
}
}
if let Some(username) = cv.get_data(DataTypes::username).cloned() {
if let Some(username) = username.as_str() {
if let Ok((user_id, iota_id, _, _, _, _, _, _, _, _, _, _)) =
get_by_username(username).await
{
if let Ok((iota_id, public_key)) = get_iota_by_id(iota_id).await {
let mut response =
CommunicationValue::new(CommunicationType::get_iota_data)
.add_data_str(DataTypes::public_key, public_key)
.add_data(
DataTypes::user_id,
JsonValue::Number(Number::from(user_id)),
)
.add_data_str(DataTypes::username, username.to_string())
.add_data(
DataTypes::iota_id,
JsonValue::Number(Number::from(iota_id)),
);
let iota_connections =
user_online_tracker::get_iota_omikron_connections(iota_id)
.await
.unwrap_or_default();
response = response.add_data(
DataTypes::omikron_connections,
JsonValue::Array(
iota_connections
.iter()
.map(|&id| JsonValue::Number(Number::from(id)))
.collect(),
),
);
self.send_message(&response).await;
return;
}
}
}
}
let response =
CommunicationValue::new(CommunicationType::error_not_found).with_id(cv.get_id());
self.send_message(&response).await;
return;
}
if cv.is_type(CommunicationType::change_user_data) {}
if cv.is_type(CommunicationType::change_iota_data) {}
}
async fn send_error_response(&self, message_id: &Uuid, error_type: CommunicationType) {
@ -265,14 +572,16 @@ 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 _ = sender.close().await;
}
pub async fn handle_close(self: Arc<Self>) {
if self.is_identified().await {
let omikron_id = self.get_user_id().await;
let omikron_id = self.get_omikron_id().await;
if omikron_id != 0 {
untrack_iota_by_omikron(omikron_id).await;
untrack_user_by_omikron(omikron_id).await;
user_online_tracker::untrack_omikron(omikron_id).await;
}
}
}

107
src/server/server.rs
View file

@ -1,12 +1,11 @@
use crate::log;
use crate::server::api;
use crate::server::omikron_connection::OmikronConnection;
use crate::server::socket;
use crate::util::file_util::load_file_buf;
use base64::Engine;
use base64::engine::general_purpose::STANDARD;
use bytes::Bytes;
use futures::StreamExt;
use futures_util::TryFutureExt;
use http_body_util::BodyExt;
use http_body_util::Full;
@ -30,9 +29,8 @@ use std::{future::Future, pin::Pin, time::Duration};
use tokio::net::TcpListener;
use tokio::sync::broadcast;
use tokio_rustls::TlsAcceptor;
use tokio_tungstenite::{WebSocketStream, accept_async};
use tower::Service;
use tungstenite::Message;
#[derive(Clone)]
struct HttpService {
peer_addr: SocketAddr,
@ -60,7 +58,7 @@ impl Service<HttpRequest<Incoming>> for HttpService {
let headers = parts.headers.clone();
let fut = async move {
let is_websocket_upgrade = path == "/ws/omikron"
let is_websocket_upgrade = path == "/ws"
&& method == Method::GET
&& headers
.get("connection")
@ -94,39 +92,7 @@ impl Service<HttpRequest<Incoming>> for HttpService {
let upgrades = upgrade::on(req_for_upgrade);
log!("Handling WebSocket upgrade");
// Spawn upgrade handling to avoid blocking the service call
tokio::spawn(async move {
match upgrades.await {
Ok(upgraded_stream) => {
log!("Valid WebSocket upgrade");
let raw_stream = TokioIo::new(upgraded_stream);
let ws_stream = WebSocketStream::from_raw_socket(
raw_stream,
tungstenite::protocol::Role::Server,
None,
)
.await;
log!(
"WebSocket handshake successful, handling connection for Omikron"
);
// Split stream for OmikronConnection
let (writer, reader) = ws_stream.split();
// INTEGRATION START
// Erstelle die Connection und starte den Handler
let connection = OmikronConnection::new(writer, reader);
// Handler in separatem Task starten
start_omikron_handler(connection).await;
// INTEGRATION END
}
Err(e) => {
log!("WebSocket upgrade failed after response: {:?}", e);
}
}
});
socket::handle(path, upgrades);
log!("Handled WebSocket connection initiation");
Ok(response)
@ -192,7 +158,6 @@ async fn run_http_server(port: u16) -> bool {
port
);
// Create a broadcast channel for graceful shutdown signal
let (shutdown_tx, _) = broadcast::channel::<()>(1);
tokio::spawn(async move {
@ -205,7 +170,6 @@ async fn run_http_server(port: u16) -> bool {
let service = HttpService { peer_addr: addr };
let io = TokioIo::new(stream);
// Subscribe to the shutdown signal for this specific connection
let mut rx = shutdown_tx.subscribe();
tokio::spawn(async move {
@ -216,7 +180,6 @@ async fn run_http_server(port: u16) -> bool {
.serve_connection(io, TowerToHyperService::new(service))
.with_upgrades();
// Wait for either the connection to finish naturally OR the shutdown signal
tokio::select! {
res = conn => {
if let Err(err) = res {
@ -276,13 +239,11 @@ async fn run_tls_server(port: u16, tls_config: Arc<ServerConfig>) -> bool {
port
);
// Create a broadcast channel for graceful shutdown signal
let (shutdown_tx, _) = broadcast::channel::<()>(1);
tokio::spawn(async move {
loop {
tokio::select! {
// Monitor for shutdown signal
_ = async {
loop {
tokio::time::sleep(Duration::from_millis(100)).await;
@ -294,7 +255,6 @@ async fn run_tls_server(port: u16, tls_config: Arc<ServerConfig>) -> bool {
break;
}
// Accept new connections
accepted = listener.accept() => {
match accepted {
std::result::Result::Ok((stream, addr)) => {
@ -317,14 +277,12 @@ async fn run_tls_server(port: u16, tls_config: Arc<ServerConfig>) -> bool {
};
let io = TokioIo::new(tls_stream);
// Prepare connection future
let conn = http1::Builder::new()
.preserve_header_case(true)
.title_case_headers(true)
.serve_connection(io, TowerToHyperService::new(service))
.with_upgrades();
// Wait for either the connection to finish naturally OR the shutdown signal
tokio::select! {
res = conn => {
if let Err(err) = res {
@ -363,7 +321,7 @@ pub async fn start(port: u16) -> bool {
match tls_result {
Ok(Some(tls_config)) => run_tls_server(port, tls_config).await,
Ok(None) => run_http_server(port).await,
Ok(_) => run_http_server(port).await,
Err(e) => {
log!("Fatal error during TLS config load: {}", e);
false
@ -376,7 +334,7 @@ fn calculate_accept_key(key: &str) -> String {
sha1.update(key.as_bytes());
sha1.update(websocket_guid.as_bytes());
let result = sha1.finalize();
STANDARD.encode(result) // Base64 encode the result
STANDARD.encode(result)
}
/// Loads TLS config. Returns Ok(None) if cert files are not found, and an error if parsing fails.
@ -391,7 +349,7 @@ fn load_tls_config() -> Result<Option<Arc<ServerConfig>>, Box<dyn Error>> {
log!("TLS certificate 'certs/cert.pem' not found.");
return Ok(None);
}
Err(e) => return Err(e.into()), // Other IO error
Err(e) => return Err(e.into()),
};
let key_file_buf = match key_file_res {
@ -400,7 +358,7 @@ fn load_tls_config() -> Result<Option<Arc<ServerConfig>>, Box<dyn Error>> {
log!("TLS key 'certs/cert.key' not found.");
return Ok(None);
}
Err(e) => return Err(e.into()), // Other IO error
Err(e) => return Err(e.into()),
};
// Continue with configuration if both files were found
@ -411,12 +369,12 @@ fn load_tls_config() -> Result<Option<Arc<ServerConfig>>, Box<dyn Error>> {
// PKCS8
let mut key_reader = BufReader::new(key_file_buf);
let mut key_ders = rustls_pemfile::pkcs8_private_keys(&mut key_reader)
.map(|r| r.map(Into::into)) // Explicit conversion
.map(|r| r.map(Into::into))
.collect::<Result<Vec<PrivateKeyDer>, io::Error>>()?;
if key_ders.is_empty() {
// RSA
key_reader = BufReader::new(load_file_buf("certs", "cert.key")?); // Re-read key file
key_reader = BufReader::new(load_file_buf("certs", "cert.key")?);
key_ders = rustls_pemfile::rsa_private_keys(&mut key_reader)
.map(|r| r.map(Into::into))
.collect::<Result<Vec<PrivateKeyDer>, io::Error>>()?;
@ -424,14 +382,14 @@ fn load_tls_config() -> Result<Option<Arc<ServerConfig>>, Box<dyn Error>> {
if key_ders.is_empty() {
// EC
key_reader = BufReader::new(load_file_buf("certs", "cert.key")?); // Re-read key file
key_reader = BufReader::new(load_file_buf("certs", "cert.key")?);
key_ders = rustls_pemfile::ec_private_keys(&mut key_reader)
.map(|r| r.map(Into::into))
.collect::<Result<Vec<PrivateKeyDer>, io::Error>>()?;
}
if key_ders.is_empty() {
return Err("No private keys found in key file. (Tried PKCS8, RSA, and EC)".into());
return Err("No valid private keys found in key file (Tried PKCS8, RSA and EC).".into());
}
let config = rustls::ServerConfig::builder()
@ -441,44 +399,3 @@ fn load_tls_config() -> Result<Option<Arc<ServerConfig>>, Box<dyn Error>> {
Ok(Some(Arc::new(config)))
}
// In deiner Server-Logik, wo OmikronConnection initialisiert wird:
pub async fn start_omikron_handler(connection: Arc<OmikronConnection>) {
loop {
let msg = match {
let mut receiver = connection.receiver.write().await;
receiver.next().await
} {
Some(Ok(msg)) => msg,
Some(Err(e)) => {
log!("WS Error: {}", e);
break;
}
None => break, // Stream ended
};
match msg {
Message::Text(text) => {
let conn_clone = connection.clone();
tokio::spawn(async move {
conn_clone.handle_message(text.to_string()).await;
});
}
Message::Ping(_) => {
let pong_response = crate::data::communication::CommunicationValue::new(
crate::data::communication::CommunicationType::pong,
);
let conn_clone = connection.clone();
tokio::spawn(async move {
conn_clone.send_message(&pong_response).await;
});
}
Message::Close(_) => {
break;
}
// Other message types like Binary, Pong are ignored.
_ => {}
}
}
connection.handle_close().await;
}

101
src/server/socket.rs
View file

@ -1,63 +1,70 @@
use std::sync::Arc;
use futures::StreamExt;
use futures::stream::SplitSink;
use futures::stream::SplitStream;
use hyper::upgrade::Upgraded;
use hyper::upgrade::OnUpgrade;
use hyper_util::rt::TokioIo;
use tokio_tungstenite::WebSocketStream;
use tungstenite::Message;
use crate::log;
use crate::server::omikron_connection::OmikronConnection;
pub fn handle(
path: String,
writer: SplitSink<tokio_tungstenite::WebSocketStream<TokioIo<Upgraded>>, Message>,
reader: SplitStream<tokio_tungstenite::WebSocketStream<TokioIo<Upgraded>>>,
) {
log!("handling");
pub fn handle(path: String, upgrades: OnUpgrade) {
tokio::spawn(async move {
if path.starts_with("/ws/phi/") {
} else if path.starts_with("/ws/omikron/") {
let community_conn: Arc<OmikronConnection> =
Arc::from(OmikronConnection::new(writer, reader));
loop {
let msg_result: Option<Result<_, _>> = {
let mut session_lock = community_conn.receiver.write().await;
session_lock.next().await
};
match upgrades.await {
Ok(upgraded_stream) => {
log!("Valid WebSocket upgrade");
let raw_stream = TokioIo::new(upgraded_stream);
match msg_result {
Some(Ok(msg)) => {
if msg.is_text() {
let text = msg.into_text().unwrap();
community_conn
.clone()
.handle_message(text.to_string())
.await;
} else if msg.is_ping() {
let pong_response = crate::data::communication::CommunicationValue::new(
crate::data::communication::CommunicationType::pong,
);
community_conn.send_message(&pong_response).await;
} else if msg.is_close() {
log!("Closing: {}", msg);
community_conn.handle_close().await;
return;
}
}
Some(Err(e)) => {
log!("Closing ERR: {}", e);
community_conn.handle_close().await;
return;
}
None => {
log!("Closed Session me!");
community_conn.handle_close().await;
return;
}
let ws_stream = WebSocketStream::from_raw_socket(
raw_stream,
tungstenite::protocol::Role::Server,
None,
)
.await;
log!(
"WebSocket handshake successful, handling connection for {}",
path
);
let (writer, reader) = ws_stream.split();
if path == "/ws/omikron" {
let connection = OmikronConnection::new(writer, reader);
start_connecteable_handler(connection).await;
}
}
Err(e) => {
log!("WebSocket upgrade failed after response: {:?}", e);
}
}
});
}
pub async fn start_connecteable_handler(connection: Arc<OmikronConnection>) {
loop {
let msg = match {
let mut receiver = connection.receiver.write().await;
receiver.next().await
} {
Some(Ok(msg)) => msg,
Some(Err(e)) => {
log!("WS Error: {}", e);
break;
}
_ => break,
};
match msg {
Message::Text(text) => {
let conn_clone = connection.clone();
tokio::spawn(async move {
conn_clone.handle_message(text.to_string()).await;
});
}
Message::Close(_) => {
break;
}
_ => {}
}
}
connection.handle_close().await;
}

38
src/sql/connection_status.rs Normal file
View file

@ -0,0 +1,38 @@
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ConnectionType {
Online,
UserOffline,
IotaOffline,
Away,
DoNotDisturb,
}
impl ConnectionType {
pub fn to_str(&self) -> &str {
match self {
ConnectionType::Online => "online",
ConnectionType::UserOffline => "user_offline",
ConnectionType::IotaOffline => "iota_offline",
ConnectionType::Away => "away",
ConnectionType::DoNotDisturb => "do_not_disturb",
}
}
pub fn to_string(&self) -> String {
match self {
ConnectionType::Online => "online".to_string(),
ConnectionType::UserOffline => "user_offline".to_string(),
ConnectionType::IotaOffline => "iota_offline".to_string(),
ConnectionType::Away => "away".to_string(),
ConnectionType::DoNotDisturb => "do_not_disturb".to_string(),
}
}
pub fn from_str(s: &str) -> Option<ConnectionType> {
match s.to_lowercase().as_str() {
"online" => Some(ConnectionType::Online),
"user_offline" => Some(ConnectionType::UserOffline),
"iota_offline" => Some(ConnectionType::IotaOffline),
"away" => Some(ConnectionType::Away),
"do_not_disturb" => Some(ConnectionType::DoNotDisturb),
_ => None,
}
}
}

1
src/sql/mod.rs
View file

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

73
src/sql/sql.rs
View file

@ -202,11 +202,11 @@ pub async fn get_by_username(
String::from_utf8_lossy(&token).to_string(),
))
}
None => Err(sqlx::Error::RowNotFound),
_ => Err(sqlx::Error::RowNotFound),
}
}
pub async fn get_by_id(
pub async fn get_by_user_id(
id: i64,
) -> Result<
(
@ -265,10 +265,73 @@ pub async fn get_by_id(
String::from_utf8_lossy(&token).to_string(),
))
}
None => Err(sqlx::Error::RowNotFound),
_ => Err(sqlx::Error::RowNotFound),
}
}
pub async fn get_users_by_iota_id(
iota_id_param: i64,
) -> Result<
Vec<(
i64,
i64,
String,
String,
String,
String,
String,
i32,
i64,
String,
String,
String,
)>,
sqlx::Error,
> {
let db_lock = SQL_DB.read().await;
let pool = db_lock.as_ref().expect("Database pool is not initialized");
let rows = sqlx::query(
"SELECT id, iota_id, username, display, status, about, avatar, sub_level, sub_end, public_key, private_key_hash, token FROM users WHERE iota_id = ?",
)
.bind(iota_id_param)
.fetch_all(pool)
.await?;
let mut users = Vec::new();
for row in rows {
let id: i64 = row.get("id");
let iota_id: i64 = row.get("iota_id");
let username: String = row.get("username");
let display: Vec<u8> = row.get("display");
let status: Vec<u8> = row.get("status");
let about: Vec<u8> = row.get("about");
let avatar: Vec<u8> = row.get("avatar");
let sub_level: i32 = row.get("sub_level");
let sub_end: i64 = row.get("sub_end");
let public_key: String = row.get("public_key");
let private_key_hash: String = row.get("private_key_hash");
let token: Vec<u8> = row.get("token");
users.push((
id,
iota_id,
username,
String::from_utf8_lossy(&display).to_string(),
String::from_utf8_lossy(&status).to_string(),
String::from_utf8_lossy(&about).to_string(),
String::from_utf8_lossy(&avatar).to_string(),
sub_level,
sub_end,
public_key,
private_key_hash,
String::from_utf8_lossy(&token).to_string(),
));
}
Ok(users)
}
pub async fn change_username(id: i64, new_username: String) -> Result<(), sqlx::Error> {
let db_lock = SQL_DB.read().await;
let pool = db_lock.as_ref().expect("Database pool is not initialized");
@ -514,7 +577,7 @@ pub async fn get_random_omikron() -> Result<(i64, String, String), sqlx::Error>
String::from_utf8_lossy(&public_key).to_string(),
String::from_utf8_lossy(&ip_address).to_string(),
)),
None => Err(sqlx::Error::RowNotFound),
_ => Err(sqlx::Error::RowNotFound),
}
}
@ -534,7 +597,7 @@ pub async fn get_omikron_by_id(id: i64) -> Result<(String, String), sqlx::Error>
String::from_utf8_lossy(&public_key).to_string(),
String::from_utf8_lossy(&ip_address).to_string(),
)),
None => Err(sqlx::Error::RowNotFound),
_ => Err(sqlx::Error::RowNotFound),
}
}

99
src/sql/user_online_tracker.rs
View file

@ -1,27 +1,100 @@
use crate::sql;
use crate::sql::connection_status::ConnectionType;
use once_cell::sync::Lazy;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
static USER_OMIKRON_MAP: Lazy<Arc<RwLock<HashMap<i64, i64>>>> =
#[derive(Debug, Clone)]
pub struct UserStatus {
pub connection_type: ConnectionType,
pub omikron_id: i64,
}
// IotaID -> Vec<OmikronID>
static IOTA_OMIKRON_CONNECTIONS: Lazy<Arc<RwLock<HashMap<i64, Vec<i64>>>>> =
Lazy::new(|| Arc::new(RwLock::new(HashMap::new())));
pub async fn track_user_omikron(user: i64, omikron: i64) {
let mut c = USER_OMIKRON_MAP.write().await;
c.insert(user, omikron);
// UserID -> UserStatus
static USER_STATUS_MAP: Lazy<Arc<RwLock<HashMap<i64, UserStatus>>>> =
Lazy::new(|| Arc::new(RwLock::new(HashMap::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 async fn get_omikron_for_user(user: i64) -> Option<i64> {
let c = USER_OMIKRON_MAP.read().await;
c.get(&user).cloned()
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) {
connections.retain(|&id| id != omikron_id);
if connections.is_empty() {
iota_map.remove(&iota_id);
return true; // Iota is now offline
}
}
false
}
pub async fn untrack_user(user: i64) {
let mut c = USER_OMIKRON_MAP.write().await;
c.remove(&user);
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 async fn untrack_by_omikron(omikron: i64) {
let mut c = USER_OMIKRON_MAP.write().await;
c.retain(|_, v| *v != omikron);
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(
user_id,
UserStatus {
connection_type: status,
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 async fn untrack_user(user_id: i64) {
let mut user_map = USER_STATUS_MAP.write().await;
user_map.remove(&user_id);
}
pub async fn untrack_many_users(user_ids: &[i64]) {
let mut user_map = USER_STATUS_MAP.write().await;
for user_id in user_ids {
user_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| {
connections.retain(|id| *id != omikron_id);
if connections.is_empty() {
offline_iotas.push(*iota_id);
false
} else {
true
}
});
user_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);
}
}
}
}