This commit is contained in:
Alois 2026-04-02 22:25:10 +02:00
commit b23eff24a3
39 changed files with 4535 additions and 4430 deletions

16
.gitignore vendored
View file

@ -8,22 +8,12 @@ target
# MSVC Windows builds of rustc generate these, which store debugging information
*.pdb
*.env
# Generated by cargo mutants
# Contains mutation testing data
**/mutants.out*/
**/mutants.out*/W
# RustRover
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
.idea/
.env
/logs
# Added by cargo
/target
/logs

2527
Cargo.lock generated

File diff suppressed because it is too large Load diff

38
Cargo.toml Executable file → Normal file
View file

@ -1,39 +1,37 @@
[package]
name = "Omega"
name = "Omikron"
version = "0.1.0"
edition = "2024"
[dependencies]
epsilon-core = { git = "https://github.com/Tensamin/Epsilon.git", package = "epsilon-core" }
epsilon-native = { git = "https://github.com/Tensamin/Epsilon.git", package = "epsilon-native" }
ttp-core = { git = "https://github.com/Tensamin/TTP.git", package = "ttp-core" }
ttp-native = { git = "https://github.com/Tensamin/TTP.git", package = "ttp-native" }
ansi_term = "*"
uuid = { version = "*", features = ["v4"] }
actix-web = { version = "4.12.1", features = ["rustls-0_23"] }
aes-gcm = "*"
ansi_term = "0.12.1"
anyhow = "1.0.101"
base64 = "0.22.1"
dashmap = "6.1.0"
dotenv = "0.15.0"
hex = "0.4.3"
hkdf = "0.12.4"
json = "0.12.4"
dashmap = "*"
futures = "*"
hex = "*"
once_cell = "1.21.3"
rand = "0.8"
rand_core = { version = "0.6", features = ["getrandom", "std"] }
reqwest = { version = "0.13.2" }
rustls = { version = "0.23.37", default-features = false, features = [
"std",
"tls12",
"aws-lc-rs",
"prefer-post-quantum",
] }
rustls-pemfile = "2.2.0"
aes-gcm = "0.10.3"
sha2 = "0.10.9"
sqlx = { version = "0.8.6", features = ["mysql", "runtime-async-std"] }
strum = "0.27.2"
strum_macros = "0.27.2"
tokio = { version = "*", features = ["full"] }
uuid = { version = "1.19.0", features = ["v4"] }
x448 = "0.6.0"
zip = "6.0.0"
x448 = { version = "*" }
log = "0.4"
dotenv = "0.15.0"
hkdf = "0.12.4"
strum = "0.28.0"
strum_macros = "0.28.0"
livekit-api = { version = "0.4.14", features = ["native-tls"] }
livekit-protocol = "0.7.1"
thiserror = "2.0.18"

View file

@ -1,2 +1 @@
# Omega
A loadballancer and Omikron-manager for Tensamin
# ⚠️ Moved to [git.methanium.net](https://git.methanium.net/tensamin/omikron) ⚠️

View file

@ -0,0 +1,519 @@
use std::str::FromStr;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use ttp_core::{CommunicationType, CommunicationValue, DataTypes, DataValue};
use ttp_native::{Receiver, Sender};
use uuid::Uuid;
use crate::anonymous_clients::anonymous_manager::{self, generate_username};
use crate::calls::call_manager;
use crate::omega::omega_connection::get_omega_connection;
use crate::rho::connection::GeneralConnection;
use crate::rho::rho_manager;
use crate::util::logger::PrintType;
use crate::{log_cv_in, log_cv_out, log_out};
pub struct AnonymousClientConnection {
user_id: u64,
pub sender: Arc<Sender>,
pub receiver: Arc<Receiver>,
pub ping: Arc<RwLock<i64>>,
pub interested_users: Arc<RwLock<Vec<i64>>>,
is_open: Arc<RwLock<bool>>,
pub user_name: Arc<RwLock<String>>,
pub display_name: Arc<RwLock<String>>,
pub avatar: Arc<RwLock<String>>,
}
impl AnonymousClientConnection {
pub async fn from_general(general: Arc<GeneralConnection>, user_id: u64) -> Arc<Self> {
let username: String = generate_username();
Arc::new(Self {
user_id: user_id,
ping: Arc::new(RwLock::new(0)),
interested_users: Arc::new(RwLock::new(Vec::new())),
is_open: Arc::new(RwLock::new(true)),
sender: general.sender.clone(),
receiver: general.receiver.clone(),
user_name: Arc::new(RwLock::new(username.to_lowercase())),
display_name: Arc::new(RwLock::new(username)),
avatar: Arc::new(RwLock::new(String::new())),
})
}
pub fn start(self: Arc<Self>) {
let self_clone = self.clone();
tokio::spawn(async move {
while let Ok(cv) = self_clone.receiver.receive().await {
self_clone.clone().handle_message(cv).await;
}
self_clone.handle_close().await;
});
}
/// Get the user ID
pub fn get_user_id(&self) -> u64 {
self.user_id
}
/// Get the user name
pub async fn get_user_name(&self) -> String {
self.user_name.read().await.clone()
}
/// Get the display name
pub async fn get_display_name(&self) -> String {
self.display_name.read().await.clone()
}
pub async fn set_display_name(&self, display: String) {
*self.display_name.write().await = display;
}
/// Get the avatar
pub async fn get_avatar(&self) -> String {
self.avatar.read().await.clone()
}
/// Send a CommunicationValue to the client
pub async fn send_message(self: Arc<Self>, cv: &CommunicationValue) {
if !*self.is_open.read().await {
log_out!(
self.user_id as i64,
PrintType::Client,
"Attempted to send message to a closed connection."
);
return;
}
if !cv.is_type(CommunicationType::pong) {
log_cv_out!(PrintType::Client, &cv);
}
if let Err(e) = self.sender.send(&cv).await {
log_out!(
self.user_id as i64,
PrintType::Client,
"Send failed: {:?}",
e
);
}
}
/// Handle incoming message from client
pub async fn handle_message(self: Arc<Self>, cv: CommunicationValue) {
tokio::spawn(async move {
if cv.is_type(CommunicationType::ping) {
self.handle_ping(cv).await;
return;
}
log_cv_in!(PrintType::Client, &cv);
if cv.is_type(CommunicationType::identification) {
let call_id =
Uuid::parse_str(cv.get_data(DataTypes::call_id).as_str().unwrap_or(""))
.unwrap_or(Uuid::new_v4());
let call = if let Some(call) = call_manager::get_call(call_id).await {
if call.is_anonymous().await {
call
} else {
self.send_error_response(
&cv.get_id(),
CommunicationType::error_not_authenticated,
)
.await;
return;
}
} else {
self.send_error_response(
&cv.get_id(),
CommunicationType::error_not_authenticated,
)
.await;
return;
};
let mut invited = Vec::new();
for call_invitee in call.members.read().await.clone() {
let call_invitee_cv = get_omega_connection()
.await_response(
&CommunicationValue::new(CommunicationType::get_user_data).add_data(
DataTypes::user_id,
DataValue::Number(call_invitee.user_id as i64),
),
Some(Duration::from_secs(2)),
)
.await
.unwrap();
let mut json_invitee = Vec::new();
let _ = json_invitee.push((
DataTypes::user_id,
call_invitee_cv.get_data(DataTypes::user_id).clone(),
));
let _ = json_invitee.push((
DataTypes::username,
call_invitee_cv.get_data(DataTypes::username).clone(),
));
let _ = json_invitee.push((
DataTypes::display,
call_invitee_cv.get_data(DataTypes::display).clone(),
));
let _ = json_invitee.push((
DataTypes::avatar,
call_invitee_cv.get_data(DataTypes::avatar).clone(),
));
let _ = invited.push(DataValue::Container(json_invitee));
}
let token = call.create_anonymous_token(self.get_user_id()).await;
let mut serialized = Vec::new();
let _ = serialized.push((DataTypes::call_id, DataValue::Str(call_id.to_string())));
let _ =
serialized.push((DataTypes::call_invited, DataValue::Array(invited.clone())));
let _ = serialized.push((DataTypes::call_members, DataValue::Array(invited)));
let _ = serialized.push((DataTypes::call_token, DataValue::Str(token.unwrap())));
self.clone()
.send_message(
&&CommunicationValue::new(CommunicationType::identification_response)
.with_id(cv.get_id())
.add_data(DataTypes::user_id, DataValue::Number(self.user_id as i64))
.add_data(
DataTypes::username,
DataValue::Str(self.clone().get_user_name().await),
)
.add_data(
DataTypes::display,
DataValue::Str(self.get_display_name().await),
)
.add_data(DataTypes::avatar, DataValue::Str(self.get_avatar().await))
.add_data(DataTypes::call_state, DataValue::Container(serialized)),
)
.await;
}
// Handle ping
if cv.is_type(CommunicationType::ping) {
self.handle_ping(cv).await;
return;
}
// Handle client status changes
if cv.is_type(CommunicationType::client_changed) {
self.handle_client_changed(cv).await;
return;
}
// Handle call invites
if cv.is_type(CommunicationType::call_invite) {
self.handle_call_invite(cv).await;
return;
}
// Handle get call requests
if cv.is_type(CommunicationType::call_token) {
self.handle_get_call(cv).await;
return;
}
if cv.is_type(CommunicationType::call_disconnect_user) {
self.handle_call_disconnect_user(cv).await;
return;
}
if cv.is_type(CommunicationType::call_timeout_user) {
self.handle_call_timeout_user(cv).await;
return;
}
if cv.is_type(CommunicationType::change_user_data) {
if let Some(display_name) = cv.get_data(DataTypes::display).as_str() {
let _ = self.set_display_name(display_name.to_string()).await;
}
return;
}
if cv.is_type(CommunicationType::get_user_data) {
if let Some(anonymous) = {
if let Some(user_id) = cv.get_data(DataTypes::user_id).as_number() {
anonymous_manager::get_anonymous_user(user_id as u64).await
} else if let Some(username) = cv.get_data(DataTypes::username).as_str() {
anonymous_manager::get_anonymous_user_by_name(username.to_string()).await
} else {
None
}
} {
let response = CommunicationValue::new(CommunicationType::get_user_data)
.with_id(cv.get_id())
.add_data(
DataTypes::username,
DataValue::Str(anonymous.get_user_name().await),
)
.add_data(
DataTypes::user_id,
DataValue::Number(anonymous.user_id as i64),
)
.add_data(
DataTypes::display,
DataValue::Str(anonymous.get_display_name().await),
)
.add_data(DataTypes::user_state, DataValue::Str("online".to_string()))
.add_data(
DataTypes::avatar,
DataValue::Str(anonymous.get_avatar().await),
);
self.send_message(&response).await;
return;
}
}
if cv.is_type(CommunicationType::get_user_data)
|| cv.is_type(CommunicationType::get_iota_data)
|| cv.is_type(CommunicationType::delete_user)
{
self.handle_omega_forward(cv).await;
return;
}
});
}
async fn handle_omega_forward(self: Arc<Self>, cv: CommunicationValue) {
let client_for_closure = self.clone();
tokio::spawn(async move {
let response_cv = get_omega_connection()
.await_response(&cv.with_sender(self.user_id), Some(Duration::from_secs(20)))
.await;
if let Ok(response_cv) = response_cv {
client_for_closure.send_message(&response_cv).await;
}
});
}
/// Handle ping message
async fn handle_ping(self: Arc<Self>, cv: CommunicationValue) {
// Update our ping if provided
if let DataValue::Number(last_ping) = cv.get_data(DataTypes::last_ping) {
if let Ok(ping_val) = last_ping.to_string().parse::<i64>() {
let mut ping_guard = self.ping.write().await;
*ping_guard = ping_val;
}
}
// Send pong response
let response = CommunicationValue::new(CommunicationType::pong).with_id(cv.get_id());
self.send_message(&response).await;
}
/// Handle client status change
async fn handle_client_changed(self: Arc<Self>, _cv: CommunicationValue) {
/*let user_id = self.get_user_id().await;
if let Some(_status_str) = cv.get_data(DataTypes::user_state) {
let user_status = UserStatus::online;
}*/
}
/// Handle call invite
async fn handle_call_invite(self: Arc<Self>, cv: CommunicationValue) {
let receiver_id: i64 = cv.get_data(DataTypes::receiver_id).as_number().unwrap_or(0);
if receiver_id == 0 {
self.send_error_response(&cv.get_id(), CommunicationType::error_no_user_id)
.await;
return;
}
let call_id = match cv.get_data(DataTypes::call_id) {
DataValue::Str(id_str) => match Uuid::parse_str(&id_str.to_string()) {
Ok(id) => id,
Err(_) => {
self.send_error_response(
&cv.get_id(),
CommunicationType::error_invalid_call_id,
)
.await;
return;
}
},
_ => {
self.send_error_response(&cv.get_id(), CommunicationType::error_no_call_id)
.await;
return;
}
};
let invited = call_manager::add_invite(call_id, self.user_id, receiver_id as u64).await;
if !invited {
self.send_error_response(&cv.get_id(), CommunicationType::error_invalid_call_id)
.await;
return;
}
// Find target RhoConnection
let target_rho = match rho_manager::get_rho_con_for_user(receiver_id).await {
Some(rho) => rho,
_ => {
self.send_error_response(&cv.get_id(), CommunicationType::error)
.await;
return;
}
};
// Get sender user ID
let sender_id = self.get_user_id();
// Create and send call distribution message
let forward = CommunicationValue::new(CommunicationType::call_invite)
.with_receiver(receiver_id as u64)
.with_sender(sender_id)
.add_data(DataTypes::call_id, DataValue::Str(call_id.to_string()))
.add_data(
DataTypes::receiver_id,
DataValue::Str(receiver_id.to_string()),
)
.add_data(DataTypes::sender_id, DataValue::Str(sender_id.to_string()));
target_rho.message_to_client(forward).await;
let response = CommunicationValue::new(CommunicationType::success).with_id(cv.get_id());
self.send_message(&response).await;
}
/// Handle get call request
async fn handle_get_call(self: Arc<Self>, cv: CommunicationValue) {
let user_id = self.get_user_id();
let call_id = match cv.get_data(DataTypes::call_id) {
DataValue::Str(id_str) => match Uuid::parse_str(&id_str.to_string()) {
Ok(id) => id,
Err(_) => {
self.send_error_response(&cv.get_id(), CommunicationType::error)
.await;
return;
}
},
_ => {
self.send_error_response(&cv.get_id(), CommunicationType::error)
.await;
return;
}
};
if let Some(token) = call_manager::get_call_token(user_id, call_id).await {
let response = CommunicationValue::new(CommunicationType::call_token)
.with_id(cv.get_id())
.with_receiver(user_id)
.add_data(DataTypes::call_token, DataValue::Str(token.to_string()));
self.send_message(&response).await;
} else {
self.send_error_response(&cv.get_id(), CommunicationType::error)
.await;
return;
}
}
async fn handle_call_timeout_user(self: Arc<Self>, cv: CommunicationValue) {
let call_id =
Uuid::from_str(cv.get_data(DataTypes::call_id).as_str().unwrap_or("")).unwrap();
let user_id = cv.get_data(DataTypes::user_id).as_number().unwrap_or(0);
let untill = cv.get_data(DataTypes::untill).as_number().unwrap_or(0);
let call = call_manager::get_call(call_id).await;
if let Some(call) = call {
if call
.get_caller(self.get_user_id())
.await
.unwrap()
.has_admin()
{
call.get_caller(user_id as u64)
.await
.unwrap()
.set_timeout(untill)
.await;
}
}
}
async fn handle_call_disconnect_user(self: Arc<Self>, cv: CommunicationValue) {
let call_id =
Uuid::from_str(cv.get_data(DataTypes::call_id).as_str().unwrap_or("")).unwrap();
let user_id = cv.get_data(DataTypes::user_id).as_number().unwrap_or(0);
let call = call_manager::get_call(call_id).await;
if let Some(call) = call {
if call
.get_caller(self.get_user_id())
.await
.unwrap()
.has_admin()
{
call.remove_caller(user_id as u64).await;
}
}
}
/// Send error response
async fn send_error_response(self: Arc<Self>, message_id: &u32, error_type: CommunicationType) {
let error = CommunicationValue::new(error_type).with_id(*message_id);
self.send_message(&error).await;
}
/// Close the connection
pub async fn close(&self) {
let mut is_open_guard = self.is_open.write().await;
if !*is_open_guard {
return;
}
*is_open_guard = false;
let _ = self.sender.close();
}
#[allow(dead_code)]
/// Set interested users list
pub async fn set_interested_users(self: Arc<Self>, interested_ids: Vec<i64>) {
let mut interested_guard = self.interested_users.write().await;
*interested_guard = interested_ids;
}
#[allow(dead_code)]
pub async fn get_interested_users(self: Arc<Self>) -> Vec<i64> {
let interested_guard = self.interested_users.read().await;
interested_guard.clone()
}
#[allow(dead_code)]
/// Check if interested in a user and send notification
pub async fn are_you_interested(self: Arc<Self>, user_id: i64) {
let interested_guard = self.clone().get_interested_users().await;
if interested_guard.contains(&user_id) {
let notification = CommunicationValue::new(CommunicationType::client_changed)
.add_data(DataTypes::user_id, DataValue::Str(user_id.to_string()))
.add_data(DataTypes::user_state, DataValue::Str("online".to_string()));
self.send_message(&notification).await;
}
}
/// Handle connection close
pub async fn handle_close(&self) {
// TODO delete temp user
}
}
// Implement Clone to make it easier to work with Arc<AnonymousClientConnection>
impl Clone for AnonymousClientConnection {
fn clone(&self) -> Self {
Self {
sender: Arc::clone(&self.sender),
receiver: Arc::clone(&self.receiver),
user_id: self.user_id,
ping: Arc::clone(&self.ping),
interested_users: Arc::clone(&self.interested_users),
is_open: Arc::clone(&self.is_open),
user_name: Arc::clone(&self.user_name),
display_name: Arc::clone(&self.display_name),
avatar: Arc::clone(&self.avatar),
}
}
}

View file

@ -0,0 +1,53 @@
use dashmap::DashMap;
use once_cell::sync::Lazy;
use rand::Rng;
use rand::seq::SliceRandom;
use std::sync::Arc;
use crate::anonymous_clients::anonymous_client_connection::AnonymousClientConnection;
static ANONYMOUS_USERS: Lazy<DashMap<u64, Arc<AnonymousClientConnection>>> =
Lazy::new(|| DashMap::new());
#[allow(dead_code)]
pub async fn add_anonymous_user(connection: Arc<AnonymousClientConnection>) {
ANONYMOUS_USERS.insert(connection.get_user_id(), connection);
}
#[allow(dead_code)]
pub async fn remove_anonymous_user(user_id: u64) {
ANONYMOUS_USERS.remove(&user_id);
}
pub async fn get_anonymous_user(user_id: u64) -> Option<Arc<AnonymousClientConnection>> {
ANONYMOUS_USERS.get(&user_id).map(|c| c.clone())
}
pub async fn get_anonymous_user_by_name(
username: String,
) -> Option<Arc<AnonymousClientConnection>> {
for user_conn in ANONYMOUS_USERS
.iter()
.map(|ref_multi| ref_multi.value().clone())
{
if user_conn.get_user_name().await == username {
return Some(user_conn);
}
}
return None;
}
// TODO: implement check if taken
pub fn generate_username() -> String {
let adjectives = ["Swift", "Clever", "Brave", "Sneaky", "Fierce"];
let nouns = ["Tiger", "Eagle", "Shark", "Wolf", "Dragon"];
let mut rng = rand::thread_rng();
let adj = adjectives.choose(&mut rng).unwrap();
let noun = nouns.choose(&mut rng).unwrap();
let number: u16 = rng.gen_range(0..10000);
format!("{}{}{}", adj, noun, number)
}

