Sync
This commit is contained in:
parent
b16045eb4f
commit
9e6850b153
31 changed files with 8340 additions and 0 deletions
19
.gitignore
vendored
Normal file
19
.gitignore
vendored
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
# Generated by Cargo
|
||||
# will have compiled files and executables
|
||||
debug
|
||||
target
|
||||
|
||||
# These are backup files generated by rustfmt
|
||||
**/*.rs.bk
|
||||
|
||||
# MSVC Windows builds of rustc generate these, which store debugging information
|
||||
*.pdb
|
||||
*.env
|
||||
# Generated by cargo mutants
|
||||
# Contains mutation testing data
|
||||
**/mutants.out*/W
|
||||
|
||||
|
||||
# Added by cargo
|
||||
/target
|
||||
/logs
|
||||
3910
Cargo.lock
generated
Normal file
3910
Cargo.lock
generated
Normal file
File diff suppressed because it is too large
Load diff
37
Cargo.toml
Normal file
37
Cargo.toml
Normal file
|
|
@ -0,0 +1,37 @@
|
|||
[package]
|
||||
name = "Omikron"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
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"] }
|
||||
|
||||
base64 = "0.22.1"
|
||||
dashmap = "*"
|
||||
futures = "*"
|
||||
hex = "*"
|
||||
once_cell = "1.21.3"
|
||||
rand = "0.8"
|
||||
rand_core = { version = "0.6", features = ["getrandom", "std"] }
|
||||
rustls = { version = "0.23.37", default-features = false, features = [
|
||||
"std",
|
||||
"tls12",
|
||||
"aws-lc-rs",
|
||||
"prefer-post-quantum",
|
||||
] }
|
||||
aes-gcm = "0.10.3"
|
||||
sha2 = "0.10.9"
|
||||
tokio = { version = "*", features = ["full"] }
|
||||
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"
|
||||
16
LICENSE
Normal file
16
LICENSE
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
Copyright (c) [2025] [Methanium]
|
||||
All rights reserved.
|
||||
|
||||
This software is protected by copyright. Copying, editing,
|
||||
distributing, publicly performing, or any other use of this software
|
||||
or its components, in source or binary form, is strictly prohibited without the express
|
||||
written permission of the copyright holder.
|
||||
|
||||
FUTURE LICENSE ACCEPTANCE:
|
||||
It is the copyright holder's intention to release this software in the future
|
||||
under a license yet to be defined, which will, among other things,
|
||||
allow private, non-commercial use. This statement does not constitute
|
||||
a current license grant and does not alter the above
|
||||
prohibition on use, copying, or modification. Until the formal
|
||||
publication of such a future license, all rights remain
|
||||
reserved.
|
||||
1
README.md
Normal file
1
README.md
Normal file
|
|
@ -0,0 +1 @@
|
|||
# ⚠️ Moved to [git.methanium.net](https://git.methanium.net/tensamin/omikron) ⚠️
|
||||
519
src/anonymous_clients/anonymous_client_connection.rs
Normal file
519
src/anonymous_clients/anonymous_client_connection.rs
Normal 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(¬ification).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),
|
||||
}
|
||||
}
|
||||
}
|
||||
53
src/anonymous_clients/anonymous_manager.rs
Normal file
53
src/anonymous_clients/anonymous_manager.rs
Normal 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)
|
||||
}
|
||||
2
src/anonymous_clients/mod.rs
Normal file
2
src/anonymous_clients/mod.rs
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
pub mod anonymous_client_connection;
|
||||
pub mod anonymous_manager;
|
||||
106
src/calls/call_group.rs
Normal file
106
src/calls/call_group.rs
Normal 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
97
src/calls/call_manager.rs
Normal 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
151
src/calls/call_util.rs
Normal 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
49
src/calls/caller.rs
Normal 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
4
src/calls/mod.rs
Normal 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
1
src/data/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod user;
|
||||
30
src/data/user.rs
Normal file
30
src/data/user.rs
Normal file
|
|
@ -0,0 +1,30 @@
|
|||
use strum::IntoEnumIterator;
|
||||
use strum_macros::EnumIter;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, EnumIter, Eq)]
|
||||
#[allow(unused, non_camel_case_types)]
|
||||
pub enum UserStatus {
|
||||
user_offline,
|
||||
user_online,
|
||||
user_dnd,
|
||||
user_idle,
|
||||
user_wc,
|
||||
user_borked,
|
||||
iota_offline,
|
||||
iota_online,
|
||||
iota_borked,
|
||||
}
|
||||
#[allow(unused)]
|
||||
impl UserStatus {
|
||||
pub fn to_string(&self) -> String {
|
||||
format!("{:?}", self)
|
||||
}
|
||||
pub fn from_str(s: &str) -> Option<UserStatus> {
|
||||
for sel in UserStatus::iter() {
|
||||
if &sel.to_string() == s {
|
||||
return Some(sel);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
}
|
||||
51
src/main.rs
Normal file
51
src/main.rs
Normal file
|
|
@ -0,0 +1,51 @@
|
|||
mod anonymous_clients;
|
||||
mod calls;
|
||||
mod data;
|
||||
mod omega;
|
||||
mod rho;
|
||||
mod util;
|
||||
|
||||
use std::env;
|
||||
|
||||
use dotenv::dotenv;
|
||||
use once_cell::sync::Lazy;
|
||||
use rustls::crypto::aws_lc_rs::default_provider;
|
||||
|
||||
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 {
|
||||
load_secret_key(&*PRIVATE_KEY).unwrap()
|
||||
}
|
||||
static PUBLIC_KEY: Lazy<String> = Lazy::new(|| env::var("PUBLIC_KEY").unwrap());
|
||||
pub fn get_public_key() -> x448::PublicKey {
|
||||
load_public_key(&*PUBLIC_KEY).unwrap()
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
if let Err(_) = default_provider().install_default() {
|
||||
println!("Error loading Provider");
|
||||
return;
|
||||
}
|
||||
dotenv().ok();
|
||||
startup();
|
||||
|
||||
get_omega_connection();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = start(959).await {
|
||||
log_err!(0, util::logger::PrintType::General, "{}", e);
|
||||
}
|
||||
});
|
||||
garbage_collect_calls();
|
||||
|
||||
tokio::signal::ctrl_c().await.unwrap();
|
||||
}
|
||||
1
src/omega/mod.rs
Normal file
1
src/omega/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod omega_connection;
|
||||
741
src/omega/omega_connection.rs
Normal file
741
src/omega/omega_connection.rs
Normal 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()
|
||||
}
|
||||
41
src/omega/ping_pong_task.rs
Normal file
41
src/omega/ping_pong_task.rs
Normal 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;
|
||||
}
|
||||
}
|
||||
}
|
||||
583
src/rho/client_connection.rs
Normal file
583
src/rho/client_connection.rs
Normal 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(¬ification).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
358
src/rho/connection.rs
Executable 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(¬ify).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(¬ify).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
461
src/rho/iota_connection.rs
Executable 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
6
src/rho/mod.rs
Normal 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
199
src/rho/rho_connection.rs
Normal 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(¬ification).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
83
src/rho/rho_manager.rs
Normal 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
25
src/rho/server.rs
Normal 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(())
|
||||
}
|
||||
149
src/util/crypto_helper.rs
Executable file
149
src/util/crypto_helper.rs
Executable file
|
|
@ -0,0 +1,149 @@
|
|||
use aes_gcm::{
|
||||
Aes256Gcm, Nonce,
|
||||
aead::{Aead, KeyInit, OsRng},
|
||||
};
|
||||
use base64::{Engine as _, engine::general_purpose::STANDARD};
|
||||
use rand_core::RngCore;
|
||||
use sha2::{Digest, Sha256};
|
||||
use x448::{PublicKey, Secret, SharedSecret};
|
||||
|
||||
/// Errors for crypto operations
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug)]
|
||||
pub enum CryptoError {
|
||||
Base64Decode(base64::DecodeError),
|
||||
InvalidKey,
|
||||
AgreementError,
|
||||
EncryptionError(aes_gcm::Error),
|
||||
DecryptionError(aes_gcm::Error),
|
||||
}
|
||||
|
||||
impl From<base64::DecodeError> for CryptoError {
|
||||
fn from(err: base64::DecodeError) -> Self {
|
||||
CryptoError::Base64Decode(err)
|
||||
}
|
||||
}
|
||||
|
||||
#[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);
|
||||
KeyPair { secret, public }
|
||||
}
|
||||
|
||||
pub fn public_key_to_base64(pubkey: &PublicKey) -> String {
|
||||
STANDARD.encode(pubkey.as_bytes().as_ref())
|
||||
}
|
||||
|
||||
pub fn secret_key_to_base64(secret: &Secret) -> String {
|
||||
STANDARD.encode(secret.as_bytes().as_ref())
|
||||
}
|
||||
|
||||
pub fn load_public_key(base64_pub: &str) -> Option<PublicKey> {
|
||||
let bytes = STANDARD.decode(base64_pub).unwrap();
|
||||
PublicKey::from_bytes(&bytes)
|
||||
}
|
||||
|
||||
pub fn load_secret_key(base64_secret: &str) -> Option<Secret> {
|
||||
let bytes = STANDARD.decode(base64_secret).unwrap();
|
||||
Secret::from_bytes(&bytes)
|
||||
}
|
||||
|
||||
fn derive_aes_key(shared: &SharedSecret) -> [u8; 32] {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(shared.as_bytes());
|
||||
let result = hasher.finalize();
|
||||
let mut key = [0u8; 32];
|
||||
key.copy_from_slice(&result[..32]);
|
||||
key
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn encrypt_b64(
|
||||
base64_secret: &str,
|
||||
base64_peer_pub: &str,
|
||||
plaintext: &str,
|
||||
) -> Result<String, CryptoError> {
|
||||
let secret = load_secret_key(base64_secret).unwrap();
|
||||
let peer_pub = load_public_key(base64_peer_pub).unwrap();
|
||||
encrypt(secret, peer_pub, plaintext)
|
||||
}
|
||||
pub fn encrypt(
|
||||
secret: Secret,
|
||||
peer_pub: PublicKey,
|
||||
plaintext: &str,
|
||||
) -> Result<String, CryptoError> {
|
||||
let shared = secret
|
||||
.to_diffie_hellman(&peer_pub)
|
||||
.ok_or(CryptoError::AgreementError)?;
|
||||
let key_bytes = derive_aes_key(&shared);
|
||||
let cipher = Aes256Gcm::new_from_slice(&key_bytes).expect("Key length should be correct");
|
||||
let mut nonce_bytes = [0u8; 12];
|
||||
OsRng.fill_bytes(&mut nonce_bytes);
|
||||
let nonce = Nonce::from_slice(&nonce_bytes);
|
||||
let ciphertext = cipher
|
||||
.encrypt(nonce, plaintext.as_bytes())
|
||||
.map_err(CryptoError::EncryptionError)?;
|
||||
// prefix nonce to ciphertext
|
||||
let mut out = Vec::with_capacity(nonce_bytes.len() + ciphertext.len());
|
||||
out.extend_from_slice(&nonce_bytes);
|
||||
out.extend_from_slice(&ciphertext);
|
||||
Ok(STANDARD.encode(&out))
|
||||
}
|
||||
|
||||
pub fn decrypt_b64(
|
||||
base64_secret: &str,
|
||||
base64_peer_pub: &str,
|
||||
encrypted_base64: &str,
|
||||
) -> Result<String, CryptoError> {
|
||||
let secret = load_secret_key(base64_secret).unwrap();
|
||||
let peer_pub = load_public_key(base64_peer_pub).unwrap();
|
||||
decrypt(secret, peer_pub, encrypted_base64)
|
||||
}
|
||||
pub fn decrypt(
|
||||
secret: Secret,
|
||||
peer_pub: PublicKey,
|
||||
encrypted_base64: &str,
|
||||
) -> Result<String, CryptoError> {
|
||||
let shared = secret
|
||||
.to_diffie_hellman(&peer_pub)
|
||||
.ok_or(CryptoError::AgreementError)?;
|
||||
let key_bytes = derive_aes_key(&shared);
|
||||
let cipher = Aes256Gcm::new_from_slice(&key_bytes).expect("Key length should be correct");
|
||||
|
||||
let encrypted = STANDARD.decode(encrypted_base64)?;
|
||||
if encrypted.len() < 12 {
|
||||
return Err(CryptoError::DecryptionError(aes_gcm::Error));
|
||||
}
|
||||
let nonce_bytes = &encrypted[..12];
|
||||
let ciphertext = &encrypted[12..];
|
||||
let nonce = Nonce::from_slice(nonce_bytes);
|
||||
let plaintext_bytes = cipher
|
||||
.decrypt(nonce, ciphertext)
|
||||
.map_err(CryptoError::DecryptionError)?;
|
||||
let plaintext = String::from_utf8(plaintext_bytes)
|
||||
.map_err(|_| CryptoError::DecryptionError(aes_gcm::Error))?;
|
||||
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()
|
||||
}
|
||||
178
src/util/crypto_util.rs
Normal file
178
src/util/crypto_util.rs
Normal file
|
|
@ -0,0 +1,178 @@
|
|||
use aes_gcm::{
|
||||
Aes256Gcm, Nonce,
|
||||
aead::{Aead, KeyInit, Payload},
|
||||
};
|
||||
use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64_STD};
|
||||
use hkdf::Hkdf;
|
||||
type HkdfSha256 = sha2::Sha256;
|
||||
use sha2::{Digest, Sha256 as HashSha256};
|
||||
use x448::{PublicKey, Secret};
|
||||
|
||||
#[derive(Debug)]
|
||||
#[allow(dead_code)]
|
||||
pub enum SecurePayloadError {
|
||||
InvalidBase64,
|
||||
InvalidHex,
|
||||
EncryptionError,
|
||||
DecryptionError,
|
||||
InvalidKeyLength,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
#[allow(dead_code)]
|
||||
pub enum DataFormat {
|
||||
Raw,
|
||||
Base64,
|
||||
Hex,
|
||||
}
|
||||
|
||||
pub struct SecurePayload {
|
||||
inner_data: Vec<u8>,
|
||||
private_key: Secret,
|
||||
}
|
||||
|
||||
impl Clone for SecurePayload {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
inner_data: self.inner_data.clone(),
|
||||
private_key: Secret::from_bytes(self.private_key.as_bytes()).unwrap(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
impl SecurePayload {
|
||||
pub fn new<S, T: AsRef<[u8]>>(
|
||||
data: T,
|
||||
format: DataFormat,
|
||||
private_key: S,
|
||||
) -> Result<Self, SecurePayloadError>
|
||||
where
|
||||
S: Into<Secret>,
|
||||
{
|
||||
let raw_data = match format {
|
||||
DataFormat::Raw => data.as_ref().to_vec(),
|
||||
DataFormat::Base64 => BASE64_STD
|
||||
.decode(data.as_ref())
|
||||
.map_err(|_| SecurePayloadError::InvalidBase64)?,
|
||||
DataFormat::Hex => {
|
||||
hex::decode(data.as_ref()).map_err(|_| SecurePayloadError::InvalidHex)?
|
||||
}
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
inner_data: raw_data,
|
||||
private_key: private_key.into(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn get_public_key(&self) -> [u8; 56] {
|
||||
*PublicKey::from(&self.private_key).as_bytes()
|
||||
}
|
||||
|
||||
pub fn export(&self, format: DataFormat) -> String {
|
||||
match format.into() {
|
||||
DataFormat::Raw => String::from_utf8_lossy(&self.inner_data).to_string(),
|
||||
DataFormat::Base64 => BASE64_STD.encode(&self.inner_data),
|
||||
DataFormat::Hex => hex::encode(&self.inner_data),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_bytes(&self) -> &[u8] {
|
||||
&self.inner_data
|
||||
}
|
||||
|
||||
pub fn get_hash(&self, format: DataFormat) -> String {
|
||||
let mut hasher = HashSha256::new();
|
||||
hasher.update(&self.inner_data);
|
||||
let result = hasher.finalize();
|
||||
|
||||
match format {
|
||||
DataFormat::Raw => String::from_utf8_lossy(&result).to_string(),
|
||||
DataFormat::Base64 => BASE64_STD.encode(result),
|
||||
DataFormat::Hex => hex::encode(result),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn encrypt_x448<S>(&self, public_key: S) -> Result<SecurePayload, SecurePayloadError>
|
||||
where
|
||||
S: Into<PublicKey>,
|
||||
{
|
||||
let peer_pub = public_key.into();
|
||||
let shared_secret = self.private_key.as_diffie_hellman(&peer_pub).unwrap();
|
||||
|
||||
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::EncryptionError)?;
|
||||
|
||||
let key = &okm[..32];
|
||||
let nonce_bytes = &okm[32..];
|
||||
|
||||
let cipher = Aes256Gcm::new(key.into());
|
||||
let nonce = Nonce::from_slice(nonce_bytes);
|
||||
|
||||
let ciphertext = cipher
|
||||
.encrypt(
|
||||
nonce,
|
||||
Payload {
|
||||
msg: &self.inner_data,
|
||||
aad: &[],
|
||||
},
|
||||
)
|
||||
.map_err(|_| SecurePayloadError::EncryptionError)?;
|
||||
|
||||
Ok(SecurePayload {
|
||||
inner_data: ciphertext,
|
||||
private_key: Secret::from_bytes(self.private_key.as_bytes()).unwrap(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn decrypt_to_format(
|
||||
&self,
|
||||
peer_public_key_bytes: &[u8; 56],
|
||||
output_format: DataFormat,
|
||||
) -> Result<String, SecurePayloadError> {
|
||||
let decrypted_instance =
|
||||
self.decrypt_x448(PublicKey::from_bytes(peer_public_key_bytes).unwrap())?;
|
||||
Ok(decrypted_instance.export(output_format))
|
||||
}
|
||||
|
||||
pub fn decrypt_x448<S>(
|
||||
&self,
|
||||
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();
|
||||
|
||||
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)?;
|
||||
|
||||
let key = &okm[..32];
|
||||
let nonce_bytes = &okm[32..];
|
||||
|
||||
let cipher = Aes256Gcm::new(key.into());
|
||||
let nonce = Nonce::from_slice(nonce_bytes);
|
||||
|
||||
let plaintext = cipher
|
||||
.decrypt(
|
||||
nonce,
|
||||
Payload {
|
||||
msg: &self.inner_data,
|
||||
aad: &[],
|
||||
},
|
||||
)
|
||||
.map_err(|_| SecurePayloadError::DecryptionError)?;
|
||||
|
||||
Ok(SecurePayload {
|
||||
inner_data: plaintext,
|
||||
private_key: Secret::from_bytes(self.private_key.as_bytes()).unwrap(),
|
||||
})
|
||||
}
|
||||
}
|
||||
185
src/util/file_util.rs
Normal file
185
src/util/file_util.rs
Normal file
|
|
@ -0,0 +1,185 @@
|
|||
use std::fs::{self, File};
|
||||
use std::io::{self, BufReader, Read};
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use crate::log;
|
||||
use crate::util::logger::PrintType;
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn delete_directory(path: &str) -> bool {
|
||||
let dir = Path::new(&get_directory()).join(path);
|
||||
delete_dir_recursive(&dir)
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
fn delete_dir_recursive(directory: &Path) -> bool {
|
||||
if !directory.exists() {
|
||||
return false;
|
||||
}
|
||||
if let Err(e) = fs::remove_dir_all(directory) {
|
||||
log!(
|
||||
0,
|
||||
PrintType::General,
|
||||
"[IMPORTANT] Couldn't delete directory {}: {}",
|
||||
directory.display(),
|
||||
e,
|
||||
);
|
||||
return false;
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
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);
|
||||
|
||||
// Ensure the directory exists, create if necessary
|
||||
if !dir.exists() {
|
||||
if let Err(_) = fs::create_dir_all(&dir) {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::NotFound,
|
||||
"Directory creation failed",
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
// Create the file if it doesn't exist
|
||||
if !file_path.exists() {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::NotFound,
|
||||
"File creation failed",
|
||||
));
|
||||
}
|
||||
|
||||
// Open the file and return a BufReader for efficient reading
|
||||
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);
|
||||
|
||||
if !dir.exists() {
|
||||
return false;
|
||||
}
|
||||
|
||||
if !file_path.exists() {
|
||||
return false;
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
#[allow(dead_code)]
|
||||
pub fn has_dir(path: &str) -> bool {
|
||||
let dir = Path::new(&get_directory()).join(path);
|
||||
|
||||
if !dir.exists() {
|
||||
return false;
|
||||
}
|
||||
|
||||
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!(
|
||||
0,
|
||||
PrintType::General,
|
||||
"[IMPORTANT] Couldn't create directories: {}",
|
||||
e
|
||||
);
|
||||
return String::new();
|
||||
}
|
||||
return String::new();
|
||||
}
|
||||
|
||||
if !file_path.exists() {
|
||||
if let Err(e) = File::create(&file_path) {
|
||||
log!(
|
||||
0,
|
||||
PrintType::General,
|
||||
"[IMPORTANT] Couldn't create file: {}",
|
||||
e
|
||||
);
|
||||
}
|
||||
return String::new();
|
||||
}
|
||||
|
||||
let mut content = String::new();
|
||||
if let Ok(mut f) = File::open(&file_path) {
|
||||
let _ = f.read_to_string(&mut content);
|
||||
}
|
||||
content
|
||||
}
|
||||
|
||||
pub fn load_file_vec(path: &str, name: &str) -> Result<Vec<u8>, std::io::Error> {
|
||||
let dir = Path::new(&get_directory()).join(path);
|
||||
let file_path = dir.join(name);
|
||||
|
||||
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!(
|
||||
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
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn get_children(path: &str) -> Vec<String> {
|
||||
let dir = Path::new(&get_directory()).join(path);
|
||||
let mut children = Vec::new();
|
||||
if let Ok(entries) = fs::read_dir(&dir) {
|
||||
for entry in entries {
|
||||
if let Ok(entry) = entry {
|
||||
children.push(entry.file_name().to_string_lossy().to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
children
|
||||
}
|
||||
|
||||
pub fn get_directory() -> String {
|
||||
let exe = std::env::current_exe().unwrap_or_else(|_| PathBuf::from("."));
|
||||
exe.parent()
|
||||
.unwrap_or(Path::new("."))
|
||||
.to_string_lossy()
|
||||
.to_string()
|
||||
}
|
||||
280
src/util/logger.rs
Normal file
280
src/util/logger.rs
Normal file
|
|
@ -0,0 +1,280 @@
|
|||
use std::{
|
||||
collections::BTreeMap,
|
||||
fs::{self, OpenOptions},
|
||||
io::Write,
|
||||
path::Path,
|
||||
sync::{OnceLock, mpsc},
|
||||
thread,
|
||||
time::{SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use ansi_term::Color;
|
||||
use ttp_core::{CommunicationValue, DataTypes, DataValue};
|
||||
|
||||
static LOGGER: OnceLock<mpsc::Sender<LogMessage>> = OnceLock::new();
|
||||
|
||||
#[allow(dead_code)]
|
||||
#[derive(Clone, Copy)]
|
||||
pub enum PrintType {
|
||||
Call,
|
||||
Client,
|
||||
Iota,
|
||||
Omikron,
|
||||
Omega,
|
||||
General,
|
||||
}
|
||||
|
||||
struct LogMessage {
|
||||
timestamp_ms: u128,
|
||||
sender: Option<i64>,
|
||||
prefix: &'static str,
|
||||
kind: PrintType,
|
||||
is_error: bool,
|
||||
message: String,
|
||||
}
|
||||
|
||||
pub fn startup() {
|
||||
let (tx, rx) = mpsc::channel::<LogMessage>();
|
||||
LOGGER.set(tx).expect("Logger already initialized");
|
||||
|
||||
thread::spawn(move || {
|
||||
let log_dir = Path::new("logs");
|
||||
fs::create_dir_all(log_dir).expect("Failed to create log directory");
|
||||
|
||||
let start_ts = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs();
|
||||
|
||||
let path = log_dir.join(format!("log_{}.txt", start_ts));
|
||||
let mut file = OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(path)
|
||||
.expect("Failed to open log file");
|
||||
|
||||
for msg in rx {
|
||||
let ts = fixed_box(&msg.timestamp_ms.to_string(), 13);
|
||||
let sender = match msg.sender {
|
||||
Some(id) => fixed_box(&id.to_string(), 19),
|
||||
_ => fixed_box("", 19),
|
||||
};
|
||||
|
||||
let line = format!("{} {} {} {}", ts, sender, msg.prefix, msg.message);
|
||||
|
||||
println!("{}", colorize(msg.kind, msg.is_error).paint(&line));
|
||||
|
||||
let _ = writeln!(file, "{}", line);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn colorize(kind: PrintType, is_error: bool) -> Color {
|
||||
if is_error {
|
||||
return Color::Red;
|
||||
}
|
||||
|
||||
match kind {
|
||||
PrintType::Call => Color::Purple,
|
||||
PrintType::Client => Color::Green,
|
||||
PrintType::Iota => Color::Yellow,
|
||||
PrintType::Omikron => Color::Blue,
|
||||
PrintType::Omega => Color::Cyan,
|
||||
PrintType::General => Color::White,
|
||||
}
|
||||
}
|
||||
|
||||
fn fixed_box(content: &str, width: usize) -> String {
|
||||
let s: String = content.chars().take(width).collect();
|
||||
let len = s.chars().count();
|
||||
if len < width {
|
||||
format!("[{}{}]", " ".repeat(width - len), s)
|
||||
} else {
|
||||
s
|
||||
}
|
||||
}
|
||||
|
||||
pub fn log_internal(
|
||||
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()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_millis(),
|
||||
sender,
|
||||
prefix,
|
||||
kind,
|
||||
is_error,
|
||||
message,
|
||||
});
|
||||
}
|
||||
}
|
||||
#[macro_export]
|
||||
macro_rules! log {
|
||||
($sender: expr, $kind:expr, $($arg:tt)*) => {
|
||||
$crate::util::logger::log_internal($sender, $kind, "", false, format!($($arg)*))
|
||||
};
|
||||
}
|
||||
#[macro_export]
|
||||
macro_rules! log_in {
|
||||
($sender: expr, $kind:expr, $($arg:tt)*) => {
|
||||
$crate::util::logger::log_internal($sender, $kind, ">", false, format!($($arg)*))
|
||||
};
|
||||
}
|
||||
#[macro_export]
|
||||
macro_rules! log_out {
|
||||
($sender:expr, $kind:expr, $($arg:tt)*) => {
|
||||
$crate::util::logger::log_internal($sender, $kind, "<", false, format!($($arg)*))
|
||||
};
|
||||
}
|
||||
#[macro_export]
|
||||
macro_rules! log_err {
|
||||
($sender:expr, $kind:expr, $($arg:tt)*) => {
|
||||
$crate::util::logger::log_internal($sender, $kind, ">>", true, format!($($arg)*))
|
||||
};
|
||||
}
|
||||
|
||||
// ******** COMMUNICATION VALUES ********
|
||||
pub fn log_cv_internal(
|
||||
prefix: &'static str,
|
||||
cv: &CommunicationValue,
|
||||
print_type: Option<PrintType>,
|
||||
) {
|
||||
let formatted = format_cv(cv);
|
||||
|
||||
log_internal(
|
||||
cv.get_sender() as i64,
|
||||
print_type.unwrap_or(PrintType::General),
|
||||
prefix,
|
||||
false,
|
||||
formatted,
|
||||
);
|
||||
}
|
||||
|
||||
pub fn format_cv(cv: &CommunicationValue) -> String {
|
||||
let mut parts = Vec::new();
|
||||
|
||||
let sender = cv.get_sender();
|
||||
let receiver = cv.get_receiver();
|
||||
|
||||
if sender > 0 && receiver > 0 {
|
||||
parts.push(format!("{} > {}", sender, receiver));
|
||||
} else if sender > 0 {
|
||||
parts.push(format!("{}", sender));
|
||||
} else if receiver > 0 {
|
||||
parts.push(format!("> {}", receiver));
|
||||
}
|
||||
|
||||
let comm_type = cv.get_type().to_string();
|
||||
parts.push(format!("{}", comm_type));
|
||||
|
||||
let data: &BTreeMap<DataTypes, DataValue> = cv.get_data_container();
|
||||
|
||||
let formated_data =
|
||||
format_data_container(data.iter().map(|(k, v)| (k.clone(), v.clone())).collect());
|
||||
|
||||
parts.push(format!("{}", formated_data));
|
||||
|
||||
parts.join(": ")
|
||||
}
|
||||
|
||||
fn format_data_container(data: Vec<(DataTypes, DataValue)>) -> String {
|
||||
let parts: Vec<String> = data
|
||||
.into_iter()
|
||||
.map(|(key, value)| {
|
||||
let key_str = key.to_string();
|
||||
|
||||
match value {
|
||||
DataValue::Str(s) => format!("{}=\"{}\"", key_str, s),
|
||||
|
||||
DataValue::Container(inner) => {
|
||||
let inner_formatted = format_data_container(inner);
|
||||
format!("{}={{ {} }}", key_str, inner_formatted)
|
||||
}
|
||||
|
||||
DataValue::Array(arr) => {
|
||||
let arr_formatted = format_array(arr);
|
||||
format!("{}=[{}]", key_str, arr_formatted)
|
||||
}
|
||||
|
||||
DataValue::Bool(b) => format!("{}={}", key_str, b),
|
||||
|
||||
DataValue::BoolTrue => format!("{}=true", key_str),
|
||||
DataValue::BoolFalse => format!("{}=false", key_str),
|
||||
|
||||
DataValue::Number(num) => format!("{}={}", key_str, num),
|
||||
|
||||
_ => "".to_string(),
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
parts.join(", ")
|
||||
}
|
||||
|
||||
fn format_array(arr: Vec<DataValue>) -> String {
|
||||
let parts: Vec<String> = arr
|
||||
.into_iter()
|
||||
.map(|value| match value {
|
||||
DataValue::Str(s) => format!("\"{}\"", s),
|
||||
|
||||
DataValue::Container(inner) => {
|
||||
let inner_formatted = format_data_container(inner);
|
||||
format!("{{ {} }}", inner_formatted)
|
||||
}
|
||||
|
||||
DataValue::Array(inner_arr) => {
|
||||
let formatted = format_array(inner_arr);
|
||||
format!("[{}]", formatted)
|
||||
}
|
||||
|
||||
DataValue::Bool(b) => b.to_string(),
|
||||
|
||||
DataValue::BoolTrue => "true".to_string(),
|
||||
DataValue::BoolFalse => "false".to_string(),
|
||||
|
||||
DataValue::Number(num) => num.to_string(),
|
||||
|
||||
_ => String::new(),
|
||||
})
|
||||
.collect();
|
||||
|
||||
parts.join(", ")
|
||||
}
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! log_cv {
|
||||
($kind:expr, $cv:expr) => {
|
||||
$crate::util::logger::log_cv_internal("", &$cv, Some($kind))
|
||||
};
|
||||
($cv:expr) => {
|
||||
$crate::util::logger::log_cv_internal("", &$cv, None)
|
||||
};
|
||||
}
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! log_cv_in {
|
||||
($kind:expr, $cv:expr) => {
|
||||
$crate::util::logger::log_cv_internal("> ", &$cv, Some($kind))
|
||||
};
|
||||
($cv:expr) => {
|
||||
$crate::util::logger::log_cv_internal("> ", &$cv, None)
|
||||
};
|
||||
}
|
||||
#[macro_export]
|
||||
macro_rules! log_cv_out {
|
||||
($kind:expr, $cv:expr) => {
|
||||
$crate::util::logger::log_cv_internal("< ", &$cv, Some($kind))
|
||||
};
|
||||
($cv:expr) => {
|
||||
$crate::util::logger::log_cv_internal("< ", &$cv, None)
|
||||
};
|
||||
}
|
||||
4
src/util/mod.rs
Normal file
4
src/util/mod.rs
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
pub mod crypto_helper;
|
||||
pub mod crypto_util;
|
||||
pub mod file_util;
|
||||
pub mod logger;
|
||||
Loading…
Reference in a new issue