View file

@ -0,0 +1,2 @@
pub mod anonymous_client_connection;
pub mod anonymous_manager;

106
src/calls/call_group.rs Normal file
View file

@ -0,0 +1,106 @@
use ttp_core::{CommunicationType, CommunicationValue, DataTypes, DataValue};
use std::{env, sync::Arc, time::Duration};
use tokio::sync::RwLock;
use uuid::Uuid;
use crate::{
calls::{call_util, caller::Caller},
omega::omega_connection::get_omega_connection,
};
pub struct CallGroup {
pub call_id: Uuid,
pub members: RwLock<Vec<Arc<Caller>>>,
pub show: RwLock<bool>,
pub anonymous_joining: RwLock<bool>,
pub short_link: RwLock<Option<String>>,
}
impl CallGroup {
pub fn new(call_id: Uuid, user: Arc<Caller>) -> Self {
CallGroup {
call_id,
members: RwLock::new(vec![user]),
show: RwLock::new(true),
anonymous_joining: RwLock::new(false),
short_link: RwLock::new(None),
}
}
pub async fn get_caller(&self, user_id: u64) -> Option<Arc<Caller>> {
self.members
.read()
.await
.iter()
.find(|caller| caller.user_id == user_id)
.cloned()
}
pub async fn is_anonymous(&self) -> bool {
*self.anonymous_joining.read().await
}
pub async fn set_anonymous_joining(&self, enable: bool) {
*self.anonymous_joining.write().await = enable;
let _ = call_util::set_room_metadata(
self.call_id,
format!("{{\"anonymous_joining\": \"{}\"}}", enable),
)
.await;
if self.short_link.read().await.is_none() {
let long_link = format!(
"https://app.tensamin.net/call/anonymous?call_id={}&omikron_id={}",
self.call_id,
env::var("ID")
.unwrap_or("0".to_string())
.parse::<i64>()
.unwrap_or(0),
);
let response_cv = get_omega_connection()
.await_response(
&CommunicationValue::new(CommunicationType::shorten_link)
.add_data(DataTypes::link, DataValue::Str(long_link)),
Some(Duration::from_secs(20)),
)
.await;
if let Ok(response) = response_cv {
*self.short_link.write().await = Some(
response
.get_data(DataTypes::link)
.as_str()
.unwrap()
.to_string(),
);
log::info!(
"Shortened link for call {} is {}",
self.call_id,
self.short_link.read().await.as_ref().unwrap()
);
}
}
}
pub async fn create_anonymous_token(&self, user_id: u64) -> Option<String> {
if self.is_anonymous().await {
if let Ok(token) = call_util::create_token(user_id, self.call_id, false) {
return Some(token);
}
}
None
}
pub async fn remove_caller(&self, user_id: u64) {
let _ = call_util::remove_participant(self.call_id, user_id).await;
self.members
.write()
.await
.retain(|caller| caller.user_id != user_id);
}
pub async fn get_short_link(self: Arc<Self>) -> Option<String> {
self.short_link.read().await.clone()
}
}

97
src/calls/call_manager.rs Normal file
View file

@ -0,0 +1,97 @@
use dashmap::DashMap;
use once_cell::sync::Lazy;
use std::sync::Arc;
use uuid::Uuid;
use crate::calls::{call_group::CallGroup, caller::Caller};
pub static CALL_GROUPS: Lazy<DashMap<Uuid, Arc<CallGroup>>> = Lazy::new(|| DashMap::new());
#[allow(dead_code)]
pub async fn get_call_invites(user_id: u64) -> Vec<Arc<Caller>> {
let mut callers = Vec::new();
for (_, cg) in CALL_GROUPS.clone().into_iter() {
let members = cg.members.read().await;
for member in members.iter() {
if member.user_id == user_id {
callers.push(member.clone());
}
}
}
callers
}
pub async fn get_call(call_id: Uuid) -> Option<Arc<CallGroup>> {
if let Some(b) = CALL_GROUPS.get(&call_id) {
Some(b.clone())
} else {
None
}
}
pub async fn get_call_groups(user_id: u64) -> Vec<Arc<CallGroup>> {
let mut call_groups = Vec::new();
for (_, cg) in CALL_GROUPS.clone().into_iter() {
let is_member = {
let members = cg.members.read().await;
members.iter().any(|m| m.user_id == user_id)
};
if is_member {
call_groups.push(cg.clone());
}
}
call_groups
}
pub async fn get_call_token(user_id: u64, call_id: Uuid) -> Option<String> {
if let Some(cg) = CALL_GROUPS.get(&call_id) {
let mut members = cg.members.write().await;
if let Some(member) = members.iter().find(|m| m.user_id == user_id) {
return Some(member.create_token());
}
let new_caller = Arc::new(Caller::new(user_id, call_id, false));
let token = new_caller.create_token();
members.push(new_caller);
return Some(token);
}
if let Some(cg) = CALL_GROUPS.get(&call_id) {
let cg_clone = cg.clone();
let mut members = cg_clone.members.write().await;
if let Some(member) = members.iter().find(|m| m.user_id == user_id) {
return Some(member.create_token());
}
let new_caller = Arc::new(Caller::new(user_id, call_id, false));
let token = new_caller.create_token();
members.push(new_caller);
return Some(token);
}
let caller = Arc::new(Caller::new(user_id, call_id, true));
let call_group = CallGroup::new(call_id, caller.clone());
CALL_GROUPS.insert(call_id, Arc::new(call_group));
Some(caller.create_token())
}
pub async fn add_invite(call_id: Uuid, inviter_id: u64, invitee_id: u64) -> bool {
if let Some(cg) = CALL_GROUPS.get(&call_id) {
let mut members = cg.members.write().await;
let is_inviter_member = members.iter().any(|m| m.user_id == inviter_id);
if is_inviter_member {
if !members.iter().any(|m| m.user_id == invitee_id) {
members.push(Arc::new(Caller::new(invitee_id, call_id, false)));
}
return true;
}
}
false
}

151
src/calls/call_util.rs Normal file
View file

@ -0,0 +1,151 @@
use livekit_api::{
access_token::{self},
services::room::RoomClient,
};
use livekit_protocol::Room;
use std::env;
use std::str::FromStr;
use std::time::Duration;
use uuid::Uuid;
use crate::{calls::call_manager::CALL_GROUPS, log, log_err, util::logger::PrintType};
pub fn get_livekit() -> Result<(String, String, String), ()> {
let hostname = match env::var("LIVEKI_HOSTNAME") {
Ok(secret) => secret,
Err(_) => {
log_err!(0, PrintType::General, "LIVEKI_HOSTNAME not set!");
return Err(());
}
};
let api_key = match env::var("LIVEKIT_API_KEY") {
Ok(key) => key,
Err(_) => {
log_err!(0, PrintType::General, "LIVEKIT_API_KEY not set!");
return Err(());
}
};
let api_secret = match env::var("LIVEKIT_API_SECRET") {
Ok(secret) => secret,
Err(_) => {
log_err!(0, PrintType::General, "LIVEKIT_API_SECRET not set!");
return Err(());
}
};
Ok((hostname, api_key, api_secret))
}
pub fn create_token(user_id: u64, call_id: Uuid, has_admin: bool) -> Result<String, ()> {
let (_, api_key, api_secret) = get_livekit()?;
let token = access_token::AccessToken::with_api_key(&api_key, &api_secret)
.with_identity(&user_id.to_string())
.with_grants(access_token::VideoGrants {
room_join: true,
room_admin: has_admin,
room: call_id.to_string(),
..Default::default()
})
.with_metadata(&format!("{{\"isAdmin\":{}}}", has_admin))
.to_jwt();
if let Ok(token) = token {
Ok(token)
} else {
Err(())
}
}
#[allow(dead_code)]
pub async fn get_room(call_id: Uuid) -> Result<(RoomClient, Room), ()> {
if let Ok((hostname, api_key, api_secret)) = get_livekit() {
let room_service = RoomClient::with_api_key(&hostname, &api_key, &api_secret);
let rooms = room_service.list_rooms(Vec::new()).await;
if let Ok(rooms) = rooms {
for room in rooms {
if room.name == call_id.to_string() {
return Ok((room_service, room));
}
}
}
}
return Err(());
}
pub async fn remove_participant(call_id: Uuid, user_id: u64) -> Result<(), ()> {
if let Ok((hostname, api_key, api_secret)) = get_livekit() {
let room_service = RoomClient::with_api_key(&hostname, &api_key, &api_secret);
if let Ok(_) = room_service
.remove_participant(&call_id.to_string(), &user_id.to_string())
.await
{
return Ok(());
}
}
return Err(());
}
#[allow(dead_code)]
pub async fn get_room_metadata(call_id: Uuid) -> Result<String, ()> {
if let Ok((_, room)) = get_room(call_id).await {
Ok(room.metadata)
} else {
Err(())
}
}
pub async fn set_room_metadata(call_id: Uuid, metadata: String) -> Result<(), ()> {
if let Ok((hostname, api_key, api_secret)) = get_livekit() {
let room_service = RoomClient::with_api_key(&hostname, &api_key, &api_secret);
if let Ok(_) = room_service
.update_room_metadata(&call_id.to_string(), &metadata)
.await
{
return Ok(());
}
}
return Err(());
}
pub fn garbage_collect_calls() {
tokio::spawn(async move {
loop {
if let Ok((hostname, api_key, api_secret)) = get_livekit() {
let room_service = RoomClient::with_api_key(&hostname, &api_key, &api_secret);
clean_calls(room_service).await;
}
tokio::time::sleep(Duration::from_secs(2)).await;
}
});
}
pub async fn clean_calls(room_service: RoomClient) {
let rooms = room_service.list_rooms(Vec::new()).await.unwrap();
let mut call_ids: Vec<Uuid> = Vec::new();
let mut no_users: Vec<Uuid> = Vec::new();
for room in rooms {
if let Ok(id) = Uuid::from_str(&room.name) {
if room.num_participants == 0 {
no_users.push(id);
}
call_ids.push(id);
}
}
let size_pre = CALL_GROUPS.len();
for (id, _) in CALL_GROUPS.clone().into_iter() {
if !call_ids.contains(&id) {
CALL_GROUPS.remove(&id);
}
}
for (_, cg) in CALL_GROUPS.clone().into_iter() {
*cg.show.write().await = !no_users.contains(&cg.call_id);
}
let size_post = CALL_GROUPS.len();
if size_pre - size_post != 0 {
log!(
0,
PrintType::Call,
"Cleaned {} calls, {} remaining",
size_pre - size_post,
size_post
);
}
}

49
src/calls/caller.rs Normal file
View file

@ -0,0 +1,49 @@
use std::time::{SystemTime, UNIX_EPOCH};
use tokio::sync::RwLock;
use uuid::Uuid;
use crate::calls::call_util;
pub struct Caller {
pub user_id: u64,
pub call_id: Uuid,
pub has_admin: bool,
pub timeout: RwLock<i64>,
}
impl Caller {
pub fn new(user_id: u64, call_id: Uuid, has_admin: bool) -> Self {
Caller {
user_id,
call_id,
has_admin,
timeout: RwLock::new(0),
}
}
#[allow(dead_code)]
pub fn set_admin(&mut self, has_admin: bool) {
self.has_admin = has_admin;
}
pub fn has_admin(&self) -> bool {
self.has_admin
}
#[allow(dead_code)]
pub async fn is_timeouted(&self) -> bool {
*self.timeout.read().await
> SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_millis() as i64
}
pub async fn set_timeout(&self, timeout: i64) {
*self.timeout.write().await = timeout;
}
pub fn create_token(&self) -> String {
if let Ok(token) = call_util::create_token(self.user_id, self.call_id, self.has_admin()) {
token
} else {
String::new()
}
}
}

4
src/calls/mod.rs Normal file
View file

@ -0,0 +1,4 @@
pub mod call_group;
pub mod call_manager;
pub mod call_util;
pub mod caller;

1
src/data/mod.rs Normal file
View file

@ -0,0 +1 @@
pub mod user;

View file

@ -1,19 +1,25 @@
mod server;
mod sql;
mod transport;
mod anonymous_clients;
mod calls;
mod data;
mod omega;
mod rho;
mod util;
use crate::sql::sql::initialize_db;
use crate::sql::sql::print_users;
use crate::transport::omikron_connection;
use crate::util::crypto_helper::load_public_key;
use crate::util::crypto_helper::load_secret_key;
use crate::util::logger::PrintType;
use crate::util::logger::startup;
use std::env;
use dotenv::dotenv;
use once_cell::sync::Lazy;
use rustls::crypto::aws_lc_rs::default_provider;
use std::env;
use crate::{
calls::call_util::garbage_collect_calls,
omega::omega_connection::get_omega_connection,
rho::server::start,
util::{
crypto_helper::{load_public_key, load_secret_key},
logger::startup,
},
};
static PRIVATE_KEY: Lazy<String> = Lazy::new(|| env::var("PRIVATE_KEY").unwrap());
pub fn get_private_key() -> x448::Secret {
@ -32,34 +38,14 @@ async fn main() {
}
dotenv().ok();
startup();
log_in!("Incoming messages");
log_out!("Outgoing messages");
get_omega_connection();
tokio::spawn(async move {
match omikron_connection::start(9187).await {
Err(e) => log_err!(0, PrintType::General, "{:?}", e),
_ => {}
if let Err(e) = start(959).await {
log_err!(0, util::logger::PrintType::General, "{}", e);
}
});
log!("Started");
log!(" .env");
if let Err(e) = initialize_db().await {
log!("[FATAL] Database initialization failed: {}", e);
log!(
"[FATAL] Please ensure the database is running and the .env file is configured correctly."
);
return;
} else {
log!(" DB");
}
if let Err(e) = print_users().await {
log!("[ERROR] Failed to print users: {}", e);
} else {
log!(" Users");
}
let _ = server::server::start(9188).await;
garbage_collect_calls();
tokio::signal::ctrl_c().await.unwrap();
}

1
src/omega/mod.rs Normal file
View file

@ -0,0 +1 @@
pub mod omega_connection;

View file

@ -0,0 +1,741 @@
use crate::{
data::user::UserStatus,
get_private_key, log, log_cv_in, log_cv_out, log_err, log_in,
rho::rho_manager::{self, RHO_CONNECTIONS, connection_count},
util::{
crypto_helper::{decrypt_b64, secret_key_to_base64},
file_util::load_file_vec,
logger::PrintType,
},
};
use dashmap::DashMap;
use once_cell::sync::Lazy;
use std::{collections::HashMap, env, sync::Arc, time::Duration};
use tokio::{
sync::{Mutex, RwLock, mpsc, watch},
task::JoinHandle,
time::{Instant, sleep},
};
use ttp_core::{CommunicationType, CommunicationValue, DataTypes, DataValue, rand_u32};
use ttp_native::{Receiver, Sender};
use uuid::Uuid;
// ============================================================================
// Configuration
// ============================================================================
const OMEGA_HOST_DEFAULT: &str = "methanium.net";
const OMEGA_PORT_DEFAULT: u16 = 9187;
const RECONNECT_DELAY: Duration = Duration::from_secs(5);
const MAX_RECONNECT_DELAY: Duration = Duration::from_secs(300);
const CONNECTION_TIMEOUT: Duration = Duration::from_secs(10);
const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);
const TASK_CLEANUP_INTERVAL: Duration = Duration::from_secs(60);
const TASK_MAX_AGE: Duration = Duration::from_secs(60);
// ============================================================================
// Waiting Task System
// ============================================================================
pub struct WaitingTask {
pub task: Box<dyn Fn(Arc<OmegaConnection>, CommunicationValue) -> bool + Send + Sync>,
pub inserted_at: Instant,
}
pub static WAITING_TASKS: Lazy<DashMap<u32, WaitingTask>> = Lazy::new(DashMap::new);
pub fn start_task_cleanup_loop() {
tokio::spawn(async {
loop {
sleep(TASK_CLEANUP_INTERVAL).await;
WAITING_TASKS.retain(|_, v| v.inserted_at.elapsed() < TASK_MAX_AGE);
}
});
}
// ============================================================================
// Connection State
// ============================================================================
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum ConnectionState {
Disconnected,
Connecting,
Connected { identified: bool },
}
#[allow(unused_variables)]
impl ConnectionState {
pub fn is_connected(&self) -> bool {
match self {
ConnectionState::Connected { identified } => true,
_ => false,
}
}
#[allow(dead_code)]
pub fn is_identified(&self) -> bool {
match self {
ConnectionState::Connected { identified: true } => true,
_ => false,
}
}
}
// ============================================================================
// Omega Connection (Client-side with auto-reconnect)
// ============================================================================
#[allow(dead_code)]
pub struct OmegaConnection {
state: Arc<RwLock<ConnectionState>>,
sender: Arc<RwLock<Option<Arc<Sender>>>>,
connection_loop_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
host: String,
port: u16,
server_cert: Vec<u8>,
last_ping: Arc<Mutex<i64>>,
heartbeat_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
message_send_times: Arc<Mutex<HashMap<Uuid, Instant>>>,
pub connection_id: Uuid,
shutdown_tx: Arc<Mutex<Option<watch::Sender<bool>>>>,
// Track if we should reconnect on close
reconnect_on_close: Arc<RwLock<bool>>,
}
impl OmegaConnection {
pub fn new() -> Self {
Self::with_host(OMEGA_HOST_DEFAULT, OMEGA_PORT_DEFAULT)
}
pub fn with_host(host: &str, port: u16) -> Self {
// Load server certificate from default location
let server_cert =
load_file_vec("certs", "cert.pem").expect("Failed to load server certificate");
Self::with_host_and_cert(host, port, server_cert)
}
// New constructor that accepts certificate directly
pub fn with_host_and_cert(host: &str, port: u16, server_cert: Vec<u8>) -> Self {
let (shutdown_tx, _) = watch::channel(false);
OmegaConnection {
state: Arc::new(RwLock::new(ConnectionState::Disconnected)),
sender: Arc::new(RwLock::new(None)),
connection_loop_handle: Arc::new(Mutex::new(None)),
host: host.to_string(),
port,
server_cert,
last_ping: Arc::new(Mutex::new(-1)),
heartbeat_handle: Arc::new(Mutex::new(None)),
message_send_times: Arc::new(Mutex::new(HashMap::new())),
connection_id: Uuid::new_v4(),
shutdown_tx: Arc::new(Mutex::new(Some(shutdown_tx))),
reconnect_on_close: Arc::new(RwLock::new(true)),
}
}
// -------------------------------------------------------------------------
// Connection Management
// -------------------------------------------------------------------------
pub async fn start(self: Arc<Self>) {
// Cancel any existing connection loop
if let Some(handle) = self.connection_loop_handle.lock().await.take() {
handle.abort();
}
// Set reconnect flag
*self.reconnect_on_close.write().await = true;
let self_clone = self.clone();
let handle = tokio::spawn(async move {
self_clone.connection_loop().await;
});
*self.connection_loop_handle.lock().await = Some(handle);
}
#[allow(dead_code)]
pub async fn stop(&self) {
// Disable reconnection
*self.reconnect_on_close.write().await = false;
if let Some(tx) = self.shutdown_tx.lock().await.take() {
let _ = tx.send(true);
}
if let Some(handle) = self.connection_loop_handle.lock().await.take() {
handle.abort();
}
if let Some(handle) = self.heartbeat_handle.lock().await.take() {
handle.abort();
}
// Close sender if connected
if let Some(sender) = self.sender.read().await.as_ref() {
sender.close();
}
*self.state.write().await = ConnectionState::Disconnected;
*self.sender.write().await = None;
}
async fn connection_loop(self: Arc<Self>) {
let mut reconnect_delay = RECONNECT_DELAY;
let shutdown_rx = self.shutdown_tx.lock().await.as_ref().unwrap().subscribe();
let mut shutdown_rx = shutdown_rx;
loop {
if *shutdown_rx.borrow() {
log_in!(0, PrintType::Omega, "Connection loop shutting down");
break;
}
// Check if reconnection is enabled
if !*self.reconnect_on_close.read().await {
log_in!(0, PrintType::Omega, "Reconnection disabled, exiting loop");
break;
}
match self.clone().connect_once().await {
Ok(()) => {
// Connection closed gracefully, check if we should reconnect
if *self.reconnect_on_close.read().await {
log_err!(
0,
PrintType::Omega,
"Connection lost, reconnecting in {:?}...",
reconnect_delay
);
} else {
log_in!(0, PrintType::Omega, "Connection closed, not reconnecting");
break;
}
}
Err(e) => {
log_err!(
0,
PrintType::Omega,
"Connection failed: {}, retrying in {:?}...",
e,
reconnect_delay
);
}
}
tokio::select! {
_ = sleep(reconnect_delay) => {}
_ = shutdown_rx.changed() => {
if *shutdown_rx.borrow() {
break;
}
}
}
reconnect_delay = std::cmp::min(reconnect_delay * 2, MAX_RECONNECT_DELAY);
}
}
async fn connect_once(self: Arc<Self>) -> Result<(), String> {
*self.state.write().await = ConnectionState::Connecting;
let addr_str = format!("https://{}:{}", self.host, self.port);
let (sender, mut receiver) = ttp_native::client::connect(&addr_str, None)
.await
.map_err(|e| format!("Connection failed: {}", e))?;
log_in!(
0,
PrintType::Omega,
"QUIC connection established to {}",
addr_str
);
// Store sender
let sender_arc = Arc::new(sender);
*self.sender.write().await = Some(sender_arc.clone());
*self.state.write().await = ConnectionState::Connected { identified: false };
// Get handle for close monitoring
let sender_handle = sender_arc.handle().clone();
// Start read loop
let read_self = self.clone();
let read_handle = tokio::spawn(async move {
read_self.read_loop(&mut receiver, sender_handle).await;
});
// Send identification
self.send_identification().await;
// Start heartbeat
let heartbeat_self = self.clone();
let heartbeat_handle = tokio::spawn(async move {
heartbeat_self.heartbeat_loop().await;
});
*self.heartbeat_handle.lock().await = Some(heartbeat_handle);
// Wait for read loop to complete (connection closed)
let result = read_handle.await;
// Cleanup
*self.sender.write().await = None;
*self.state.write().await = ConnectionState::Disconnected;
if let Some(handle) = self.heartbeat_handle.lock().await.take() {
handle.abort();
}
match result {
Ok(()) => {
// Check if we should reconnect
if *self.reconnect_on_close.read().await {
Err("Connection closed, will reconnect".to_string())
} else {
Ok(())
}
}
Err(e) => Err(format!("Read loop error: {}", e)),
}
}
// -------------------------------------------------------------------------
// Identification Handshake
// -------------------------------------------------------------------------
async fn send_identification(&self) {
let id = rand_u32();
let omikron_id = env::var("ID")
.unwrap_or("0".to_string())
.parse::<i64>()
.unwrap_or(0);
let identify_msg = CommunicationValue::new(CommunicationType::identification)
.with_id(id)
.add_data(DataTypes::omikron_id, DataValue::Number(omikron_id));
WAITING_TASKS.insert(
id,
WaitingTask {
task: Box::new(|selfc, cv| {
if cv.is_type(CommunicationType::error_not_found) {
log_err!(
0,
PrintType::Omega,
"Identification failed: Omikron ID not found"
);
return false;
}
if !cv.is_type(CommunicationType::challenge) {
return false;
}
tokio::spawn(async move {
if let Err(e) = selfc.handle_challenge(cv).await {
log_err!(0, PrintType::Omega, "Challenge handling failed: {}", e);
}
});
true
}),
inserted_at: Instant::now(),
},
);
self.send_message(&identify_msg).await;
}
async fn handle_challenge(&self, cv: CommunicationValue) -> Result<(), String> {
let challenge = cv
.get_data(DataTypes::challenge)
.as_str()
.ok_or("Challenge not found")?;
let server_pub_key = cv
.get_data(DataTypes::public_key)
.as_str()
.ok_or("Public key not found")?;
let decrypted_challenge = decrypt_b64(
&secret_key_to_base64(&get_private_key()),
server_pub_key,
challenge,
)
.map_err(|e| format!("Decryption failed: {:?}", e))?;
let response_msg = CommunicationValue::new(CommunicationType::challenge_response)
.with_id(cv.get_id())
.add_data(DataTypes::challenge, DataValue::Str(decrypted_challenge));
let response_id = response_msg.get_id();
WAITING_TASKS.insert(
response_id,
WaitingTask {
task: Box::new(|selfc, final_cv| {
if !final_cv.is_type(CommunicationType::identification_response) {
log_err!(0, PrintType::Omega, "Expected identification_response");
return false;
}
let accepted = final_cv
.get_data(DataTypes::accepted)
.as_bool()
.unwrap_or(false);
if !accepted {
log_err!(0, PrintType::Omega, "Omega did not accept identification");
return false;
}
tokio::spawn(async move {
let mut state = selfc.state.write().await;
if let ConnectionState::Connected { identified: _ } = *state {
*state = ConnectionState::Connected { identified: true };
}
drop(state);
selfc.sync_client_iota_status().await;
});
log!(0, PrintType::Omega, "Successfully identified with Omega");
true
}),
inserted_at: Instant::now(),
},
);
self.send_message(&response_msg).await;
Ok(())
}
async fn sync_client_iota_status(self: Arc<Self>) {
let mut connected_iota_ids: Vec<DataValue> = Vec::new();
let mut connected_user_ids: Vec<DataValue> = Vec::new();
let rho_connections_reader = RHO_CONNECTIONS.read().await;
for iota_id in rho_connections_reader.keys() {
connected_iota_ids.push(DataValue::Number(*iota_id));
}
for rho in rho_connections_reader.values() {
for client_conn in rho.get_client_connections().await {
connected_user_ids.push(DataValue::Number(client_conn.get_user_id().await as i64));
}
}
drop(rho_connections_reader);
let sync_msg = CommunicationValue::new(CommunicationType::sync_client_iota_status)
.add_data(DataTypes::iota_ids, DataValue::Array(connected_iota_ids))
.add_data(DataTypes::user_ids, DataValue::Array(connected_user_ids))
.add_data(
DataTypes::rho_connections,
DataValue::Number(connection_count().await as i64),
);
self.send_message(&sync_msg).await;
}
// -------------------------------------------------------------------------
// Read Loop & Heartbeat
// -------------------------------------------------------------------------
async fn read_loop(
self: Arc<Self>,
receiver: &mut Receiver,
sender_handle: Arc<ttp_native::ConnectionHandle>,
) {
// Monitor both receiver and sender handle for close
let mut close_rx = sender_handle.subscribe_close();
loop {
tokio::select! {
result = receiver.receive() => {
match result {
Ok(cv) => {
if !cv.is_type(CommunicationType::pong) && !cv.is_type(CommunicationType::ping) {
log_cv_in!(PrintType::Omega, &cv);
}
if cv.is_type(CommunicationType::pong) || cv.is_type(CommunicationType::ping) {
self.handle_pong(&cv).await;
continue;
}
let msg_id = cv.get_id();
if let Some((_, task)) = WAITING_TASKS.remove(&msg_id) {
if (task.task)(self.clone(), cv.clone()) {
continue;
}
}
if cv.is_type(CommunicationType::iota_user_data) {
if let DataValue::Array(users) = cv.get_data(DataTypes::user_ids) {
let mut user_ids: Vec<u64> = Vec::new();
for value in users {
if let DataValue::Number(user_id) = value {
user_ids.push(*user_id as u64);
}
}
let connections = crate::rho::rho_manager::RHO_CONNECTIONS.read().await;
if let Some(iota_id) = cv.get_data(DataTypes::iota_id).as_number() {
if let Some(rho) = connections.get(&iota_id) {
rho.get_iota_connection().set_user_ids(user_ids).await;
}
} else {
for rho in connections.values() {
rho.get_iota_connection().set_user_ids(user_ids.clone()).await;
}
}
}
}
}
Err(e) => {
log_err!(0, PrintType::Omega, "Receive error: {}", e);
break;
}
}
}
_ = close_rx.changed() => {
// Connection was closed by either side
if let Some(reason) = close_rx.borrow().clone() {
log_err!(0, PrintType::Omega, "Connection closed: {:?}", reason);
} else {
log_in!(0, PrintType::Omega, "Connection closed cleanly");
}
break;
}
}
}
}
async fn heartbeat_loop(self: Arc<Self>) {
loop {
sleep(HEARTBEAT_INTERVAL).await;
// Check if still connected
if !self.state.read().await.is_connected() {
break;
}
// Check if sender is closed
if let Some(sender) = self.sender.read().await.as_ref() {
if sender.is_closed() {
log_err!(0, PrintType::Omega, "Sender closed, stopping heartbeat");
break;
}
} else {
break;
}
self.send_ping().await;
}
}
async fn send_ping(&self) {
let ping = CommunicationValue::new(CommunicationType::ping).add_data(
DataTypes::send_time,
DataValue::Number(
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs() as i64,
),
);
self.send_message(&ping).await;
}
async fn handle_pong(&self, cv: &CommunicationValue) {
let timestamp = cv
.get_data(DataTypes::send_time)
.as_number()
.unwrap_or_else(|| {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs() as i64
});
*self.last_ping.lock().await = timestamp;
}
// -------------------------------------------------------------------------
// Public API
// -------------------------------------------------------------------------
pub async fn send_message(&self, cv: &CommunicationValue) {
if !cv.is_type(CommunicationType::pong) && !cv.is_type(CommunicationType::ping) {
log_cv_out!(PrintType::Omega, &cv);
}
let sender_guard = self.sender.read().await;
if let Some(sender) = sender_guard.as_ref() {
// Check if closed before sending
if sender.is_closed() {
log_err!(0, PrintType::Omega, "Cannot send: connection closed");
drop(sender_guard);
// Trigger reconnection by closing the connection state
if let Some(sender) = self.sender.write().await.take() {
sender.close();
}
return;
}
let sender_clone = Arc::clone(sender);
drop(sender_guard);
if let Err(e) = sender_clone.send(cv).await {
log_err!(0, PrintType::Omega, "Send failed: {}", e);
}
} else {
log_err!(0, PrintType::Omega, "Cannot send: not connected");
}
}
pub async fn await_connection(&self, timeout_duration: Option<Duration>) -> Result<(), String> {
if self.state.read().await.is_connected() {
return Ok(());
}
let timeout = timeout_duration.unwrap_or(CONNECTION_TIMEOUT);
let start = Instant::now();
loop {
if self.state.read().await.is_connected() {
return Ok(());
}
if start.elapsed() >= timeout {
return Err(format!(
"Connection not established within {} seconds",
timeout.as_secs()
));
}
sleep(Duration::from_millis(100)).await;
}
}
pub async fn await_response(
&self,
cv: &CommunicationValue,
timeout_duration: Option<Duration>,
) -> Result<CommunicationValue, String> {
self.await_connection(timeout_duration).await?;
let (tx, mut rx) = mpsc::channel(1);
let msg_id = cv.get_id();
WAITING_TASKS.insert(
msg_id,
WaitingTask {
task: Box::new(move |_, response_cv| {
let inner_tx = tx.clone();
tokio::spawn(async move {
let _ = inner_tx.send(response_cv).await;
});
true
}),
inserted_at: Instant::now(),
},
);
self.send_message(cv).await;
let timeout = timeout_duration.unwrap_or(Duration::from_secs(10));
match tokio::time::timeout(timeout, rx.recv()).await {
Ok(Some(response_cv)) => Ok(response_cv),
Ok(_) => Err("Channel closed".to_string()),
Err(_) => {
WAITING_TASKS.remove(&msg_id);
Err("Request timed out".to_string())
}
}
}
#[allow(dead_code)]
pub async fn is_connected(&self) -> bool {
self.state.read().await.is_connected()
}
#[allow(dead_code)]
pub async fn is_identified(&self) -> bool {
self.state.read().await.is_identified()
}
#[allow(dead_code)]
pub async fn close_iota(iota_id: i64) {
let cv = CommunicationValue::new(CommunicationType::iota_disconnected)
.add_data(DataTypes::iota_id, DataValue::Number(iota_id));
OMEGA_CONNECTION.send_message(&cv).await;
}
pub async fn client_changed(_iota_id: i64, user_id: i64, state: UserStatus) {
let msg_type = match state {
UserStatus::iota_offline => CommunicationType::user_disconnected,
UserStatus::user_offline => CommunicationType::user_disconnected,
_ => CommunicationType::user_connected,
};
let cv = CommunicationValue::new(msg_type)
.add_data(DataTypes::user_id, DataValue::Number(user_id));
OMEGA_CONNECTION.send_message(&cv).await;
}
pub async fn user_states(user_id: i64, user_ids: Vec<i64>) {
let user_ids = user_ids.iter().map(|v| DataValue::Number(*v)).collect();
let cv = CommunicationValue::new(CommunicationType::get_states)
.add_data(DataTypes::user_ids, DataValue::Array(user_ids));
let msg_id = cv.get_id();
WAITING_TASKS.insert(
msg_id,
WaitingTask {
task: Box::new(
move |_: Arc<OmegaConnection>, response: CommunicationValue| {
tokio::spawn(async move {
let rho = rho_manager::get_rho_con_for_user(user_id).await;
if let Some(rho) = rho {
for client in rho.get_client_connections_for_user(user_id).await {
client.send_message(&response).await;
}
}
});
true
},
),
inserted_at: Instant::now(),
},
);
OMEGA_CONNECTION.send_message(&cv).await;
}
}
// ============================================================================
// Global Instance
// ============================================================================
static OMEGA_CONNECTION: Lazy<Arc<OmegaConnection>> = Lazy::new(|| {
let conn = Arc::new(OmegaConnection::new());
// Start the connection manager immediately
let conn_clone = conn.clone();
tokio::spawn(async move {
conn_clone.start().await;
});
start_task_cleanup_loop();
conn
});
pub fn get_omega_connection() -> Arc<OmegaConnection> {
OMEGA_CONNECTION.clone()
}

View file

@ -0,0 +1,41 @@
use std::time::Duration;
use tokio::time::Instant;
use ttp_core::{CommunicationType, CommunicationValue, DataTypes, DataValue, rand_u32};
use crate::omega::omega_connection::OmegaConnection;
const PING_TIMEOUT: Duration = Duration::from_secs(30);
impl OmegaConnection {
pub async fn send_ping(&self) {
let id = rand_u32();
let send_time = Instant::now();
let mut message_send_times = self.message_send_times.lock().await;
message_send_times.retain(|_uuid, time| time.elapsed() < PING_TIMEOUT);
message_send_times.insert(id as i64, send_time);
self.send_ping_message(id).await;
}
pub async fn send_ping_message(&self, id: u32) {
let ping_message = CommunicationValue::new(CommunicationType::ping)
.with_id(id)
.add_data(
DataTypes::last_ping,
DataValue::Number(self.last_ping.lock().await.unwrap()),
);
self.send_message(&ping_message).await;
}
/// Handles incoming pong and calculates latency
pub async fn handle_pong(&self, cv: &CommunicationValue, _log: bool) {
let id = cv.get_id();
let mut message_send_times = self.message_send_times.lock().await;
if let Some(send_time) = message_send_times.remove(&(id as i64)) {
let ping = Instant::now().duration_since(send_time).as_millis() as i64;
*self.last_ping.lock().await = ping;
}
}
}

View file

@ -0,0 +1,583 @@
use crate::anonymous_clients::anonymous_manager;
use crate::calls::{call_manager, call_util};
use crate::omega::omega_connection::get_omega_connection;
use crate::rho::connection::GeneralConnection;
use crate::rho::{rho_connection::RhoConnection, rho_manager};
use crate::util::logger::PrintType;
use crate::{data::user::UserStatus, omega::omega_connection::OmegaConnection};
use crate::{log_cv_in, log_cv_out, log_err, log_in, log_out};
use std::str::FromStr;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::sync::RwLock;
use ttp_core::{CommunicationType, CommunicationValue, DataTypes, DataValue};
use ttp_native::{Receiver, Sender};
use uuid::Uuid;
pub struct ClientConnection {
pub user_id: u64,
pub sender: Arc<Sender>,
pub receiver: Arc<Receiver>,
pub ping: Arc<RwLock<i64>>,
pub_key: Arc<RwLock<Option<Vec<u8>>>>,
pub rho_connection: Arc<RwLock<Option<Arc<RhoConnection>>>>,
pub interested_users: Arc<RwLock<Vec<i64>>>,
is_open: Arc<RwLock<bool>>,
}
impl ClientConnection {
pub async fn from_general(general: Arc<GeneralConnection>, user_id: u64) -> Arc<Self> {
Arc::new(Self {
ping: Arc::new(RwLock::new(0)),
pub_key: Arc::new(RwLock::new(None)),
rho_connection: general.rho_connection.clone(),
interested_users: Arc::new(RwLock::new(Vec::new())),
is_open: Arc::new(RwLock::new(true)),
sender: general.sender.clone(),
receiver: general.receiver.clone(),
user_id: user_id,
})
}
pub fn start(self: Arc<Self>) {
let self_clone = self.clone();
tokio::spawn(async move {
while let Ok(cv) = self_clone.receiver.receive().await {
self_clone.clone().handle_message(cv).await;
}
self_clone.handle_close().await;
});
let self_clone2 = self.clone();
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(50)).await;
if self_clone2.get_rho_connection().await.is_none() {
self_clone2
.send_error_response(0, CommunicationType::error)
.await;
}
});
}
/// Get the user ID
pub async fn get_user_id(&self) -> u64 {
self.user_id
}
/// Get current ping
pub async fn get_ping(&self) -> i64 {
*self.ping.read().await
}
/// Get RhoConnection if available
pub async fn get_rho_connection(&self) -> Option<Arc<RhoConnection>> {
self.rho_connection.read().await.clone()
}
/// Send a CommunicationValue to the client
pub async fn send_message(self: Arc<Self>, cv: &CommunicationValue) {
if !*self.is_open.read().await {
log_out!(
self.user_id as i64,
PrintType::Client,
"Attempted to send message to a closed connection."
);
return;
}
if !cv.is_type(CommunicationType::pong) && !cv.is_type(CommunicationType::ping) {
log_cv_out!(PrintType::Client, &cv);
}
let _ = self.sender.send(&cv).await;
}
/// Handle incoming message from client
pub async fn handle_message(self: Arc<Self>, cv: CommunicationValue) {
tokio::spawn(async move {
if cv.is_type(CommunicationType::ping) {
self.handle_ping(cv).await;
return;
}
log_cv_in!(PrintType::Client, cv);
// Handle client status changes
if cv.is_type(CommunicationType::client_changed) {
self.handle_client_changed(cv).await;
return;
}
// Handle call invites
if cv.is_type(CommunicationType::call_invite) {
self.handle_call_invite(cv).await;
return;
}
// Handle get call requests
if cv.is_type(CommunicationType::call_token) {
self.handle_get_call(cv).await;
return;
}
if cv.is_type(CommunicationType::call_disconnect_user) {
self.handle_call_disconnect_user(cv).await;
return;
}
if cv.is_type(CommunicationType::call_timeout_user) {
self.handle_call_timeout_user(cv).await;
return;
}
if cv.is_type(CommunicationType::call_set_anonymous_joining) {
self.handle_call_set_anonymous_joining(cv).await;
return;
}
if cv.is_type(CommunicationType::get_user_data) {
if let Some(anonymous) = {
if let Some(user_id) = cv.get_data(DataTypes::user_id).as_number() {
anonymous_manager::get_anonymous_user(user_id as u64).await
} else if let Some(username) = cv.get_data(DataTypes::username).as_str() {
anonymous_manager::get_anonymous_user_by_name(username.to_string()).await
} else {
None
}
} {
let response = CommunicationValue::new(CommunicationType::get_user_data)
.with_id(cv.get_id())
.add_data(
DataTypes::username,
DataValue::Str(anonymous.get_user_name().await),
)
.add_data(
DataTypes::user_id,
DataValue::Number(anonymous.get_user_id() as i64),
)
.add_data(
DataTypes::display,
DataValue::Str(anonymous.get_display_name().await),
)
.add_data(
DataTypes::avatar,
DataValue::Str(anonymous.get_avatar().await),
)
.add_data(DataTypes::user_state, DataValue::Str("online".to_string()));
self.send_message(&response).await;
return;
}
}
if cv.is_type(CommunicationType::change_user_data)
|| cv.is_type(CommunicationType::read_notification)
|| cv.is_type(CommunicationType::get_notifications)
|| cv.is_type(CommunicationType::get_user_data)
|| cv.is_type(CommunicationType::get_iota_data)
|| cv.is_type(CommunicationType::delete_user)
{
let sender = self.get_user_id().await;
self.handle_omega_forward(cv.with_sender(sender as u64))
.await;
return;
}
// Forward other messages to Iota
self.forward_to_iota(cv).await;
});
}
async fn handle_omega_forward(self: Arc<Self>, cv: CommunicationValue) {
let client_for_closure = self.clone();
tokio::spawn(async move {
let response_cv = get_omega_connection()
.await_response(&cv.with_sender(self.user_id), Some(Duration::from_secs(20)))
.await;
if let Ok(response_cv) = response_cv {
client_for_closure.send_message(&response_cv).await;
}
});
}
/// Handle ping message
async fn handle_ping(self: Arc<Self>, cv: CommunicationValue) {
// Update our ping if provided
if let DataValue::Number(last_ping) = cv.get_data(DataTypes::last_ping) {
let current = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_millis();
let mut ping_guard = self.ping.write().await;
*ping_guard = current as i64 - last_ping;
}
// Get Iota ping from RhoConnection
let iota_ping = if let Some(rho_conn) = self.get_rho_connection().await {
rho_conn.get_iota_connection().get_ping().await
} else {
-1
};
// Send pong response
let response = CommunicationValue::new(CommunicationType::pong)
.with_id(cv.get_id())
.add_data(DataTypes::ping_iota, DataValue::Number(iota_ping));
self.send_message(&response).await;
}
/// Handle client status change
async fn handle_client_changed(self: Arc<Self>, cv: CommunicationValue) {
let user_id = self.get_user_id().await;
if let DataValue::Str(_status_str) = cv.get_data(DataTypes::user_state) {
let user_status = UserStatus::user_online;
if let Some(rho_conn) = self.get_rho_connection().await {
OmegaConnection::client_changed(
rho_conn.get_iota_id().await as i64,
user_id as i64,
user_status,
)
.await;
}
}
}
/// Handle call invite
async fn handle_call_invite(self: Arc<Self>, cv: CommunicationValue) {
let receiver_id: i64 = cv.get_data(DataTypes::receiver_id).as_number().unwrap_or(0);
if receiver_id == 0 {
self.send_error_response(cv.get_id(), CommunicationType::error_no_user_id)
.await;
return;
}
let call_id = match cv.get_data(DataTypes::call_id) {
DataValue::Str(id_str) => match Uuid::parse_str(id_str.as_str()) {
Ok(id) => id,
Err(_) => {
self.send_error_response(cv.get_id(), CommunicationType::error_invalid_call_id)
.await;
return;
}
},
_ => {
self.send_error_response(cv.get_id(), CommunicationType::error_no_call_id)
.await;
return;
}
};
let invited = call_manager::add_invite(call_id, self.user_id, receiver_id as u64).await;
if !invited {
self.send_error_response(cv.get_id(), CommunicationType::error_invalid_call_id)
.await;
return;
}
// Find target RhoConnection
let target_rho = match rho_manager::get_rho_con_for_user(receiver_id).await {
Some(rho) => rho,
_ => {
self.send_error_response(cv.get_id(), CommunicationType::error)
.await;
return;
}
};
// Get sender user ID
let sender_id = self.get_user_id().await;
// Create and send call distribution message
let forward = CommunicationValue::new(CommunicationType::call_invite)
.with_receiver(receiver_id as u64)
.with_sender(sender_id as u64)
.add_data(DataTypes::call_id, DataValue::Str(call_id.to_string()))
.add_data(
DataTypes::receiver_id,
DataValue::Str(receiver_id.to_string()),
)
.add_data(DataTypes::sender_id, DataValue::Str(sender_id.to_string()));
target_rho.message_to_client(forward).await;
let response = CommunicationValue::new(CommunicationType::success).with_id(cv.get_id());
self.send_message(&response).await;
}
/// Handle get call request
async fn handle_get_call(self: Arc<Self>, cv: CommunicationValue) {
let user_id = self.get_user_id().await;
let call_id = match cv.get_data(DataTypes::call_id) {
DataValue::Str(id_str) => match Uuid::parse_str(id_str.as_str()) {
Ok(id) => id,
Err(_) => {
self.send_error_response(cv.get_id(), CommunicationType::error)
.await;
return;
}
},
_ => {
self.send_error_response(cv.get_id(), CommunicationType::error)
.await;
return;
}
};
if let Some(token) = call_manager::get_call_token(user_id, call_id).await {
let response = CommunicationValue::new(CommunicationType::call_token)
.with_id(cv.get_id())
.with_receiver(user_id as u64)
.add_data(DataTypes::call_token, DataValue::Str(token));
self.send_message(&response).await;
} else {
self.send_error_response(cv.get_id(), CommunicationType::error)
.await;
return;
}
}
async fn handle_call_timeout_user(self: Arc<Self>, cv: CommunicationValue) {
let call_id =
Uuid::from_str(cv.get_data(DataTypes::call_id).as_str().unwrap_or("")).unwrap();
let user_id = cv.get_data(DataTypes::user_id).as_number().unwrap_or(0);
let untill = cv.get_data(DataTypes::untill).as_number().unwrap_or(0);
let call = call_manager::get_call(call_id).await;
if let Some(call) = call {
if call
.get_caller(self.get_user_id().await)
.await
.unwrap()
.has_admin()
{
let _ = call_util::remove_participant(call_id, user_id as u64).await;
call.get_caller(user_id as u64)
.await
.unwrap()
.set_timeout(untill)
.await;
}
}
}
async fn handle_call_disconnect_user(self: Arc<Self>, cv: CommunicationValue) {
let call_id =
Uuid::from_str(cv.get_data(DataTypes::call_id).as_str().unwrap_or("")).unwrap();
let user_id = cv.get_data(DataTypes::user_id).as_number().unwrap_or(0);
let call = call_manager::get_call(call_id).await;
if let Some(call) = call {
if call
.get_caller(self.get_user_id().await)
.await
.unwrap()
.has_admin()
{
call.remove_caller(user_id as u64).await;
}
}
}
async fn handle_call_set_anonymous_joining(self: Arc<Self>, cv: CommunicationValue) {
let call_id =
Uuid::from_str(cv.get_data(DataTypes::call_id).as_str().unwrap_or("")).unwrap();
let enable = cv.get_data(DataTypes::enabled).as_bool().unwrap_or(true);
let call = call_manager::get_call(call_id).await;
let mut short_link = None;
if let Some(call) = call {
if call
.get_caller(self.get_user_id().await)
.await
.unwrap()
.has_admin()
{
call.set_anonymous_joining(enable).await;
}
short_link = call.get_short_link().await;
}
let mut response_cv =
CommunicationValue::new(CommunicationType::call_set_anonymous_joining)
.with_id(cv.get_id())
.add_data(DataTypes::call_id, DataValue::Str(call_id.to_string()))
.add_data(DataTypes::enabled, DataValue::Bool(enable));
if let Some(short_link) = short_link {
response_cv = response_cv.add_data(DataTypes::link, DataValue::Str(short_link));
}
self.send_message(&response_cv).await;
}
/// Forward message to Iota
async fn forward_to_iota(self: Arc<Self>, cv: CommunicationValue) {
let sender_user_id = self.get_user_id().await;
let msg_id = cv.get_id();
let msg_type = cv.get_type();
log_in!(
sender_user_id as i64,
PrintType::Client,
"Forwarding client->iota: sender={} type={:?} id={} receiver={}",
sender_user_id,
msg_type,
msg_id,
cv.get_receiver()
);
if cv.is_type(CommunicationType::add_conversation)
&& cv
.get_data(DataTypes::chat_partner_id)
.as_number()
.is_none()
{
let chat_partner_name = cv
.get_data(DataTypes::chat_partner_name)
.as_str()
.unwrap_or("")
.to_string();
if anonymous_manager::get_anonymous_user_by_name(chat_partner_name.to_string())
.await
.is_some()
{
self.send_error_response(cv.get_id(), CommunicationType::error_anonymous)
.await;
return;
}
let load_uuid_response = get_omega_connection()
.await_response(
&CommunicationValue::new(CommunicationType::get_user_data)
.with_id(cv.clone().get_id())
.add_data(
DataTypes::username,
DataValue::Str(chat_partner_name.clone()),
),
Some(Duration::from_secs(20)),
)
.await;
let chat_partner_id = {
if let Ok(load_uuid_response) = load_uuid_response {
load_uuid_response.get_data(DataTypes::user_id).clone()
} else {
DataValue::Null
}
};
if let Some(rho_conn) = self.get_rho_connection().await {
let iota_id = rho_conn.get_iota_id().await;
log_in!(
sender_user_id as i64,
PrintType::Client,
"Resolved rho for add_conversation: sender={} -> iota_id={} id={}",
sender_user_id,
iota_id,
msg_id
);
let updated_cv = cv
.with_sender(sender_user_id as u64)
.add_data(DataTypes::chat_partner_id, chat_partner_id);
rho_conn.message_to_iota(updated_cv).await;
} else {
log_err!(
sender_user_id as i64,
PrintType::Client,
"No rho/iota mapping found for add_conversation sender={} type={:?} id={}",
sender_user_id,
msg_type,
msg_id
);
}
return;
}
if let Some(rho_conn) = self.get_rho_connection().await {
let iota_id = rho_conn.get_iota_id().await;
log_in!(
sender_user_id as i64,
PrintType::Client,
"Resolved rho for forward: sender={} -> iota_id={} type={:?} id={}",
sender_user_id,
iota_id,
msg_type,
msg_id
);
let updated_cv = cv.with_sender(sender_user_id as u64);
rho_conn.message_to_iota(updated_cv).await;
} else {
log_err!(
sender_user_id as i64,
PrintType::Client,
"No rho/iota mapping found for sender={} type={:?} id={}",
sender_user_id,
msg_type,
msg_id
);
self.send_error_response(msg_id, CommunicationType::error)
.await;
}
}
/// Send error response
async fn send_error_response(self: Arc<Self>, message_id: u32, error_type: CommunicationType) {
let error = CommunicationValue::new(error_type).with_id(message_id);
self.send_message(&error).await;
}
/// Close the connection
pub async fn close(&self) {
let mut is_open_guard = self.is_open.write().await;
if !*is_open_guard {
return;
}
*is_open_guard = false;
let _ = self.sender.close();
}
/// Set interested users list
pub async fn set_interested_users(self: Arc<Self>, interested_ids: Vec<i64>) {
let mut interested_guard = self.interested_users.write().await;
*interested_guard = interested_ids;
}
#[allow(dead_code)]
pub async fn get_interested_users(self: Arc<Self>) -> Vec<i64> {
let interested_guard = self.interested_users.read().await;
interested_guard.clone()
}
/// Check if interested in a user and send notification
#[allow(dead_code)]
pub async fn are_you_interested(self: Arc<Self>, user_id: i64) {
let interested_guard = self.clone().get_interested_users().await;
if interested_guard.contains(&user_id) {
let notification = CommunicationValue::new(CommunicationType::client_changed)
.add_data(DataTypes::user_id, DataValue::Str(user_id.to_string()))
.add_data(DataTypes::user_state, DataValue::Str("online".to_string()));
self.send_message(&notification).await;
}
}
/// Handle connection close
pub async fn handle_close(&self) {
let user_id = self.get_user_id().await;
if let Some(rho_conn) = rho_manager::get_rho_con_for_user(user_id as i64).await {
rho_conn
.close_client_connection(Arc::new(self.clone()))
.await;
}
}
}
// Implement Clone to make it easier to work with Arc<ClientConnection>
impl Clone for ClientConnection {
fn clone(&self) -> Self {
Self {
sender: Arc::clone(&self.sender),
receiver: Arc::clone(&self.receiver),
user_id: self.user_id,
ping: Arc::clone(&self.ping),
pub_key: Arc::clone(&self.pub_key),
rho_connection: Arc::clone(&self.rho_connection),
interested_users: Arc::clone(&self.interested_users),
is_open: Arc::clone(&self.is_open),
}
}
}

358
src/rho/connection.rs Executable file
View file

@ -0,0 +1,358 @@
use rand::{Rng, distributions::Alphanumeric};
use std::{sync::Arc, time::Duration};
use tokio::sync::RwLock;
use ttp_core::{CommunicationType, CommunicationValue, DataTypes, DataValue};
use ttp_native::{Receiver, Sender};
use crate::{
anonymous_clients::anonymous_client_connection::AnonymousClientConnection,
get_private_key, get_public_key, log_cv_in, log_cv_out, log_err, log_in, log_out,
omega::omega_connection::get_omega_connection,
rho::{
client_connection::ClientConnection, iota_connection::IotaConnection,
rho_connection::RhoConnection, rho_manager,
},
util::{
crypto_helper::{load_public_key, public_key_to_base64},
crypto_util::{DataFormat, SecurePayload},
logger::PrintType,
},
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[allow(dead_code)]
pub enum ConnectionKind {
Client,
Iota,
AnonymousClient,
Phi,
}
pub struct GeneralConnection {
pub sender: Arc<Sender>,
pub receiver: Arc<Receiver>,
identified: Arc<RwLock<bool>>,
challenged: Arc<RwLock<bool>>,
challenge: Arc<RwLock<String>>,
connection_kind: Arc<RwLock<Option<ConnectionKind>>>,
pub rho_connection: Arc<RwLock<Option<Arc<RhoConnection>>>>,
id: Arc<RwLock<u64>>,
pub_key: Arc<RwLock<Option<Vec<u8>>>>,
}
impl GeneralConnection {
pub fn new(sender: Sender, receiver: Receiver) -> Arc<Self> {
Arc::new(Self {
sender: Arc::new(sender),
receiver: Arc::new(receiver),
identified: Arc::new(RwLock::new(false)),
challenged: Arc::new(RwLock::new(false)),
challenge: Arc::new(RwLock::new(String::new())),
connection_kind: Arc::new(RwLock::new(None)),
rho_connection: Arc::new(RwLock::new(None)),
id: Arc::new(RwLock::new(0)),
pub_key: Arc::new(RwLock::new(None)),
})
}
}
impl GeneralConnection {
pub async fn handle(self: Arc<Self>) {
log_in!(0, PrintType::General, "General connection handler started");
loop {
let cv = match self.receiver.receive().await {
Ok(v) => v,
Err(_) => {
break;
}
};
log_cv_in!(cv);
if !*self.identified.read().await {
self.handle_identification(cv).await;
continue;
}
if !*self.challenged.read().await {
self.handle_challenge_response(cv).await;
}
if *self.challenged.read().await {
let self_clone = self.clone();
tokio::spawn(async move {
self_clone.migrate().await;
});
break;
}
}
log_out!(0, PrintType::General, "General connection handler stopped");
}
async fn handle_identification(self: &Arc<Self>, cv: CommunicationValue) {
if !cv.is_type(CommunicationType::identification) {
return;
}
if let DataValue::Number(iota_id) = cv.get_data(DataTypes::iota_id) {
*self.id.write().await = *iota_id as u64;
*self.connection_kind.write().await = Some(ConnectionKind::Iota);
let get_pub_key_msg = CommunicationValue::new(CommunicationType::get_iota_data)
.add_data(DataTypes::iota_id, DataValue::Number(*iota_id));
let response_cv = get_omega_connection()
.await_response(&get_pub_key_msg, Some(Duration::from_secs(20)))
.await;
let response_cv = match response_cv {
Ok(r) => r,
Err(_) => {
return;
}
};
let base64_pub = response_cv
.get_data(DataTypes::public_key)
.as_str()
.unwrap_or("");
let pub_key = match load_public_key(base64_pub) {
Some(pk) => pk,
None => {
return;
}
};
*self.pub_key.write().await = Some(pub_key.as_bytes().to_vec());
let challenge: String = rand::thread_rng()
.sample_iter(&Alphanumeric)
.take(32)
.map(char::from)
.collect();
*self.challenge.write().await = challenge.clone();
*self.identified.write().await = true;
let encrypted_challenge =
SecurePayload::new(challenge.as_bytes(), DataFormat::Raw, get_private_key())
.unwrap()
.encrypt_x448(pub_key)
.unwrap()
.export(DataFormat::Base64);
let response = CommunicationValue::new(CommunicationType::challenge)
.with_id(cv.get_id())
.add_data(
DataTypes::public_key,
DataValue::Str(public_key_to_base64(&get_public_key())),
)
.add_data(DataTypes::challenge, DataValue::Str(encrypted_challenge));
log_cv_out!(response);
let _ = self.sender.send(&response).await;
} else if let DataValue::Number(user_id) = cv.get_data(DataTypes::user_id) {
*self.id.write().await = *user_id as u64;
*self.connection_kind.write().await = Some(ConnectionKind::Client);
let get_pub_key_msg = CommunicationValue::new(CommunicationType::get_user_data)
.add_data(DataTypes::user_id, DataValue::Number(*user_id));
let response_cv = get_omega_connection()
.await_response(&get_pub_key_msg, Some(Duration::from_secs(20)))
.await;
let response_cv = match response_cv {
Ok(r) => r,
Err(_) => {
return;
}
};
let base64_pub = response_cv
.get_data(DataTypes::public_key)
.as_str()
.unwrap_or("");
let pub_key = match load_public_key(base64_pub) {
Some(pk) => pk,
None => {
return;
}
};
*self.pub_key.write().await = Some(pub_key.as_bytes().to_vec());
let challenge: String = rand::thread_rng()
.sample_iter(&Alphanumeric)
.take(32)
.map(char::from)
.collect();
*self.challenge.write().await = challenge.clone();
*self.identified.write().await = true;
let encrypted_challenge =
SecurePayload::new(challenge.as_bytes(), DataFormat::Raw, get_private_key())
.unwrap()
.encrypt_x448(pub_key)
.unwrap()
.export(DataFormat::Base64);
let response = CommunicationValue::new(CommunicationType::challenge)
.with_id(cv.get_id())
.add_data(
DataTypes::public_key,
DataValue::Str(public_key_to_base64(&get_public_key())),
)
.add_data(DataTypes::challenge, DataValue::Str(encrypted_challenge));
log_cv_out!(response);
let _ = self.sender.send(&response).await;
}
}
async fn handle_challenge_response(self: &Arc<Self>, cv: CommunicationValue) {
let id = *self.id.read().await as i64;
if !cv.is_type(CommunicationType::challenge_response) {
return;
}
if let DataValue::Str(response) = cv.get_data(DataTypes::challenge) {
let expected = self.challenge.read().await.clone();
if *response == expected {
*self.challenged.write().await = true;
let response = CommunicationValue::new(CommunicationType::identification_response)
.with_id(cv.get_id())
.add_data(DataTypes::accepted, DataValue::Bool(true));
log_cv_out!(response);
if let Err(_) = self.sender.send(&response).await {
return;
}
} else {
log_err!(
id,
PrintType::Iota,
"Challenge response mismatch expected={} actual={}",
expected,
response
);
}
} else {
log_err!(
id,
PrintType::Iota,
"Challenge response missing challenge payload"
);
}
}
async fn migrate(self: &Arc<Self>) -> bool {
let kind = match *self.connection_kind.read().await {
Some(kind) => kind,
None => {
return false;
}
};
let id = *self.id.read().await;
match kind {
ConnectionKind::Client => {
let notify = CommunicationValue::new(CommunicationType::user_connected)
.add_data(DataTypes::user_id, DataValue::Number(id as i64));
get_omega_connection().send_message(&notify).await;
let user_id = id as i64;
let client = ClientConnection::from_general(self.clone(), id).await;
let mut rho = rho_manager::get_rho_con_for_user(user_id).await;
if rho.is_none() {
let get_user_msg = CommunicationValue::new(CommunicationType::get_user_data)
.add_data(DataTypes::user_id, DataValue::Number(user_id));
if let Ok(user_data_cv) = get_omega_connection()
.await_response(&get_user_msg, Some(Duration::from_secs(20)))
.await
{
if let DataValue::Number(iota_id) =
user_data_cv.get_data(DataTypes::iota_id)
{
if let Some(bound_rho) =
rho_manager::bind_user_to_iota(user_id, *iota_id).await
{
bound_rho.bind_user_id(user_id).await;
rho = Some(bound_rho);
}
}
}
}
*self.rho_connection.write().await = rho.clone();
if let Some(rho_conn) = rho {
rho_conn.bind_user_id(user_id).await;
rho_conn.add_client_connection(client.clone()).await;
} else {
log_err!(
user_id,
PrintType::Client,
"No RhoConnection found for user {}, client not attached to iota",
id
);
}
client.start();
}
ConnectionKind::Iota => {
let notify = CommunicationValue::new(CommunicationType::iota_connected)
.add_data(DataTypes::iota_id, DataValue::Number(id as i64));
get_omega_connection().send_message(&notify).await;
let iota = IotaConnection::from_general(self.clone(), id).await;
let rho = Arc::new(RhoConnection::new(iota.clone(), Vec::new()).await);
iota.set_rho_connection(rho.clone()).await;
rho_manager::add_rho(rho).await;
let get_iota_msg = CommunicationValue::new(CommunicationType::get_iota_data)
.add_data(DataTypes::iota_id, DataValue::Number(id as i64));
if let Ok(iota_data_cv) = get_omega_connection()
.await_response(&get_iota_msg, Some(Duration::from_secs(20)))
.await
{
if let DataValue::Array(users) = iota_data_cv.get_data(DataTypes::user_ids) {
let mut user_ids: Vec<u64> = Vec::new();
for value in users {
if let DataValue::Number(user_id) = value {
user_ids.push(*user_id as u64);
}
}
iota.set_user_ids(user_ids).await;
}
}
iota.clone().start();
}
ConnectionKind::AnonymousClient => {
let client = AnonymousClientConnection::from_general(self.clone(), id).await;
client.start();
}
ConnectionKind::Phi => {
let phi = ClientConnection::from_general(self.clone(), id).await;
phi.start();
}
}
true
}
}

461
src/rho/iota_connection.rs Executable file
View file

@ -0,0 +1,461 @@
use crate::calls::call_group::CallGroup;
use crate::calls::call_manager;
use crate::log_cv_in;
use crate::log_cv_out;
use crate::log_err;
use crate::log_in;
use crate::omega::omega_connection::get_omega_connection;
use crate::rho::connection::GeneralConnection;
use crate::util::logger::PrintType;
use dashmap::DashMap;
use std::collections::BTreeMap;
use std::{collections::HashMap, sync::Arc, time::Duration};
use tokio::sync::RwLock;
use tokio::sync::mpsc;
use ttp_core::CommunicationType;
use ttp_core::CommunicationValue;
use ttp_core::DataTypes;
use ttp_core::DataValue;
use ttp_native::Receiver;
use ttp_native::Sender;
use x448::PublicKey;
use super::{rho_connection::RhoConnection, rho_manager};
use crate::omega::omega_connection::OmegaConnection;
#[allow(dead_code)]
pub struct IotaConnection {
pub iota_id: u64,
pub sender: Arc<Sender>,
pub receiver: Arc<Receiver>,
pub user_ids: Arc<RwLock<Vec<u64>>>,
pub ping: Arc<RwLock<i64>>,
pub_key: Arc<RwLock<Option<Vec<u8>>>>,
pub waiting_tasks:
DashMap<u32, Box<dyn Fn(Arc<IotaConnection>, CommunicationValue) -> bool + Send + Sync>>,
pub rho_connection: Arc<RwLock<Option<Arc<RhoConnection>>>>,
}
impl IotaConnection {
pub async fn from_general(general: Arc<GeneralConnection>, iota_id: u64) -> Arc<Self> {
Arc::new(Self {
ping: Arc::new(RwLock::new(0)),
pub_key: Arc::new(RwLock::new(None)),
rho_connection: general.rho_connection.clone(),
user_ids: Arc::new(RwLock::new(Vec::new())),
sender: general.sender.clone(),
receiver: general.receiver.clone(),
iota_id: iota_id,
waiting_tasks: DashMap::new(),
})
}
pub fn start(self: Arc<Self>) {
let self_clone = self.clone();
tokio::spawn(async move {
loop {
match self_clone.receiver.receive().await {
Ok(cv) => {
self_clone.clone().handle_message(cv).await;
}
Err(_) => {
break;
}
}
}
self_clone.handle_close().await;
});
}
/// Get the Iota ID
pub async fn get_iota_id(&self) -> u64 {
self.iota_id
}
#[allow(dead_code)]
pub async fn get_public_key(&self) -> Option<PublicKey> {
if let Some(public_key) = self.pub_key.read().await.clone() {
PublicKey::from_bytes(&public_key)
} else {
None
}
}
/// Get the user IDs
pub async fn get_user_ids(&self) -> Vec<u64> {
self.user_ids.read().await.clone()
}
/// Replace all users linked to this iota and synchronize the attached rho mapping.
pub async fn set_user_ids(&self, user_ids: Vec<u64>) {
{
let mut guard = self.user_ids.write().await;
*guard = user_ids.clone();
}
if let Some(rho_conn) = self.get_rho_connection().await {
let user_ids_i64: Vec<i64> = user_ids.into_iter().map(|u| u as i64).collect();
rho_conn.set_user_ids(user_ids_i64).await;
}
}
pub async fn add_user_id(&self, user_id: u64) {
let mut should_sync = false;
{
let mut guard = self.user_ids.write().await;
if !guard.contains(&user_id) {
guard.push(user_id);
should_sync = true;
}
}
if should_sync {
if let Some(rho_conn) = self.get_rho_connection().await {
rho_conn.add_user_id(user_id as i64).await;
}
}
}
/// Get current ping
pub async fn get_ping(&self) -> i64 {
*self.ping.read().await
}
/// Set the RhoConnection reference
pub async fn set_rho_connection(&self, rho_connection: Arc<RhoConnection>) {
let mut rho_ref = self.rho_connection.write().await;
*rho_ref = Some(rho_connection);
}
/// Get RhoConnection if available
pub async fn get_rho_connection(&self) -> Option<Arc<RhoConnection>> {
let rho_ref = self.rho_connection.read().await;
if let Some(weak_ref) = rho_ref.as_ref() {
Some(weak_ref.clone())
} else {
None
}
}
/// Send a CommunicationValue to the Iota
pub async fn send_message(&self, cv: &CommunicationValue) {
if !cv.is_type(CommunicationType::pong) {
log_cv_out!(PrintType::Iota, cv);
}
if let Err(e) = self.sender.send(&cv).await {
log_err!(
self.iota_id as i64,
PrintType::Iota,
"Failed to send message: {:?}",
e
);
}
}
/// Handle incoming message from Iota
pub async fn handle_message(self: Arc<Self>, cv: CommunicationValue) {
let msg_id = cv.get_id();
if let Some((_, task)) = self.waiting_tasks.remove(&msg_id) {
if (task)(self.clone(), cv.clone()) {
return;
}
}
// Handle ping
if cv.is_type(CommunicationType::ping) || cv.is_type(CommunicationType::pong) {
self.handle_ping(cv).await;
return;
}
log_cv_in!(PrintType::Iota, cv);
// Handle GET_CHATS
if cv.is_type(CommunicationType::get_chats) {
self.handle_get_chats(cv).await;
return;
}
// Handle forwarding to other Iotas or clients
let receiver_id = cv.get_receiver();
if (receiver_id != 0 && !self.get_user_ids().await.contains(&(receiver_id as u64)))
|| cv.is_type(CommunicationType::message_other_iota)
|| cv.is_type(CommunicationType::send_chat)
{
self.handle_forward_message(cv).await;
return;
}
if cv.is_type(CommunicationType::change_iota_data)
|| cv.is_type(CommunicationType::push_notification)
|| cv.is_type(CommunicationType::get_user_data)
|| cv.is_type(CommunicationType::get_iota_data)
|| cv.is_type(CommunicationType::get_register)
|| cv.is_type(CommunicationType::complete_register_user)
|| cv.is_type(CommunicationType::delete_iota)
{
let sender = self.get_iota_id().await;
self.handle_omega_forward(cv.with_sender(sender as u64))
.await;
return;
}
self.forward_to_client(cv).await;
}
#[allow(dead_code)]
async fn send_error_response(&self, message_id: u32, error_type: CommunicationType) {
let error = CommunicationValue::new(error_type).with_id(message_id);
self.send_message(&error).await;
}
#[allow(dead_code)]
async fn close(&self) {
let _ = self.sender.close();
}
async fn handle_omega_forward(self: Arc<Self>, cv: CommunicationValue) {
let iota_for_closure = self.clone();
tokio::spawn(async move {
let response_cv = get_omega_connection()
.await_response(&cv.with_sender(self.iota_id), Some(Duration::from_secs(20)))
.await;
if let Ok(response_cv) = response_cv {
iota_for_closure.send_message(&response_cv).await;
}
});
}
/// Handle ping message
async fn handle_ping(&self, cv: CommunicationValue) {
if let DataValue::Number(last_ping) = cv.get_data(DataTypes::last_ping) {
if let Ok(ping_val) = last_ping.to_string().parse::<i64>() {
let mut ping_guard = self.ping.write().await;
*ping_guard = ping_val;
}
}
let client_pings = if let Some(rho_conn) = self.get_rho_connection().await {
rho_conn.get_client_pings().await
} else {
HashMap::new()
};
let pings: Vec<(DataTypes, DataValue)> = client_pings
.into_iter()
.map(|(k, v)| (DataTypes::parse(k), DataValue::Number(v)))
.collect();
let response = CommunicationValue::new(CommunicationType::pong)
.with_id(cv.get_id())
.add_data(DataTypes::ping_clients, DataValue::Container(pings));
self.send_message(&response).await;
}
/// Handle message forwarding to other Iotas
async fn handle_forward_message(&self, cv: CommunicationValue) {
let receiver_id = cv.get_receiver();
let sender_id = cv.get_sender();
let my_user_ids = self.get_user_ids().await;
log_in!(
self.iota_id as i64,
PrintType::Iota,
"Authority check: sender_id={} receiver_id={} iota_user_ids={:?} msg_type={:?} msg_id={}",
sender_id,
receiver_id,
my_user_ids,
cv.get_type(),
cv.get_id()
);
if my_user_ids.contains(&(sender_id as u64)) {
if let Some(target_rho) = rho_manager::get_rho_con_for_user(receiver_id as i64).await {
target_rho.message_to_iota(cv).await;
} else {
let error = CommunicationValue::new(CommunicationType::error_no_iota)
.with_id(cv.get_id())
.with_sender(cv.get_sender());
self.send_message(&error).await;
}
} else {
log_err!(
self.iota_id as i64,
PrintType::Iota,
"Rejected client->iota forward: sender_id={} is not authorized for this iota. Known users={:?}",
sender_id,
my_user_ids
);
self.send_message(
&CommunicationValue::new(CommunicationType::error_invalid_user_id).add_data(
DataTypes::error_type,
DataValue::Str(
"You are sending to another User without authority.".to_string(),
),
),
)
.await;
}
}
/// Handle GET_CHATS message
async fn handle_get_chats(&self, cv: CommunicationValue) {
let receiver_id = cv.get_receiver();
let mut interested_ids: Vec<i64> = Vec::new();
// ============================
// Load Calls
// ============================
let calls: Vec<Arc<CallGroup>> = call_manager::get_call_groups(receiver_id).await;
let mut invites: HashMap<i64, Vec<DataValue>> = HashMap::new();
let empty = calls.is_empty();
for call in calls {
for inviter in call.members.read().await.iter() {
let call_self = call.get_caller(receiver_id).await.unwrap();
let inviter_id = inviter.user_id;
let timeout = *call_self.timeout.read().await;
let admin = call_self.has_admin();
// Build call container
let mut call_map: BTreeMap<DataTypes, DataValue> = BTreeMap::new();
call_map.insert(DataTypes::call_id, DataValue::Str(call.call_id.to_string()));
if timeout > 0 {
call_map.insert(DataTypes::timeout, DataValue::Number(timeout as i64));
}
if admin {
call_map.insert(DataTypes::has_admin, DataValue::Bool(true));
}
let call_container = DataValue::container_from_map(&call_map);
invites
.entry(inviter_id as i64)
.or_insert_with(Vec::new)
.push(call_container);
}
}
// ============================
// Enrich Contacts
// ============================
let enriched_contacts = if empty {
match cv.get_data(DataTypes::user_ids) {
DataValue::Array(arr) => DataValue::Array(arr.clone()),
_ => DataValue::Array(vec![]),
}
} else {
let mut enriched: Vec<DataValue> = Vec::new();
if let DataValue::Array(users) = cv.get_data(DataTypes::user_ids) {
for user_val in users {
if let DataValue::Container(entries) = user_val {
let mut user_map: BTreeMap<DataTypes, DataValue> =
entries.iter().cloned().collect();
// extract user_id
if let Some(DataValue::Number(user_id)) = user_map.get(&DataTypes::user_id)
{
interested_ids.push(*user_id);
// attach calls if exists
if let Some(call_list) = invites.get(user_id) {
user_map
.insert(DataTypes::calls, DataValue::Array(call_list.clone()));
}
}
enriched.push(DataValue::container_from_map(&user_map));
}
}
}
DataValue::Array(enriched)
};
// ============================
// Notify Omega
// ============================
OmegaConnection::user_states(receiver_id as i64, interested_ids.clone()).await;
// ============================
// Notify Rho
// ============================
if let Some(rho_conn) = self.get_rho_connection().await {
rho_conn
.set_interested(receiver_id as i64, interested_ids)
.await;
}
// ============================
// Forward to client
// ============================
self.forward_to_client(cv.add_data(DataTypes::user_ids, enriched_contacts))
.await;
}
/// Forward message to client
async fn forward_to_client(&self, cv: CommunicationValue) {
if let Some(rho_conn) = self.get_rho_connection().await {
let updated_cv = cv.with_sender(self.get_iota_id().await);
rho_conn.message_to_client(updated_cv).await;
} else {
}
}
pub async fn handle_close(&self) {
if let Some(rho_conn) = self.get_rho_connection().await {
rho_conn.close_iota_connection().await;
}
}
#[allow(dead_code)]
pub async fn await_response(
self: Arc<IotaConnection>,
cv: &CommunicationValue,
timeout_duration: Option<Duration>,
) -> Result<CommunicationValue, String> {
let (tx, mut rx) = mpsc::channel(1);
let msg_id = cv.get_id();
let task_tx = tx.clone();
self.waiting_tasks.insert(
msg_id,
Box::new(move |_, response_cv| {
let inner_tx = task_tx.clone();
tokio::spawn(async move {
let _ = inner_tx.send(response_cv).await;
});
true
}),
);
self.send_message(cv).await;
let timeout = timeout_duration.unwrap_or(Duration::from_secs(10));
match tokio::time::timeout(timeout, rx.recv()).await {
Ok(Some(response_cv)) => Ok(response_cv),
Ok(_) => Err("Failed to receive response, channel was closed.".to_string()),
Err(_) => {
self.waiting_tasks.remove(&msg_id);
Err(format!(
"Request timed out after {} seconds.",
timeout.as_secs()
))
}
}
}
}
impl std::fmt::Debug for IotaConnection {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("IotaConnection")
.field("iota_id", &"[async]")
.field("identified", &"[async]")
.field("ping", &"[async]")
.finish()
}
}

6
src/rho/mod.rs Normal file
View file

@ -0,0 +1,6 @@
pub mod client_connection;
pub mod connection;
pub mod iota_connection;
pub mod rho_connection;
pub mod rho_manager;
pub mod server;

199
src/rho/rho_connection.rs Normal file
View file

@ -0,0 +1,199 @@
use super::{client_connection::ClientConnection, iota_connection::IotaConnection, rho_manager};
use crate::data::user::UserStatus;
use crate::omega::omega_connection::OmegaConnection;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use ttp_core::{CommunicationType, CommunicationValue, DataTypes, DataValue};
pub struct RhoConnection {
iota_connection: Arc<IotaConnection>,
user_ids: Arc<RwLock<Vec<i64>>>,
client_connections: Arc<RwLock<Vec<Arc<ClientConnection>>>>,
}
impl RhoConnection {
/// Create a new RhoConnection
pub async fn new(iota_connection: Arc<IotaConnection>, user_ids: Vec<i64>) -> Self {
let rho_connection = Self {
iota_connection,
user_ids: Arc::new(RwLock::new(user_ids.clone())),
client_connections: Arc::new(RwLock::new(Vec::new())),
};
rho_connection
}
pub async fn get_iota_id(&self) -> u64 {
self.iota_connection.iota_id
}
pub async fn get_user_ids(&self) -> Vec<i64> {
self.user_ids.read().await.clone()
}
pub async fn set_user_ids(&self, user_ids: Vec<i64>) {
let mut guard = self.user_ids.write().await;
*guard = user_ids;
}
pub async fn add_user_id(&self, user_id: i64) {
let mut guard = self.user_ids.write().await;
if !guard.contains(&user_id) {
guard.push(user_id);
}
}
pub async fn bind_user_id(&self, user_id: i64) {
self.add_user_id(user_id).await;
self.iota_connection.add_user_id(user_id as u64).await;
}
pub fn get_iota_connection(&self) -> &Arc<IotaConnection> {
&self.iota_connection
}
pub async fn get_client_connections(&self) -> Vec<Arc<ClientConnection>> {
let connections = self.client_connections.read().await;
connections.clone()
}
/// Get client connections for a specific user
pub async fn get_client_connections_for_user(
&self,
user_id: i64,
) -> Vec<Arc<ClientConnection>> {
let connections = self.client_connections.read().await;
let mut collections = Vec::new();
for con in connections.iter() {
if con.get_user_id().await == user_id as u64 {
collections.push(con.clone());
}
}
collections
}
/// Add a client connection
#[allow(dead_code)]
pub async fn add_client_connection(&self, connection: Arc<ClientConnection>) {
let notification = CommunicationValue::new(CommunicationType::client_connected).add_data(
DataTypes::user_id,
DataValue::Number(connection.get_user_id().await as i64),
);
self.iota_connection.send_message(&notification).await;
{
let mut connections = self.client_connections.write().await;
connections.push(Arc::clone(&connection));
}
OmegaConnection::client_changed(
self.get_iota_id().await as i64,
connection.get_user_id().await as i64,
UserStatus::user_online,
)
.await;
}
/// Remove a client connection
pub async fn close_client_connection(&self, connection: Arc<ClientConnection>) {
let target_user_id = connection.get_user_id().await;
{
let mut connections = self.client_connections.write().await;
connections.retain(|con| {
futures::executor::block_on(async { con.get_user_id().await != target_user_id })
});
}
// Notify OmegaConnection
OmegaConnection::client_changed(
self.get_iota_id().await as i64,
connection.get_user_id().await as i64,
UserStatus::user_offline,
)
.await;
}
/// Close the Iota connection and all associated client connections
pub async fn close_iota_connection(&self) {
// Close all client connections
let connections = self.get_client_connections().await;
for connection in connections {
connection.close().await;
}
// Remove from manager
rho_manager::remove_rho(self.get_iota_id().await as i64).await;
// Notify OmegaConnection
OmegaConnection::close_iota(self.get_iota_id().await as i64).await;
}
/// Send message from Iota to specific client
pub async fn message_to_client(&self, cv: CommunicationValue) {
let connections = self.client_connections.read().await;
let receiver_id = cv.get_receiver();
for connection in connections.iter() {
if connection.get_user_id().await == receiver_id {
connection.clone().send_message(&cv).await;
}
}
}
/// Send message to Iota
pub async fn message_to_iota(&self, cv: CommunicationValue) {
self.iota_connection.send_message(&cv).await;
}
/// Set interested users for a specific client
pub async fn set_interested(&self, user_id: i64, interested_ids: Vec<i64>) {
let connections = self.client_connections.read().await;
for connection in connections.iter() {
let conn_user_id = connection.get_user_id().await;
if conn_user_id == user_id as u64 {
connection
.clone()
.set_interested_users(interested_ids.clone())
.await;
break;
}
}
}
/// Check if clients are interested in a user
#[allow(dead_code)]
pub async fn are_they_interested(&self, user_id: i64) {
let connections = self.client_connections.read().await;
for connection in connections.iter() {
connection.clone().are_you_interested(user_id).await;
}
}
/// Get ping information for all clients
pub async fn get_client_pings(&self) -> HashMap<String, i64> {
let connections = self.client_connections.read().await;
let mut pings = HashMap::new();
for connection in connections.iter() {
let user_id = connection.get_user_id().await;
pings.insert(user_id.to_string(), connection.get_ping().await);
}
pings
}
/// Check if this RhoConnection contains a specific user ID
#[allow(dead_code)]
pub async fn contains_user(&self, user_id: &i64) -> bool {
self.user_ids.read().await.contains(user_id)
}
/// Get count of active client connections
#[allow(dead_code)]
pub async fn client_count(&self) -> usize {
let connections = self.client_connections.read().await;
connections.len()
}
}

83
src/rho/rho_manager.rs Normal file
View file

@ -0,0 +1,83 @@
use super::rho_connection::RhoConnection;
use crate::log_in;
use crate::util::logger::PrintType;
use std::{
collections::HashMap,
sync::{Arc, LazyLock},
};
use tokio::sync::RwLock;
pub static RHO_CONNECTIONS: LazyLock<Arc<RwLock<HashMap<i64, Arc<RhoConnection>>>>> =
LazyLock::new(|| Arc::new(RwLock::new(HashMap::new())));
pub async fn get_rho_con_for_user(user_id: i64) -> Option<Arc<RhoConnection>> {
let connections = RHO_CONNECTIONS.read().await;
for rho_connection in connections.values() {
let rho_user_ids = rho_connection.get_user_ids().await;
log_in!(
user_id,
PrintType::Client,
"Comparing user IDs: {:?}",
rho_user_ids
);
if rho_user_ids.contains(&user_id) {
return Some(Arc::clone(rho_connection));
}
}
None
}
#[allow(dead_code)]
pub async fn contains_iota(iota_id: i64) -> bool {
let connections = RHO_CONNECTIONS.read().await;
connections.contains_key(&iota_id)
}
/// Bind a user ID to an already tracked iota/rho connection.
pub async fn bind_user_to_iota(user_id: i64, iota_id: i64) -> Option<Arc<RhoConnection>> {
let connections = RHO_CONNECTIONS.read().await;
if let Some(rho_connection) = connections.get(&iota_id) {
let rho = Arc::clone(rho_connection);
drop(connections);
rho.add_user_id(user_id).await;
log_in!(
user_id,
PrintType::Client,
"Bound user {} to iota {}",
user_id,
iota_id
);
Some(rho)
} else {
None
}
}
/// Remove a RhoConnection by Iota ID
pub async fn remove_rho(iota_id: i64) -> Option<Arc<RhoConnection>> {
let mut connections = RHO_CONNECTIONS.write().await;
connections.remove(&iota_id)
}
/// Add a RhoConnection to the manager
pub async fn add_rho(rho_connection: Arc<RhoConnection>) {
let mut connections = RHO_CONNECTIONS.write().await;
let iota_id = rho_connection.get_iota_id().await;
connections.insert(iota_id as i64, rho_connection);
}
/// Get a RhoConnection by Iota ID directly
#[allow(dead_code)]
pub async fn get_rho_by_iota(iota_id: i64) -> Option<Arc<RhoConnection>> {
let connections = RHO_CONNECTIONS.read().await;
connections.get(&iota_id).map(Arc::clone)
}
/// Get the count of active connections
pub async fn connection_count() -> usize {
let connections = RHO_CONNECTIONS.read().await;
connections.len()
}

25
src/rho/server.rs Normal file
View file

@ -0,0 +1,25 @@
use crate::{
log,
rho::connection::GeneralConnection,
util::{file_util::load_file_vec, logger::PrintType},
};
use ttp_native::Host;
pub async fn start(port: u16) -> Result<(), Box<dyn std::error::Error>> {
let cert_pem = load_file_vec("certs", "cert.pem").expect("Error loading Pemfile");
let key_pem = load_file_vec("certs", "key.pem").expect("Error loading Keyfile");
let mut host: Host = ttp_native::host(port, cert_pem, key_pem).await?;
log!(0, PrintType::General, "Server listening on port {}", port);
while let Some((sender, receiver)) = host.next().await {
tokio::spawn(async move {
let conn = GeneralConnection::new(sender, receiver);
conn.handle().await;
});
}
log!(0, PrintType::General, "Server stopped");
Ok(())
}

View file

@ -1,275 +0,0 @@
use crate::get_public_key;
use crate::sql::sql;
use crate::sql::user_online_tracker::get_iota_primary_omikron_connection;
use crate::transport::omikron_manager::get_random_omikron;
use crate::util::file_util::get_directory;
use crate::{
sql::sql::{get_by_user_id, get_omikron_by_id},
util::crypto_helper::public_key_to_base64,
};
use actix_web::HttpResponse;
use actix_web::http::{StatusCode, header};
use base64::Engine as _;
use json::JsonValue;
pub async fn handle(path: &str, body_string: Option<String>) -> HttpResponse {
if path == "OPTIONS" {
return HttpResponse::Ok()
.insert_header(("Access-Control-Allow-Origin", "*"))
.insert_header(("Access-Control-Allow-Methods", "GET, POST, OPTIONS"))
.insert_header(("Access-Control-Allow-Headers", "*"))
.finish();
}
let path_parts: Vec<&str> = path.split("/").filter(|s| !s.is_empty()).collect();
let _body: Option<JsonValue> = if body_string.is_some() {
if let Ok(body_json) = json::parse(&body_string.unwrap()) {
Some(body_json)
} else {
None
}
} else {
None
};
let (status, body_text) = match path_parts.as_slice() {
// ==================================================
// DOWNLOAD IOTA FRONTEND
// ==================================================
["api", "download", "iota_frontend"] => {
let file_path = format!("{}/downloads/iota_frontend.zip", get_directory());
match std::fs::read(file_path) {
Ok(file_bytes) => {
return HttpResponse::Ok()
.insert_header(("Access-Control-Allow-Origin", "*"))
.insert_header(("Content-Type", "application/zip"))
.insert_header((
"Content-Disposition",
"attachment; filename=\"iota_frontend.zip\"",
))
.body(file_bytes);
}
Err(_) => {
let mut res = JsonValue::new_object();
res["status"] = "error_not_found".into();
return HttpResponse::NotFound()
.insert_header(("Access-Control-Allow-Origin", "*"))
.body(res.dump());
}
}
}
// ==================================================
// GET RANDOM OMIKRON
// ==================================================
["api", "get", "omikron"] => {
if let Ok(omikron_conn) = get_random_omikron().await {
if let Some(id) = omikron_conn.get_omikron_id().await {
if let Ok((public_key, ip_address)) = sql::get_omikron_by_id(id).await {
let mut res = JsonValue::new_object();
res["status"] = "success".into();
res["id"] = id.into();
res["public_key"] = public_key.into();
res["ip_address"] = ip_address.into();
(StatusCode::OK, res.dump())
} else {
let mut res = JsonValue::new_object();
res["status"] = "error".into();
(StatusCode::INTERNAL_SERVER_ERROR, res.dump())
}
} else {
let mut res = JsonValue::new_object();
res["status"] = "error".into();
(StatusCode::INTERNAL_SERVER_ERROR, res.dump())
}
} else {
let mut res = JsonValue::new_object();
res["status"] = "error_not_found".into();
(StatusCode::NOT_FOUND, res.dump())
}
}
// ==================================================
// GET OMIKRON BY ID
// ==================================================
["api", "get", "omikron", id] => {
let id = id.parse::<i64>().unwrap_or(0);
if id == 0 {
let mut res = JsonValue::new_object();
res["status"] = "error_bad_request".into();
(StatusCode::BAD_REQUEST, res.dump())
} else if let Ok((public_key, ip_address)) = get_omikron_by_id(id).await {
let mut res = JsonValue::new_object();
res["status"] = "success".into();
res["id"] = id.into();
res["public_key"] = public_key.into();
res["ip_address"] = ip_address.into();
(StatusCode::OK, res.dump())
} 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 {
let mut res = JsonValue::new_object();
res["status"] = "success".into();
res["id"] = omikron_id.into();
res["public_key"] = public_key.into();
res["ip_address"] = ip_address.into();
(StatusCode::OK, res.dump())
} else {
let mut res = JsonValue::new_object();
res["status"] = "error_not_found".into();
(StatusCode::NOT_FOUND, res.dump())
}
} else if let Ok((_, iota_id, _, _, _, _, _, _, _, _, _, _)) = get_by_user_id(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 {
let mut res = JsonValue::new_object();
res["status"] = "success".into();
res["id"] = omikron_id.into();
res["public_key"] = public_key.into();
res["ip_address"] = ip_address.into();
(StatusCode::OK, res.dump())
} else {
let mut res = JsonValue::new_object();
res["status"] = "error_not_found".into();
(StatusCode::NOT_FOUND, res.dump())
}
} else {
let mut res = JsonValue::new_object();
res["status"] = "error_not_found".into();
(StatusCode::NOT_FOUND, res.dump())
}
} else {
let mut res = JsonValue::new_object();
res["status"] = "error_not_found".into();
(StatusCode::NOT_FOUND, res.dump())
}
}
// ==================================================
// GET ID BY USERNAME
// ==================================================
["api", "get", "id", username] => {
if username.is_empty() {
let mut res = JsonValue::new_object();
res["status"] = "error_bad_request".into();
(StatusCode::BAD_REQUEST, res.dump())
} else if let Ok((
id,
iota_id,
username,
_,
_,
_,
_,
sub_level,
sub_end,
public_key,
_,
_,
)) = sql::get_by_username(username).await
{
let mut res = JsonValue::new_object();
res["status"] = "success".into();
res["username"] = username.into();
res["public_key"] = public_key.into();
res["user_id"] = id.into();
res["iota_id"] = iota_id.into();
res["sub_level"] = sub_level.into();
res["sub_end"] = sub_end.into();
(StatusCode::OK, res.dump())
} else {
let mut res = JsonValue::new_object();
res["status"] = "error_not_found".into();
(StatusCode::OK, res.dump())
}
}
// ==================================================
// GET SERVER PUBLIC KEY
// ==================================================
["api", "get", "public_key"] => {
let mut res = JsonValue::new_object();
res["status"] = "success".into();
res["public_key"] = public_key_to_base64(&get_public_key()).into();
(StatusCode::OK, res.dump())
}
// ==================================================
// GET USER BY ID
// ==================================================
["api", "get", "user", id] => {
let id: i64 = id.parse().unwrap_or(0);
if id == 0 {
let mut res = JsonValue::new_object();
res["status"] = "error_bad_request".into();
(StatusCode::BAD_REQUEST, res.dump())
} else if let Ok((
id,
iota_id,
username,
display,
status_msg,
about,
avatar,
sub_level,
sub_end,
public_key,
_,
_,
)) = sql::get_by_user_id(id).await
{
let mut res = JsonValue::new_object();
res["status"] = "success".into();
res["username"] = username.into();
res["public_key"] = public_key.into();
res["user_id"] = id.into();
res["iota_id"] = iota_id.into();
res["sub_level"] = sub_level.into();
res["sub_end"] = sub_end.into();
if let Some(display) = display {
res["display"] = display.into();
}
if let Some(status_msg) = status_msg {
res["status_message"] = status_msg.into();
}
if let Some(about) = about {
res["about"] = about.into();
}
if let Some(avatar) = avatar {
res["avatar"] = base64::engine::general_purpose::STANDARD
.encode(avatar)
.into();
}
(StatusCode::OK, res.dump())
} else {
let mut res = JsonValue::new_object();
res["status"] = "error_not_found".into();
(StatusCode::OK, res.dump())
}
}
// ==================================================
// DEFAULT
// ==================================================
_ => {
let mut res = JsonValue::new_object();
res["status"] = "error".into();
(StatusCode::INTERNAL_SERVER_ERROR, res.dump())
}
};
let body_bytes = body_text.into_bytes();
HttpResponse::build(status)
.insert_header((header::ACCESS_CONTROL_ALLOW_ORIGIN, "*"))
.insert_header((header::ACCESS_CONTROL_ALLOW_HEADERS, "*"))
.insert_header((header::ACCESS_CONTROL_ALLOW_METHODS, "GET, POST, OPTIONS"))
.body(body_bytes)
}

View file

@ -1,3 +0,0 @@
pub mod api;
pub mod server;
pub mod short_link;

View file

@ -1,67 +0,0 @@
use crate::{
log,
server::{api, short_link::get_short_link},
util::file_util::load_file_buf,
};
use actix_web::{App, HttpRequest, HttpResponse, HttpServer, Responder, http::header, web};
use rustls::ServerConfig;
use rustls::pki_types::{CertificateDer, PrivateKeyDer};
use rustls_pemfile::{certs, pkcs8_private_keys};
pub async fn start(port: u16) -> anyhow::Result<()> {
let mut cert_reader = load_file_buf("certs", "server_cert.pem")?;
let mut key_reader = load_file_buf("certs", "server_key.pem")?;
let cert_chain: Vec<CertificateDer<'static>> =
certs(&mut cert_reader).collect::<Result<_, _>>()?;
let mut keys: Vec<PrivateKeyDer<'static>> = pkcs8_private_keys(&mut key_reader)
.map(|res| res.map(Into::into))
.collect::<Result<_, _>>()?;
let key = keys.remove(0);
let mut config = ServerConfig::builder()
.with_no_client_auth()
.with_single_cert(cert_chain, key)?;
config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
let addr = format!("0.0.0.0:{port}");
log!(" Server on {}", addr);
HttpServer::new(move || {
App::new()
.route("/api/{path:.*}", web::to(api_handler))
.route("/direct/{path:.*}", web::to(direct_handler))
})
.bind_rustls_0_23(addr, config)?
.run()
.await?;
Ok(())
}
async fn direct_handler(req: HttpRequest) -> impl Responder {
let path = req.uri().path().to_string();
let short = path.replace("/direct/", "");
if let Ok(long) = get_short_link(&short).await {
HttpResponse::TemporaryRedirect()
.append_header((header::LOCATION, long))
.finish()
} else {
HttpResponse::TemporaryRedirect()
.append_header((header::LOCATION, "https://tensamin.net"))
.finish()
}
}
async fn api_handler(req: HttpRequest, body: web::Bytes) -> HttpResponse {
let path = req.uri().path().to_string();
let body_string = String::from_utf8_lossy(&body).to_string();
api::handle(&path, Some(body_string)).await
}

View file

@ -1,87 +0,0 @@
use dashmap::DashMap;
use once_cell::sync::Lazy;
use rand::{Rng, thread_rng};
static LINKS: Lazy<DashMap<String, String>> = Lazy::new(DashMap::new);
const CHARSET: &[u8] = b"abcdefghijklmnopqrstuvwxyzABCDEFGHJKLMNPRSTUVWXYZ1234567890";
pub async fn add_short_link(long: &str) -> Result<String, ()> {
let raw = generate_unique_short_link().await;
LINKS.insert(raw.clone(), long.to_string());
Ok(format!(
"https://omega.tensamin.net/direct/{}",
format_with_dashes(&raw)
))
}
async fn generate_unique_short_link() -> String {
loop {
let short = generate_short_link().await;
if !LINKS.contains_key(&short) {
return short;
}
}
}
pub async fn generate_short_link() -> String {
let len = short_length();
let mut rng = thread_rng();
(0..len)
.map(|_| {
let idx = rng.gen_range(0..CHARSET.len());
CHARSET[idx] as char
})
.collect()
}
pub async fn get_short_link(short: &str) -> Result<String, ()> {
let key = if short.contains("/") {
short.split("/").nth(1).unwrap_or_default()
} else {
short
};
let frag = short.replace(key, "");
let normalized = normalize_short(&key);
if let Ok(t) = LINKS.get(&normalized).map(|v| v.value().clone()).ok_or(()) {
Ok(format!("{}{}", t, frag))
} else {
Err(())
}
}
/* ---------------- helpers ---------------- */
fn short_length() -> usize {
let count = LINKS.len();
match count {
0..=1_999 => 4,
2_000..=999_999 => 8,
_ => 12,
}
}
fn format_with_dashes(s: &str) -> String {
s.chars()
.collect::<Vec<_>>()
.chunks(4)
.map(|c| c.iter().collect::<String>())
.collect::<Vec<_>>()
.join("-")
}
fn normalize_short(input: &str) -> String {
input
.chars()
.filter(|c| *c != '-')
.map(|c| match c {
'Q' | 'O' => '0',
'I' => 'l',
_ => c,
})
.collect()
}

View file

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

View file

@ -1,765 +0,0 @@
use crate::log;
use once_cell::sync::Lazy;
use sqlx::{MySql, Pool, Row, mysql::MySqlPoolOptions};
use std::sync::atomic::{AtomicU64, Ordering};
use std::{
env,
sync::Arc,
time::{SystemTime, UNIX_EPOCH},
};
use tokio::sync::RwLock;
/*
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_register_id, register_complete_iota, register_complete_user,
},
};
*/
static SQL_DB: Lazy<Arc<RwLock<Option<Pool<MySql>>>>> = Lazy::new(|| Arc::new(RwLock::new(None)));
pub async fn connect() -> Result<Pool<MySql>, sqlx::Error> {
let user = env::var("DB_USERNAME").expect("DB_USERNAME is not set");
let passwd = env::var("DB_PASSWD").expect("DB_PASSWD is not set");
let table = env::var("DB_TABLE").expect("DB_TABLE is not set");
MySqlPoolOptions::new()
.max_connections(200)
.connect(&format!(
"mysql://{}:{}@127.0.0.1:3306/{}",
user, passwd, table
))
.await
}
// Omega
// - Omikron
// - Iota
// - User
// - User
// - Iota
// - User
// - User
// - Omikron
// - Iota
// - User
// - User
// - Iota
// - User
// - User
pub async fn initialize_db() -> Result<(), sqlx::Error> {
let pool = connect().await?;
let mut db_lock = SQL_DB.write().await;
// create tables
// with indexes
let _ = sqlx::query(
"CREATE TABLE IF NOT EXISTS
users (
id BIGINT UNSIGNED NOT NULL PRIMARY KEY,
username VARCHAR(15) NOT NULL UNIQUE COLLATE utf8mb4_bin,
display VARCHAR(15) COLLATE utf8mb4_bin,
status VARCHAR(15) COLLATE utf8mb4_bin,
about VARCHAR(200) COLLATE utf8mb4_bin,
avatar MEDIUMBLOB,
sub_level INT(11) NOT NULL DEFAULT 0,
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 DEFAULT '',
iota_id BIGINT UNSIGNED NOT NULL,
token VARCHAR(255) NOT NULL UNIQUE COLLATE utf8mb4_bin
)",
)
.execute(&pool)
.await;
let _ = sqlx::query(
"CREATE TABLE IF NOT EXISTS
iotas (
id BIGINT UNSIGNED NOT NULL PRIMARY KEY,
public_key VARCHAR(255) NOT NULL COLLATE utf8mb4_bin
)",
)
.execute(&pool)
.await;
let _ = sqlx::query(
"CREATE TABLE IF NOT EXISTS
omikrons (
id BIGINT UNSIGNED NOT NULL PRIMARY KEY,
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
)",
)
.execute(&pool)
.await;
let _ = sqlx::query(
"CREATE TABLE IF NOT EXISTS
notifications (
id BIGINT UNSIGNED NOT NULL PRIMARY KEY AUTO INCREMENT,
sender_id BIGINT UNSIGNED NOT NULL,
receiver_id BIGINT UNSIGNED NOT NULL,
amount BIGINT UNSIGNED NOT NULL DEFAULT 0
)",
)
.execute(&pool)
.await;
*db_lock = Some(pool);
Ok(())
}
// ==========================================================================================
// REGISTER
// ==========================================================================================
pub static CURRENT_MILLI_USED: Lazy<Arc<AtomicU64>> = Lazy::new(|| Arc::new(AtomicU64::new(0)));
pub static CURRENT_REGISTER_PROCESS: Lazy<Arc<RwLock<Vec<u64>>>> =
Lazy::new(|| Arc::new(RwLock::new(Vec::new())));
pub async fn get_register_id() -> u64 {
let mut current_time = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64;
loop {
let current_locked = CURRENT_MILLI_USED.load(Ordering::SeqCst);
if current_locked < current_time {
let result = CURRENT_MILLI_USED.compare_exchange(
current_locked, // expected value
current_time, // new value
Ordering::SeqCst, // acquire/release ordering
Ordering::SeqCst, // failure ordering
);
match result {
Ok(_) => {
CURRENT_REGISTER_PROCESS.write().await.push(current_time);
return current_time;
}
Err(_) => {
continue;
}
}
} else {
current_time = current_locked + 1;
}
}
}
// ==========================================================================================
// USERS
// ==========================================================================================
pub async fn get_by_username(
username: &str,
) -> Result<
(
i64,
i64,
String,
Option<String>,
Option<String>,
Option<String>,
Option<Vec<u8>>,
i32,
i64,
String,
String,
String,
),
sqlx::Error,
> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
let row = sqlx::query(
"SELECT id, iota_id, username, display, status, about, avatar, sub_level, sub_end, public_key, private_key_hash, token FROM users WHERE username = ?",
)
.bind(username)
.fetch_optional(&pool)
.await?;
match row {
Some(row) => {
let id: i64 = row.get("id");
let iota_id: i64 = row.get("iota_id");
let username: String = row.get("username");
let display: Option<Vec<u8>> = row.get("display");
let status: Option<Vec<u8>> = row.get("status");
let about: Option<Vec<u8>> = row.get("about");
let avatar: Option<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");
Ok((
id,
iota_id,
username,
display.map(|d| String::from_utf8_lossy(&d).to_string()),
status.map(|s| String::from_utf8_lossy(&s).to_string()),
about.map(|a| String::from_utf8_lossy(&a).to_string()),
avatar,
sub_level,
sub_end,
public_key,
private_key_hash,
String::from_utf8_lossy(&token).to_string(),
))
}
_ => Err(sqlx::Error::RowNotFound),
}
}
pub async fn get_by_user_id(
id: i64,
) -> Result<
(
i64,
i64,
String,
Option<String>,
Option<String>,
Option<String>,
Option<Vec<u8>>,
i32,
i64,
String,
String,
String,
),
sqlx::Error,
> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
let row = sqlx::query(
"SELECT id, iota_id, username, display, status, about, avatar, sub_level, sub_end, public_key, private_key_hash, token FROM users WHERE id = CAST(? AS UNSIGNED)",
)
.bind(id)
.fetch_optional(&pool)
.await?;
match row {
Some(row) => {
let id: i64 = row.get("id");
let iota_id: i64 = row.get("iota_id");
let username: String = row.get("username");
let display: Option<Vec<u8>> = row.get("display");
let status: Option<Vec<u8>> = row.get("status");
let about: Option<Vec<u8>> = row.get("about");
let avatar: Option<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");
Ok((
id,
iota_id,
username,
display.map(|d| String::from_utf8_lossy(&d).to_string()),
status.map(|s| String::from_utf8_lossy(&s).to_string()),
about.map(|a| String::from_utf8_lossy(&a).to_string()),
avatar,
sub_level,
sub_end,
public_key,
private_key_hash,
String::from_utf8_lossy(&token).to_string(),
))
}
_ => Err(sqlx::Error::RowNotFound),
}
}
pub async fn get_users_by_iota_id(
iota_id_param: i64,
) -> Result<
Vec<(
i64,
i64,
String,
Option<String>,
Option<String>,
Option<String>,
Option<Vec<u8>>,
i32,
i64,
String,
String,
String,
)>,
sqlx::Error,
> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool 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 = CAST(? AS UNSIGNED)",
)
.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: Option<Vec<u8>> = row.get("display");
let status: Option<Vec<u8>> = row.get("status");
let about: Option<Vec<u8>> = row.get("about");
let avatar: Option<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,
display.map(|d| String::from_utf8_lossy(&d).to_string()),
status.map(|s| String::from_utf8_lossy(&s).to_string()),
about.map(|a| String::from_utf8_lossy(&a).to_string()),
avatar,
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 pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("UPDATE users SET username = ? WHERE id = CAST(? AS UNSIGNED)")
.bind(new_username)
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
pub async fn change_display_name(id: i64, new_display: String) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("UPDATE users SET display = ? WHERE id = CAST(? AS UNSIGNED)")
.bind(new_display)
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
pub async fn change_avatar(id: i64, new_avatar: String) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("UPDATE users SET avatar = ? WHERE id = CAST(? AS UNSIGNED)")
.bind(new_avatar)
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
pub async fn change_about(id: i64, new_about: String) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("UPDATE users SET about = ? WHERE id = CAST(? AS UNSIGNED)")
.bind(new_about)
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
pub async fn change_status(id: i64, new_status: String) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("UPDATE users SET status = ? WHERE id = CAST(? AS UNSIGNED)")
.bind(new_status)
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
pub async fn delete_user(id: i64) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("DELETE FROM users WHERE id = CAST(? AS UNSIGNED)")
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
pub async fn change_iota_id(id: i64, new_iota_id: i64) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("UPDATE users SET iota_id = CAST(? AS UNSIGNED) WHERE id = CAST(? AS UNSIGNED)")
.bind(new_iota_id)
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
pub async fn change_keys(
id: i64,
new_public_key: String,
new_private_key_hash: String,
) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query(
"UPDATE users SET public_key = ?, private_key_hash = ? WHERE id = CAST(? AS UNSIGNED)",
)
.bind(new_public_key)
.bind(new_private_key_hash)
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
pub async fn change_token(id: i64, new_token: String) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("UPDATE users SET token = ? WHERE id = CAST(? AS UNSIGNED)")
.bind(new_token)
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
pub async fn register_complete_user(
id: i64,
username: String,
public_key: String,
iota_id: i64,
token: String,
) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query(
"INSERT INTO users (id, username, public_key, iota_id, token) VALUES (?, ?, ?, ?, ?)",
)
.bind(id)
.bind(username)
.bind(public_key)
.bind(iota_id)
.bind(token)
.execute(&pool)
.await?;
Ok(())
}
pub async fn print_users() -> Result<(), Box<dyn std::error::Error>> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
log!("Printing users...");
for row in sqlx::query(
"SELECT id, iota_id, username, display, status, about, sub_level, sub_end, public_key, private_key_hash, token FROM users",
)
.fetch_all(&pool)
.await?
.iter()
{
let id: i64 = row.get("id");
let iota_id: i64 = row.get("iota_id");
let username: String = row.get("username");
let display: Option<Vec<u8>> = row.get("display");
let status: Option<Vec<u8>> = row.get("status");
let about: Option<Vec<u8>> = row.get("about");
let sub_level: i32 = row.get("sub_level");
let sub_end: i64 = row.get("sub_end");
log!(
"User: {:?}",
(
id,
iota_id,
username,
display.map_or("".to_string(), |d| String::from_utf8_lossy(&d).to_string()),
status.map_or("".to_string(), |s| String::from_utf8_lossy(&s).to_string()),
about.map_or("".to_string(), |a| String::from_utf8_lossy(&a).to_string()),
sub_level,
sub_end
)
);
}
Ok(())
}
// ==========================================================================================
// IOTA
// ==========================================================================================
pub async fn create_new_iota(public_key: String) -> Result<i64, sqlx::Error> {
let new_id = get_register_id().await as i64;
register_complete_iota(new_id, public_key).await?;
Ok(new_id)
}
pub async fn register_complete_iota(id: i64, public_key: String) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("INSERT INTO iotas (id, public_key) VALUES (?, ?)")
.bind(id)
.bind(public_key)
.execute(&pool)
.await?;
Ok(())
}
pub async fn get_iota_by_id(id: i64) -> Result<(i64, String), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
let result = sqlx::query_as::<_, (u64, Vec<u8>)>(
"SELECT id, public_key FROM iotas WHERE id = CAST(? AS UNSIGNED)",
)
.bind(id)
.fetch_optional(&pool)
.await;
match result {
Ok(optional_row) => match optional_row {
Some((id_u64, public_key)) => Ok((
id_u64 as i64,
String::from_utf8_lossy(&public_key).to_string(),
)),
_ => Err(sqlx::Error::RowNotFound),
},
Err(e) => Err(e),
}
}
pub async fn change_iota_key(id: i64, new_key: String) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("UPDATE iotas SET public_key = ? WHERE id = CAST(? AS UNSIGNED)")
.bind(new_key)
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
pub async fn delete_iota(id: i64) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("DELETE FROM iotas WHERE id = CAST(? AS UNSIGNED)")
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
// ==========================================================================================
// OMIKRONS
// ==========================================================================================
pub async fn get_omikron_by_id(id: i64) -> Result<(String, String), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
let row = sqlx::query_as::<_, (Vec<u8>, Vec<u8>)>(
"SELECT public_key, ip_address FROM omikrons WHERE id = CAST(? AS UNSIGNED)",
)
.bind(id)
.fetch_optional(&pool)
.await?;
match row {
Some((public_key, ip_address)) => Ok((
String::from_utf8_lossy(&public_key).to_string(),
String::from_utf8_lossy(&ip_address).to_string(),
)),
_ => Err(sqlx::Error::RowNotFound),
}
}
// ==========================================================================================
// PHI
// ==========================================================================================
pub async fn add_notification(sender_id: i64, receiver_id: i64) -> Result<(), sqlx::Error> {
let db_lock = SQL_DB.read().await;
let pool = db_lock.as_ref().expect("Database pool is not initialized");
sqlx::query(
r#"
INSERT INTO notifications (sender_id, receiver_id, amount)
VALUES (?, ?, 1)
ON DUPLICATE KEY UPDATE amount = amount + 1
"#,
)
.bind(sender_id)
.bind(receiver_id)
.execute(pool)
.await?;
Ok(())
}
pub async fn read_notification(sender_id: i64, receiver_id: i64) -> Result<(), sqlx::Error> {
let db_lock = SQL_DB.read().await;
let pool = db_lock.as_ref().expect("Database pool is not initialized");
sqlx::query(
r#"
DELETE FROM notifications WHERE sender_id = ? AND receiver_id = ?
"#,
)
.bind(sender_id)
.bind(receiver_id)
.execute(pool)
.await?;
Ok(())
}
pub async fn get_notifications(user_id: i64) -> Result<Vec<(i64, i64)>, sqlx::Error> {
let db_lock = SQL_DB.read().await;
let pool = db_lock.as_ref().expect("Database pool is not initialized");
sqlx::query_as::<_, (i64, i64)>(
r#"
SELECT sender_id, amount FROM notifications WHERE receiver_id = ?
"#,
)
.bind(user_id)
.fetch_all(pool)
.await
}

View file

@ -1,144 +0,0 @@
use crate::sql;
use crate::sql::connection_status::UserStatus;
use dashmap::DashMap;
use once_cell::sync::Lazy;
#[derive(Debug, Clone)]
pub struct UserConnection {
pub connection_type: 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<DashMap<i64, Vec<i64>>> = Lazy::new(DashMap::new);
// UserID -> UserStatus
static USER_STATUS_MAP: Lazy<DashMap<i64, UserConnection>> = Lazy::new(DashMap::new);
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 fn untrack_iota_connection(iota_id: i64, omikron_id: i64) -> bool {
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 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_empty {
IOTA_OMIKRON_CONNECTIONS.remove(&iota_id);
}
connections_empty
}
pub fn get_iota_primary_omikron_connection(iota_id: i64) -> Option<i64> {
IOTA_PRIMARY_OMIKRON_CONNECTION.get(&iota_id).map(|v| *v)
}
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: UserStatus, omikron_id: i64) {
USER_STATUS_MAP.insert(
user_id,
UserConnection {
connection_type: status,
omikron_id,
},
);
}
pub fn get_user_status(user_id: i64) -> Option<UserConnection> {
USER_STATUS_MAP.get(&user_id).map(|v| v.clone())
}
pub fn untrack_many_users(user_ids: &[i64]) {
for user_id in user_ids {
USER_STATUS_MAP.remove(user_id);
}
}
pub async fn untrack_omikron(omikron_id: i64) {
let primary_keys_to_remove: Vec<i64> = IOTA_PRIMARY_OMIKRON_CONNECTION
.iter()
.filter(|entry| *entry.value() == omikron_id)
.map(|entry| *entry.key())
.collect();
for key in primary_keys_to_remove {
IOTA_PRIMARY_OMIKRON_CONNECTION.remove(&key);
}
let mut offline_iotas = Vec::new();
let mut primary_to_remove = Vec::new();
// 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)
.map(|p| *p == omikron_id)
.unwrap_or(false)
{
primary_to_remove.push(iota_id);
}
// 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);
}
}

View file

@ -1,2 +0,0 @@
pub mod omikron_connection;
pub mod omikron_manager;

File diff suppressed because it is too large Load diff

View file

@ -1,39 +0,0 @@
use crate::transport::omikron_connection::OmikronConnection;
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(conn: Arc<OmikronConnection>) {
let id = match conn.clone().get_omikron_id().await {
Some(id) => id,
_ => {
conn.close().await;
return;
}
};
if let Some(old) = OMIKRON_CONNECTIONS.insert(id, conn.clone()) {
old.close().await;
}
}
pub async fn remove_omikron(omikron_id: i64) {
OMIKRON_CONNECTIONS.remove(&omikron_id);
}
pub async fn get_random_omikron() -> Result<Arc<OmikronConnection>, ()> {
let mut rng = rand::thread_rng();
let keys: Vec<_> = OMIKRON_CONNECTIONS.iter().map(|e| *e.key()).collect();
if let Some(key) = keys.into_iter().choose(&mut rng) {
if let Some(entry) = OMIKRON_CONNECTIONS.get(&key) {
return Ok(entry.clone());
}
}
Err(())
}

22
src/util/crypto_helper.rs Normal file → Executable file
View file

@ -8,27 +8,36 @@ use sha2::{Digest, Sha256};
use x448::{PublicKey, Secret, SharedSecret};
/// Errors for crypto operations
#[allow(dead_code)]
#[derive(Debug)]
pub enum CryptoError {
Base64Decode,
Base64Decode(base64::DecodeError),
InvalidKey,
AgreementError,
EncryptionError(aes_gcm::Error),
DecryptionError(aes_gcm::Error),
}
impl From<base64::DecodeError> for CryptoError {
fn from(_: base64::DecodeError) -> Self {
CryptoError::Base64Decode
fn from(err: base64::DecodeError) -> Self {
CryptoError::Base64Decode(err)
}
}
pub fn generate_keypair() -> (Secret, PublicKey) {
#[allow(dead_code)]
pub struct KeyPair {
pub secret: Secret,
pub public: PublicKey,
}
#[allow(dead_code)]
pub fn generate_keypair() -> KeyPair {
let mut buf = [0u8; 56];
let mut rng = OsRng;
rng.fill_bytes(&mut buf);
let secret = Secret::from_bytes(&buf).unwrap();
let public = PublicKey::from(&secret);
(secret, public)
KeyPair { secret, public }
}
pub fn public_key_to_base64(pubkey: &PublicKey) -> String {
@ -58,6 +67,7 @@ fn derive_aes_key(shared: &SharedSecret) -> [u8; 32] {
key
}
#[allow(dead_code)]
pub fn encrypt_b64(
base64_secret: &str,
base64_peer_pub: &str,
@ -125,12 +135,14 @@ pub fn decrypt(
Ok(plaintext)
}
#[allow(dead_code)]
pub fn hash_it(input: &str) -> Vec<u8> {
let mut hasher = Sha256::new();
hasher.update(input.as_bytes());
hasher.finalize().to_vec()
}
#[allow(dead_code)]
pub fn hex_hash(input: &str) -> String {
let digest = hash_it(input);
digest.iter().map(|b| format!("{:02x}", b)).collect()

View file

@ -4,11 +4,12 @@ use aes_gcm::{
};
use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64_STD};
use hkdf::Hkdf;
use sha2::{Digest, Sha256};
type HkdfSha256 = sha2::Sha256;
use sha2::{Digest, Sha256 as HashSha256};
use x448::{PublicKey, Secret};
// --- Custom Errors ---
#[derive(Debug)]
#[allow(dead_code)]
pub enum SecurePayloadError {
InvalidBase64,
InvalidHex,
@ -17,19 +18,16 @@ pub enum SecurePayloadError {
InvalidKeyLength,
}
// --- Data Format Enum ---
#[derive(Clone, Copy, Debug)]
#[allow(dead_code)]
pub enum DataFormat {
Raw,
Base64,
Hex,
}
// --- Main Class Structure ---
pub struct SecurePayload {
/// The internal canonical representation is always raw bytes.
inner_data: Vec<u8>,
/// The private key of the user associated with this payload instance.
private_key: Secret,
}
@ -42,8 +40,8 @@ impl Clone for SecurePayload {
}
}
#[allow(dead_code)]
impl SecurePayload {
/// Clear Constructor: Takes data in any format and the user's private key.
pub fn new<S, T: AsRef<[u8]>>(
data: T,
format: DataFormat,
@ -68,12 +66,10 @@ impl SecurePayload {
})
}
/// Helper to get the public key associated with this instance's private key.
pub fn get_public_key(&self) -> [u8; 56] {
*PublicKey::from(&self.private_key).as_bytes()
}
/// Exports the internal data to the requested format
pub fn export(&self, format: DataFormat) -> String {
match format.into() {
DataFormat::Raw => String::from_utf8_lossy(&self.inner_data).to_string(),
@ -82,14 +78,12 @@ impl SecurePayload {
}
}
/// Access raw bytes directly
pub fn get_bytes(&self) -> &[u8] {
&self.inner_data
}
/// Returns the SHA-256 Hash of the data in the requested format
pub fn get_hash(&self, format: DataFormat) -> String {
let mut hasher = Sha256::new();
let mut hasher = HashSha256::new();
hasher.update(&self.inner_data);
let result = hasher.finalize();
@ -100,8 +94,6 @@ impl SecurePayload {
}
}
/// Encrypts the held data for a specific recipient using AES-256-GCM.
/// The message will contain ONLY the ciphertext.
pub fn encrypt_x448<S>(&self, public_key: S) -> Result<SecurePayload, SecurePayloadError>
where
S: Into<PublicKey>,
@ -109,22 +101,15 @@ impl SecurePayload {
let peer_pub = public_key.into();
let shared_secret = self.private_key.as_diffie_hellman(&peer_pub).unwrap();
println!(
"Encryption Shared Secret (Hex): {}",
hex::encode(shared_secret.as_bytes())
);
let hkdf = Hkdf::<HkdfSha256>::new(None, shared_secret.as_bytes());
let mut okm = [0u8; 44];
// 3. Key & Nonce Derivation (HKDF)
// We derive 32 bytes for the key and 12 bytes for a deterministic nonce.
let hkdf = Hkdf::<Sha256>::new(None, shared_secret.as_bytes());
let mut okm = [0u8; 44]; // 32 (Key) + 12 (Nonce)
hkdf.expand(b"x448-aes-gcm-no-overhead", &mut okm)
.map_err(|_| SecurePayloadError::EncryptionError)?;
let key = &okm[..32];
let nonce_bytes = &okm[32..];
// 4. Encrypt with AES-256-GCM
let cipher = Aes256Gcm::new(key.into());
let nonce = Nonce::from_slice(nonce_bytes);
@ -138,40 +123,33 @@ impl SecurePayload {
)
.map_err(|_| SecurePayloadError::EncryptionError)?;
// 5. Result is ONLY the ciphertext. No key or nonce is packed.
Ok(SecurePayload {
inner_data: ciphertext,
private_key: Secret::from_bytes(self.private_key.as_bytes()).unwrap(),
})
}
/// Decrypts the held data providing the sender's public key manually.
pub fn decrypt_to_format(
&self,
peer_public_key_bytes: &[u8; 56],
output_format: DataFormat,
) -> Result<String, SecurePayloadError> {
let decrypted_instance = self.decrypt_x448(peer_public_key_bytes)?;
let decrypted_instance =
self.decrypt_x448(PublicKey::from_bytes(peer_public_key_bytes).unwrap())?;
Ok(decrypted_instance.export(output_format))
}
/// Decrypts the held data using the internal Private Key and the provided Peer Public Key.
pub fn decrypt_x448(
pub fn decrypt_x448<S>(
&self,
peer_public_key_bytes: &[u8; 56],
) -> Result<SecurePayload, SecurePayloadError> {
// 1. Perform Exchange
let peer_pub = PublicKey::from_bytes(peer_public_key_bytes).unwrap();
peer_public_key_bytes: S,
) -> Result<SecurePayload, SecurePayloadError>
where
S: Into<PublicKey>,
{
let peer_pub = peer_public_key_bytes.into();
let shared_secret = self.private_key.as_diffie_hellman(&peer_pub).unwrap();
// LOGGING: Shared Secret
println!(
"Decryption Shared Secret (Hex): {}",
hex::encode(shared_secret.as_bytes())
);
// 2. Key & Nonce Derivation (Must match encryption exactly)
let hkdf = Hkdf::<Sha256>::new(None, shared_secret.as_bytes());
let hkdf = Hkdf::<HkdfSha256>::new(None, shared_secret.as_bytes());
let mut okm = [0u8; 44];
hkdf.expand(b"x448-aes-gcm-no-overhead", &mut okm)
.map_err(|_| SecurePayloadError::DecryptionError)?;
@ -179,7 +157,6 @@ impl SecurePayload {
let key = &okm[..32];
let nonce_bytes = &okm[32..];
// 3. Decrypt with AES-256-GCM
let cipher = Aes256Gcm::new(key.into());
let nonce = Nonce::from_slice(nonce_bytes);

View file

@ -1,20 +1,9 @@
use std::fs::{self, File};
use std::io::{self, BufReader, Read};
use std::path::{Path, PathBuf};
use tokio::io::AsyncWriteExt;
use uuid::Uuid;
use zip::ZipArchive;
use crate::log;
pub fn delete_file(path: &str, name: &str) -> bool {
let dir = Path::new(&get_directory()).join(path);
let file = dir.join(name);
if !file.exists() {
return false;
}
fs::remove_file(file).is_ok()
}
use crate::util::logger::PrintType;
#[allow(dead_code)]
pub fn delete_directory(path: &str) -> bool {
@ -29,9 +18,11 @@ fn delete_dir_recursive(directory: &Path) -> bool {
}
if let Err(e) = fs::remove_dir_all(directory) {
log!(
0,
PrintType::General,
"[IMPORTANT] Couldn't delete directory {}: {}",
directory.display(),
e
e,
);
return false;
}
@ -39,13 +30,14 @@ fn delete_dir_recursive(directory: &Path) -> bool {
}
#[allow(dead_code)]
pub fn delete_user_directory(user_id: Uuid) {
pub fn delete_user_directory(user_id: i64) {
let user_dir = Path::new(&get_directory())
.join("users")
.join(user_id.to_string());
let _ = delete_dir_recursive(&user_dir);
}
#[allow(dead_code)]
pub fn load_file_buf(path: &str, name: &str) -> io::Result<BufReader<File>> {
let dir = Path::new(&get_directory()).join(path);
let file_path = dir.join(name);
@ -72,6 +64,7 @@ pub fn load_file_buf(path: &str, name: &str) -> io::Result<BufReader<File>> {
let file = File::open(&file_path)?;
Ok(BufReader::new(file))
}
#[allow(dead_code)]
pub fn has_file(path: &str, name: &str) -> bool {
let dir = Path::new(&get_directory()).join(path);
let file_path = dir.join(name);
@ -86,6 +79,7 @@ pub fn has_file(path: &str, name: &str) -> bool {
true
}
#[allow(dead_code)]
pub fn has_dir(path: &str) -> bool {
let dir = Path::new(&get_directory()).join(path);
@ -96,13 +90,19 @@ pub fn has_dir(path: &str) -> bool {
true
}
#[allow(dead_code)]
pub fn load_file(path: &str, name: &str) -> String {
let dir = Path::new(&get_directory()).join(path);
let file_path = dir.join(name);
if !dir.exists() {
if let Err(e) = fs::create_dir_all(&dir) {
log!("[IMPORTANT] Couldn't create directories: {}", e);
log!(
0,
PrintType::General,
"[IMPORTANT] Couldn't create directories: {}",
e
);
return String::new();
}
return String::new();
@ -110,7 +110,12 @@ pub fn load_file(path: &str, name: &str) -> String {
if !file_path.exists() {
if let Err(e) = File::create(&file_path) {
log!("[IMPORTANT] Couldn't create file: {}", e);
log!(
0,
PrintType::General,
"[IMPORTANT] Couldn't create file: {}",
e
);
}
return String::new();
}
@ -129,19 +134,27 @@ pub fn load_file_vec(path: &str, name: &str) -> Result<Vec<u8>, std::io::Error>
std::fs::read(file_path)
}
#[allow(dead_code)]
pub fn save_file(path: &str, name: &str, value: &str) {
let dir = Path::new(&get_directory()).join(path);
let file_path = dir.join(name);
if !dir.exists() {
if let Err(e) = fs::create_dir_all(&dir) {
log!("[IMPORTANT] Couldn't create directories: {}", e);
log!(
0,
PrintType::General,
"[IMPORTANT] Couldn't create directories: {}",
e
);
return;
}
}
if let Err(e) = fs::write(&file_path, value) {
log!(
0,
PrintType::General,
"[IMPORTANT] Couldn't write file {}: {}",
file_path.display(),
e
@ -149,6 +162,7 @@ pub fn save_file(path: &str, name: &str, value: &str) {
}
}
#[allow(dead_code)]
pub fn get_children(path: &str) -> Vec<String> {
let dir = Path::new(&get_directory()).join(path);
let mut children = Vec::new();
@ -169,121 +183,3 @@ pub fn get_directory() -> String {
.to_string_lossy()
.to_string()
}
// Helper to download the zip file content to a file on disk
#[allow(dead_code)]
async fn download_zip(url: &str, as_name: &Path) -> Result<(), Box<dyn std::error::Error>> {
let response = reqwest::get(url).await?;
// Check for successful response status
if !response.status().is_success() {
return Err(format!("Failed to download file: Status {}", response.status()).into());
}
let mut zip_file = tokio::fs::File::create(as_name).await?;
let body = response.bytes().await?;
zip_file.write_all(&body).await?;
Ok(())
}
#[allow(dead_code, deprecated)]
fn extract_zip_contents_to_folder(
zip_path: &Path,
target_dir: &Path,
) -> Result<(), Box<dyn std::error::Error>> {
let file = File::open(zip_path)?;
let mut archive = ZipArchive::new(file)?;
let staging_dir = target_dir.with_extension("staging");
let _ = fs::remove_dir_all(&staging_dir);
fs::create_dir_all(&staging_dir)?;
let mut first_item_name: Option<PathBuf> = None;
for i in 0..archive.len() {
let mut file = archive.by_index(i)?;
let entry_path = staging_dir.join(file.sanitized_name());
if i == 0 {
if file.name().ends_with('/') || file.sanitized_name().components().count() == 1 {
first_item_name = Some(file.sanitized_name());
}
}
if file.name().ends_with('/') {
fs::create_dir_all(&entry_path)?;
} else {
if let Some(parent) = entry_path.parent() {
fs::create_dir_all(parent)?;
}
let mut out_file = File::create(entry_path)?;
io::copy(&mut file, &mut out_file)?;
}
}
if let Some(root_path) = first_item_name {
let root_dir = staging_dir.join(&root_path);
if root_dir.is_dir() {
let root_contents_count = fs::read_dir(&staging_dir)?.count();
if root_contents_count == 1
|| (root_contents_count > 1 && fs::metadata(&root_dir).is_ok())
{
let _ = fs::remove_dir_all(target_dir);
fs::create_dir_all(target_dir)?;
for entry in fs::read_dir(root_dir)? {
let entry = entry?;
let src = entry.path();
let dest = target_dir.join(entry.file_name());
if let Err(_) = fs::rename(&src, &dest) {
if src.is_file() {
fs::copy(&src, &dest)?;
} else {
if entry.path().is_dir() {
fs::rename(&src, &dest)?;
}
}
}
}
let _ = fs::remove_dir_all(&staging_dir);
return Ok(());
}
}
}
log!("Extracting directly (no single root folder detected).");
let _ = fs::remove_dir_all(target_dir);
fs::rename(&staging_dir, target_dir)?;
Ok(())
}
#[allow(dead_code)]
pub async fn download_and_extract_zip(url: &str, as_name: &str) {
let base_dir = PathBuf::from(get_directory());
let zip_filename = format!("{}.zip", Uuid::new_v4());
let zip_path = base_dir.join(&zip_filename);
let target_dir = base_dir.join(as_name);
if let Err(e) = download_zip(url, &zip_path).await {
log!("Error downloading file: {}", e);
return;
}
let zip_path_clone = zip_path.clone();
let target_dir_clone = target_dir.clone();
let extract_result = extract_zip_contents_to_folder(&zip_path_clone, &target_dir_clone);
if let Err(e) = extract_result {
log!("Panic during ZIP extraction: {}", e);
}
if let Err(e) = tokio::fs::remove_file(&zip_path).await {
log!("Error cleaning up ZIP file {}: {}", zip_path.display(), e);
}
}

View file

@ -1,5 +1,5 @@
use std::{
collections::{BTreeMap, HashMap},
collections::BTreeMap,
fs::{self, OpenOptions},
io::Write,
path::Path,
@ -9,13 +9,12 @@ use std::{
};
use ansi_term::Color;
use epsilon_core::{CommunicationValue, DataTypes, DataValue};
use json::JsonValue;
use ttp_core::{CommunicationValue, DataTypes, DataValue};
static LOGGER: OnceLock<mpsc::Sender<LogMessage>> = OnceLock::new();
#[allow(dead_code)]
#[derive(Clone, Copy)]
#[allow(unused)]
pub enum PrintType {
Call,
Client,
@ -96,12 +95,14 @@ fn fixed_box(content: &str, width: usize) -> String {
}
pub fn log_internal(
sender: Option<i64>,
sender: i64,
kind: PrintType,
prefix: &'static str,
is_error: bool,
message: String,
) {
let sender = if sender == 0 { None } else { Some(sender) };
if let Some(tx) = LOGGER.get() {
let _ = tx.send(LogMessage {
timestamp_ms: SystemTime::now()
@ -116,102 +117,28 @@ pub fn log_internal(
});
}
}
#[macro_export]
macro_rules! log {
// plain
($($arg:tt)*) => {
$crate::util::logger::log_internal(
None,
$crate::util::logger::PrintType::General,
"",
false,
format!($($arg)*)
)
};
// sender + actor
($sender:expr, $kind:expr, $($arg:tt)*) => {
$crate::util::logger::log_internal(Some($sender), $kind, "", false, format!($($arg)*))
};
// actor only
($kind:expr, $($arg:tt)*) => {
$crate::util::logger::log_internal(None, $kind, "", false, format!($($arg)*))
($sender: expr, $kind:expr, $($arg:tt)*) => {
$crate::util::logger::log_internal($sender, $kind, "", false, format!($($arg)*))
};
}
#[macro_export]
macro_rules! log_in {
// sender + actor
($sender:expr, $kind:expr, $($arg:tt)*) => {
$crate::util::logger::log_internal(Some($sender), $kind, ">", false, format!($($arg)*))
};
// actor only
($kind:expr, $($arg:tt)*) => {
$crate::util::logger::log_internal(None, $kind, ">", false, format!($($arg)*))
};
// plain
($($arg:tt)*) => {
$crate::util::logger::log_internal(
None,
$crate::util::logger::PrintType::General,
">",
false,
format!($($arg)*)
)
($sender: expr, $kind:expr, $($arg:tt)*) => {
$crate::util::logger::log_internal($sender, $kind, ">", false, format!($($arg)*))
};
}
#[macro_export]
macro_rules! log_out {
// sender + actor
($sender:expr, $kind:expr, $($arg:tt)*) => {
$crate::util::logger::log_internal(Some($sender), $kind, "<", false, format!($($arg)*))
};
// actor only
($kind:expr, $($arg:tt)*) => {
$crate::util::logger::log_internal(None, $kind, "<", false, format!($($arg)*))
};
// plain
($($arg:tt)*) => {
$crate::util::logger::log_internal(
None,
$crate::util::logger::PrintType::General,
"<",
false,
format!($($arg)*)
)
$crate::util::logger::log_internal($sender, $kind, "<", false, format!($($arg)*))
};
}
#[macro_export]
macro_rules! log_err {
// sender + actor
($sender:expr, $kind:expr, $($arg:tt)*) => {
$crate::util::logger::log_internal(Some($sender), $kind, ">>", true, format!($($arg)*))
};
// actor only
($kind:expr, $($arg:tt)*) => {
$crate::util::logger::log_internal(None, $kind, ">>", true, format!($($arg)*))
};
// plain
($($arg:tt)*) => {
$crate::util::logger::log_internal(
None,
$crate::util::logger::PrintType::General,
">>",
true,
format!($($arg)*)
)
$crate::util::logger::log_internal($sender, $kind, ">>", true, format!($($arg)*))
};
}
@ -224,7 +151,7 @@ pub fn log_cv_internal(
let formatted = format_cv(cv);
log_internal(
Some(cv.get_sender() as i64),
cv.get_sender() as i64,
print_type.unwrap_or(PrintType::General),
prefix,
false,
@ -322,6 +249,7 @@ fn format_array(arr: Vec<DataValue>) -> String {
parts.join(", ")
}
#[macro_export]
macro_rules! log_cv {
($kind:expr, $cv:expr) => {