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

Inital Commit (Errors, I need to move through devices)

Browse source 
Signed-off-by: Alex Emmet <111742636+Alex-Emmet@users.noreply.github.com>
This commit is contained in:
Alex Emmet 2025-10-09 19:47:32 +00:00
commit 044eeb67a6
23 changed files with 5687 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

178
src/auth/auth_connector.rs Normal file
View file

@ -0,0 +1,178 @@
use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes};
use crate::util::config_util::CONFIG;
use hex;
use json::JsonValue;
use reqwest::header::CONTENT_TYPE;
use reqwest::{Client, Response};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::time::Duration;
use uuid::Uuid;
#[derive(Debug, Clone)]
pub struct AuthUser {
pub created_at: i64,
pub username: String,
pub display: String,
pub avatar: String,
pub about: String,
pub status: String,
pub public_key: String,
pub sub_level: i32,
pub sub_end: i32,
}
fn client() -> Client {
Client::builder()
.connect_timeout(Duration::from_secs(100))
.timeout(Duration::from_secs(150))
.build()
.unwrap()
}
pub async fn get_user(user_id: Uuid) -> Option<AuthUser> {
let url = format!("https://auth.tensamin.methanium.net/api/get/{}/", user_id);
let client = client();
let res = client.get(&url).send().await.ok()?;
let json = res.text().await.ok()?;
let mut cv = CommunicationValue::from_json(&json);
if cv.comm_type != CommunicationType::success {
return None;
}
Some(AuthUser {
created_at: cv
.get_data(DataTypes::created_at)
.unwrap()
.to_string()
.parse::<i64>()
.unwrap_or(-1),
username: cv.get_data(DataTypes::username).unwrap().to_string(),
display: cv.get_data(DataTypes::display).unwrap().to_string(),
avatar: cv.get_data(DataTypes::avatar).unwrap().to_string(),
about: cv.get_data(DataTypes::about).unwrap().to_string(),
status: cv.get_data(DataTypes::status).unwrap().to_string(),
public_key: cv.get_data(DataTypes::public_key).unwrap().to_string(),
sub_level: cv
.get_data(DataTypes::sub_level)
.unwrap()
.to_string()
.parse::<i32>()
.unwrap_or(-1),
sub_end: cv
.get_data(DataTypes::sub_end)
.unwrap()
.to_string()
.parse::<i32>()
.unwrap_or(-1),
})
}
pub async fn get_iota_id(user_id: Uuid) -> Option<Uuid> {
let url = format!(
"https://auth.tensamin.methanium.net/api/get/iota-id/{}",
user_id
);
let client = client();
let res = client
.get(&url)
.header(
"Authorization",
CONFIG.read().unwrap().omikron_id.to_string(),
)
.header("Content-Type", "application/json")
.send()
.await
.ok()?;
let json = res.text().await.ok()?;
let json = json.replace("iota_uuid", "iota_id");
let mut cv = CommunicationValue::from_json(&json);
if cv.comm_type != CommunicationType::success {
return None;
}
let iota_id_str = cv.get_data(DataTypes::iota_id)?.to_string();
Uuid::parse_str(&iota_id_str).ok()
}
pub async fn is_private_key_valid(user_id: Uuid, pk_hash: &str) -> bool {
let url = format!(
"https://auth.tensamin.methanium.net/api/get/private-key-hash/{}",
user_id
);
let client = client();
let res = client
.get(&url)
.header(
"Authorization",
CONFIG.read().unwrap().omikron_id.to_string(),
)
.header("PrivateKeyHash", pk_hash)
.header("Accept", "application/json")
.send()
.await;
let Ok(response) = res else {
return false;
};
let Ok(body) = response.text().await else {
return false;
};
let mut cv = CommunicationValue::from_json(&body);
if cv.comm_type != CommunicationType::success {
return false;
}
match cv.get_data(DataTypes::matches) {
Some(val) => val.as_bool().unwrap_or(false),
None => false,
}
}
pub async fn get_public_key(user_id: Uuid) -> Option<String> {
let url = format!(
"https://auth.tensamin.methanium.net/api/{}/public-key/",
user_id
);
let client = client();
let res = client
.get(&url)
.header("Accept", "application/json")
.send()
.await
.ok()?;
let body = res.text().await.ok()?;
let mut cv = CommunicationValue::from_json(&body);
if cv.comm_type != CommunicationType::message_send {
return None;
}
Some(cv.get_data(DataTypes::ping_clients)?.to_string())
}
pub async fn get_register() -> Option<Uuid> {
let url = "https://auth.tensamin.methanium.net/api/register/init/".to_string();
let client = client();
let res = client.get(&url).send().await.ok()?;
let json = res.text().await.ok()?;
let mut cv = CommunicationValue::from_json(&json);
Uuid::parse_str(&*cv.get_data(DataTypes::user_id).unwrap().to_string()).ok()
}
async fn handle_response(resp: Response) -> bool {
match resp.text().await {
Ok(text) => {
let cv = CommunicationValue::from_json(&text.to_string());
cv.comm_type == CommunicationType::success
}
Err(_) => false,
}
}

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

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

189
src/calls/call_connection.rs Normal file
View file

@ -0,0 +1,189 @@
// part of calls or in your ws-server file
use std::sync::Arc;
use futures::{SinkExt, StreamExt, channel::mpsc::UnboundedSender, lock::Mutex};
use json::JsonValue;
use tokio::sync::mpsc::unbounded_channel;
use tungstenite::{Message, Utf8Bytes};
use uuid::Uuid;
use crate::{
calls::{call_group::CallGroup, call_group::Caller, call_manager::CallManagerState},
data::communication::{CommunicationType, CommunicationValue, DataTypes},
};
pub async fn handle_connection(
raw_stream: tokio::net::TcpStream,
state: Arc<Mutex<CallManagerState>>,
) {
let ws_stream = tokio_tungstenite::accept_async(raw_stream)
.await
.expect("Error during the websocket handshake");
let (mut outgoing, mut incoming) = ws_stream.split();
// We create a channel so other parts can send to this client
pub type Tx = UnboundedSender<Utf8Bytes>;
let (tx, mut rx): (_, _) = unbounded_channel();
// Spawn a task to forward from rx → outgoing
let mut outgoing_clone = outgoing;
tokio::spawn(async move {
while let Some(msg) = rx.recv().await {
let _ = outgoing_clone
.send(tokio_tungstenite::tungstenite::Message::Text(msg))
.await;
}
});
// Each connection will track its user_id, call_group, etc.
let mut maybe_user_id: Option<Uuid> = None;
let mut maybe_call_id: Option<Uuid> = None;
while let Some(msg) = incoming.next().await {
let msg = match msg {
Ok(m) => m,
Err(_) => break,
};
if let tokio_tungstenite::tungstenite::Message::Text(s) = msg {
// parse CommunicationValue
if let mut cv = CommunicationValue::from_json(&s) {
match cv.comm_type {
CommunicationType::identification => {
// extract fields
let user_str = cv.get_data(DataTypes::user_id).unwrap().as_str().unwrap();
let call_str = cv.get_data(DataTypes::call_id).unwrap().as_str().unwrap();
let _secret_sha = cv
.get_data(DataTypes::call_secret_sha)
.unwrap()
.as_str()
.unwrap();
let user_id = Uuid::parse_str(user_str).ok();
let call_id = Uuid::parse_str(call_str).ok();
if let (Some(uid), Some(cid)) = (user_id, call_id) {
maybe_user_id = Some(uid);
maybe_call_id = Some(cid);
// Add to call group
let mut st = state.lock().await;
let group = state.lock().await.get_or_create_group(cid, &_secret_sha);
group.clone().add_member(uid, tx.clone());
// Response: identification_response plus states
let mut response =
CommunicationValue::new(CommunicationType::identification_response)
.with_id(cv.get_id().clone());
// build user states
let mut users = JsonValue::new_object();
for c in group.callers.values() {
let mut user_info = JsonValue::new_object();
user_info.insert("state", JsonValue::from("muted".to_string()));
user_info.insert("streaming", JsonValue::from(false));
users.insert(&c.user_id.to_string(), JsonValue::from(user_info));
}
response = response.add_data(DataTypes::about, JsonValue::from(users));
// send back
let _ = tx.send(Utf8Bytes::from(response.to_json().to_string()));
}
}
CommunicationType::ping => {
// optionally parse LAST_PING
// reply with PONG with same message_id
let mut resp = CommunicationValue::new(CommunicationType::pong)
.with_id(cv.get_id().clone());
let _ = tx.send(Utf8Bytes::from(resp.to_json().to_string()));
}
CommunicationType::client_changed => {
match (maybe_user_id, cv.get_data(DataTypes::call_state)) {
(Some(uid), Some(JsonValue::String(state_str))) => {
if let Some(cid) = maybe_call_id {
{
let mut st = state.lock().await;
if let Some(group) = st.call_groups.get_mut(&cid) {
if let Some(caller) = group.caller_state_mut(&uid) {
caller_state_change(caller, &state_str);
}
let mut bc = cv.clone();
bc = bc.add_data(
DataTypes::sender_id,
JsonValue::String(uid.to_string()),
);
group.broadcast(&bc.to_json().to_string());
}
}
}
}
_ => (),
}
}
CommunicationType::start_stream | CommunicationType::end_stream => {
if let Some(uid) = maybe_user_id {
if let Some(cid) = maybe_call_id {
let mut st = state.lock().await;
if let Some(group) = st.call_groups.get_mut(&cid) {
if let Some(caller) = group.caller_state_mut(&uid) {
let streaming =
cv.comm_type == CommunicationType::start_stream;
// set streaming
// caller.streaming = streaming; // if you store streaming
}
let mut bc = cv.clone();
bc = bc.add_data(
DataTypes::sender_id,
JsonValue::String(uid.to_string()),
);
group.broadcast(&bc.to_json().to_string());
}
}
}
}
CommunicationType::webrtc_sdp
| CommunicationType::webrtc_ice
| CommunicationType::watch_stream => {
if let Some(uid) = maybe_user_id {
if let Some(JsonValue::String(receiver_str)) =
cv.get_data(DataTypes::receiver_id)
{
if let Ok(receiver_id) = Uuid::parse_str(&receiver_str) {
if let Some(cid) = maybe_call_id {
let mut st = state.lock().await;
if let Some(group) = st.call_groups.get(&cid) {
let mut bc = cv.clone();
bc = bc.add_data(
DataTypes::sender_id,
JsonValue::String(uid.to_string()),
);
group.send_to(&receiver_id, &bc.to_json().to_string());
}
}
}
}
}
}
_ => {
// other types you may handle
}
}
}
}
}
// on close / disconnection:
if let (Some(uid), Some(cid)) = (maybe_user_id, maybe_call_id) {
let mut st = state.lock().await;
if let Some(group) = st.call_groups.get_mut(&cid) {
group.remove_member(&uid);
}
st.remove_inactive();
}
}
fn caller_state_change(_caller: &mut Caller, _state_str: &str) {
// parse and set your enum, e.g. match _state_str { "active" => ..., etc. }
}

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

@ -0,0 +1,106 @@
use std::{
collections::{HashMap, HashSet},
sync::Arc,
};
use tokio::sync::mpsc::UnboundedSender;
use json::JsonValue;
use tungstenite::Utf8Bytes;
use uuid::Uuid;
use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes};
pub type Tx = UnboundedSender<Utf8Bytes>;
#[derive(Debug)]
pub struct CallGroup {
pub call_id: Uuid,
secret: String,
pub callers: HashMap<Uuid, Caller>,
pub inform_on_end: HashSet<Uuid>,
pub started_at: u64,
pub ended_at: Option<u64>,
}
impl CallGroup {
pub fn new(call_id: Uuid, secret: String) -> Self {
CallGroup {
call_id,
secret,
callers: HashMap::new(),
inform_on_end: HashSet::new(),
started_at: chrono::Utc::now().timestamp_millis() as u64,
ended_at: None,
}
}
pub fn add_member(mut self, user_id: Uuid, tx: Tx) {
// If already present, drop old
if let Some(old) = self.callers.remove(&user_id) {
// optionally notify or close old
// no direct close since we don’t hold session here
}
// Notify existing about new client
let mut cv = CommunicationValue::new(CommunicationType::client_connected);
cv = cv
.add_data(DataTypes::user_id, JsonValue::String(user_id.to_string()))
.add_data(
DataTypes::call_state,
JsonValue::String("muted".to_string()),
);
// broadcast
let msg = cv.to_json().to_string();
for c in self.callers.values() {
let _ = c.tx.send(Utf8Bytes::from(msg.clone()));
}
self.callers.insert(user_id, Caller { user_id, tx });
}
pub fn remove_member(&mut self, user_id: &Uuid) {
if self.callers.remove(user_id).is_some() {
let cv = CommunicationValue::new(CommunicationType::client_closed)
.add_data(DataTypes::user_id, JsonValue::String(user_id.to_string()))
.to_json()
.to_string();
for c in self.callers.values() {
let _ = c.tx.send(Utf8Bytes::from(cv.clone()));
}
}
if self.callers.is_empty() {
self.ended_at = Some(chrono::Utc::now().timestamp_millis() as u64);
// Optionally notify inform_on_end
let end_msg = CommunicationValue::new(CommunicationType::end_call)
.add_data(
DataTypes::call_id,
JsonValue::String(self.call_id.to_string()),
)
.to_json()
.to_string();
for &uid in &self.inform_on_end {
// here you’d send via your Rho / other channel to user
// e.g. RhoManager::message_to(uid, end_msg.clone());
}
}
}
pub fn broadcast(&self, msg: &str) {
for c in self.callers.values() {
let _ = c.tx.send(Utf8Bytes::from(msg.clone()));
}
}
pub fn send_to(&self, target: &Uuid, msg: &str) {
if let Some(c) = self.callers.get(target) {
let _ = c.tx.send(Utf8Bytes::from(msg.clone()));
}
}
pub fn caller_state_mut(&mut self, uid: &Uuid) -> Option<&mut Caller> {
self.callers.get_mut(uid)
}
}
#[derive(Debug)]
pub struct Caller {
pub user_id: Uuid,
pub tx: Tx,
}

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

@ -0,0 +1,40 @@
use crate::calls::call_group::CallGroup;
use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes};
use futures::{SinkExt, StreamExt};
use json::JsonValue;
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex};
use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel};
use tungstenite::Utf8Bytes;
use uuid::Uuid;
use warp::Filter;
pub type Tx = UnboundedSender<Utf8Bytes>;
#[derive(Debug)]
pub struct CallManagerState {
pub call_groups: HashMap<Uuid, Arc<CallGroup>>,
}
impl CallManagerState {
pub fn new() -> Self {
CallManagerState {
call_groups: HashMap::new(),
}
}
pub fn get_or_create_group(&mut self, call_id: Uuid, secret: &str) -> Arc<CallGroup> {
let g = self.call_groups.get_mut(&call_id);
if let Some(group) = g {
group.clone()
} else {
let cg = Arc::new(CallGroup::new(call_id, secret.to_string()));
self.call_groups.insert(call_id, cg);
self.call_groups.get(&call_id).unwrap().clone()
}
}
pub fn remove_inactive(&mut self) {
self.call_groups.retain(|_k, v| v.callers.len() > 0);
}
}

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

@ -0,0 +1,3 @@
pub mod call_connection;
pub mod call_group;
pub mod call_manager;

402
src/data/communication.rs Normal file
View file

@ -0,0 +1,402 @@
use axum::Json;
use json::number::Number;
use json::{Array, JsonValue, object, parse, stringify};
use std::any::Any;
use std::collections::HashMap;
use std::env::VarsOs;
use std::str::FromStr;
use std::time::{SystemTime, UNIX_EPOCH};
use uuid::Uuid;
#[derive(Eq, Hash, PartialEq, Clone, Debug)]
pub enum DataTypes {
error_type,
uuid,
chat_partner_id,
iota_id,
user_id,
user_ids,
user_state,
user_states,
user_pings,
call_state,
screen_share,
private_key_hash,
accepted,
accepted_profiles,
denied_profiles,
content,
messages,
send_time,
get_time,
get_variant,
shared_secret_own,
shared_secret_other,
shared_secret_sign,
shared_secret,
call_id,
call_name,
call_secret_sha,
call_secret,
shared_call_secret,
start_date,
end_date,
receiver_id,
sender_id,
signature,
signed,
message,
last_ping,
ping_iota,
ping_clients,
matches,
omikron,
offset,
amount,
position,
name,
path,
codec,
function,
payload,
result,
interactables,
want_to_watch,
watcher,
created_at,
username,
display,
avatar,
about,
status,
public_key,
sub_level,
sub_end,
community_address,
challenge,
community_title,
communities,
}
impl DataTypes {
pub fn parse(p0: String) -> DataTypes {
// normalize: lowercase + remove underscores
let normalized = p0.to_lowercase().replace('_', "");
match normalized.as_str() {
"errortype" => DataTypes::error_type,
"chatpartnerid" => DataTypes::chat_partner_id,
"iotaid" => DataTypes::iota_id,
"userid" => DataTypes::user_id,
"userids" => DataTypes::user_ids,
"userstate" => DataTypes::user_state,
"userstates" => DataTypes::user_states,
"userpings" => DataTypes::user_pings,
"callstate" => DataTypes::call_state,
"screenshare" => DataTypes::screen_share,
"privatekeyhash" => DataTypes::private_key_hash,
"accepted" => DataTypes::accepted,
"acceptedprofiles" => DataTypes::accepted_profiles,
"deniedprofiles" => DataTypes::denied_profiles,
"content" => DataTypes::content,
"messages" => DataTypes::messages,
"sendtime" => DataTypes::send_time,
"gettime" => DataTypes::get_time,
"getvariant" => DataTypes::get_variant,
"sharedsecretown" => DataTypes::shared_secret_own,
"sharedsecretother" => DataTypes::shared_secret_other,
"sharedsecretsign" => DataTypes::shared_secret_sign,
"sharedsecret" => DataTypes::shared_secret,
"callid" => DataTypes::call_id,
"callname" => DataTypes::call_name,
"callsecretsha" => DataTypes::call_secret_sha,
"callsecret" => DataTypes::call_secret,
"sharedcallsecret" => DataTypes::shared_call_secret,
"startdate" => DataTypes::start_date,
"enddate" => DataTypes::end_date,
"receiverid" => DataTypes::receiver_id,
"senderid" => DataTypes::sender_id,
"signature" => DataTypes::signature,
"signed" => DataTypes::signed,
"message" => DataTypes::message,
"lastping" => DataTypes::last_ping,
"pingiota" => DataTypes::ping_iota,
"pingclients" => DataTypes::ping_clients,
"matches" => DataTypes::matches,
"omikron" => DataTypes::omikron,
"offset" => DataTypes::offset,
"amount" => DataTypes::amount,
"position" => DataTypes::position,
"name" => DataTypes::name,
"path" => DataTypes::path,
"codec" => DataTypes::codec,
"function" => DataTypes::function,
"payload" => DataTypes::payload,
"result" => DataTypes::result,
"interactables" => DataTypes::interactables,
"wanttowatch" => DataTypes::want_to_watch,
"watcher" => DataTypes::watcher,
"createdat" => DataTypes::created_at,
"username" => DataTypes::username,
"display" => DataTypes::display,
"avatar" => DataTypes::avatar,
"about" => DataTypes::about,
"status" => DataTypes::status,
"publickey" => DataTypes::public_key,
"sublevel" => DataTypes::sub_level,
"subend" => DataTypes::sub_end,
"communityaddress" => DataTypes::community_address,
"challenge" => DataTypes::challenge,
"communitytitle" => DataTypes::community_title,
"communities" => DataTypes::communities,
_ => DataTypes::error_type, // fallback if unknown
}
}
}
#[derive(PartialEq, Clone, Debug)]
pub enum CommunicationType {
error,
success,
message,
message_send,
message_live,
message_other_iota,
message_chunk,
messages_get,
change_confirm,
confirm_receive,
confirm_read,
get_chats,
get_states,
add_community,
remove_community,
get_communities,
challenge,
challenge_response,
register,
register_response,
identification,
identification_response,
ping,
pong,
add_chat,
send_chat,
iota_connected,
iota_closed,
client_changed,
client_connected,
client_closed,
public_key,
private_key,
webrtc_sdp,
webrtc_ice,
start_stream,
end_stream,
watch_stream,
get_call,
new_call,
call_invite,
end_call,
function,
update,
}
impl CommunicationType {
pub fn parse(p0: String) -> CommunicationType {
let normalized = p0.to_lowercase().replace('_', "");
match normalized.as_str() {
"error" => CommunicationType::error,
"success" => CommunicationType::success,
"message" => CommunicationType::message,
"messagelive" => CommunicationType::message_live,
"messageotheriota" => CommunicationType::message_other_iota,
"messagechunk" => CommunicationType::message_chunk,
"messagesget" => CommunicationType::messages_get,
"message_send" => CommunicationType::message_send,
"changeconfirm" => CommunicationType::change_confirm,
"confirmreceive" => CommunicationType::confirm_receive,
"confirmread" => CommunicationType::confirm_read,
"getchats" => CommunicationType::get_chats,
"getstates" => CommunicationType::get_states,
"addcommunity" => CommunicationType::add_community,
"removecommunity" => CommunicationType::remove_community,
"getcommunities" => CommunicationType::get_communities,
"challenge" => CommunicationType::challenge,
"challengeresponse" => CommunicationType::challenge_response,
"register" => CommunicationType::register,
"registerresponse" => CommunicationType::register_response,
"identification" => CommunicationType::identification,
"identificationresponse" => CommunicationType::identification_response,
"ping" => CommunicationType::ping,
"pong" => CommunicationType::pong,
"addchat" => CommunicationType::add_chat,
"sendchat" => CommunicationType::send_chat,
"iotaconnected" => CommunicationType::iota_connected,
"iotaclosed" => CommunicationType::iota_closed,
"clientchanged" => CommunicationType::client_changed,
"clientconnected" => CommunicationType::client_connected,
"clientclosed" => CommunicationType::client_closed,
"publickey" => CommunicationType::public_key,
"privatekey" => CommunicationType::private_key,
"webrtcsdp" => CommunicationType::webrtc_sdp,
"webrtcice" => CommunicationType::webrtc_ice,
"startstream" => CommunicationType::start_stream,
"endstream" => CommunicationType::end_stream,
"watchstream" => CommunicationType::watch_stream,
"getcall" => CommunicationType::get_call,
"newcall" => CommunicationType::new_call,
"callinvite" => CommunicationType::call_invite,
"endcall" => CommunicationType::end_call,
"function" => CommunicationType::function,
"update" => CommunicationType::update,
_ => CommunicationType::error,
}
}
}
#[derive(Debug, Clone)]
pub struct CommunicationValue {
pub id: Uuid,
pub comm_type: CommunicationType,
pub sender: Uuid,
pub receiver: Uuid,
pub data: HashMap<DataTypes, JsonValue>,
}
impl CommunicationValue {
pub fn new(comm_type: CommunicationType) -> Self {
Self {
id: Uuid::new_v4(),
comm_type,
sender: Uuid::new_v4(),
receiver: Uuid::new_v4(),
data: HashMap::new(),
}
}
pub fn with_id(mut self, p0: Uuid) -> Self {
self.id = p0;
self
}
pub fn get_id(&self) -> Uuid {
self.id.clone()
}
pub fn with_sender(mut self, sender: Uuid) -> Self {
self.sender = sender;
self
}
pub fn get_sender(&self) -> Uuid {
self.sender.clone()
}
pub fn with_receiver(mut self, receiver: Uuid) -> Self {
self.receiver = receiver;
self
}
pub fn get_receiver(&self) -> Uuid {
self.receiver.clone()
}
pub fn add_data_num(mut self, key: DataTypes, value: Number) -> Self {
self.data.insert(key, JsonValue::Number(value));
self
}
pub fn add_data_str(mut self, key: DataTypes, value: String) -> Self {
self.data.insert(key, JsonValue::String(value));
self
}
pub fn add_data(mut self, key: DataTypes, value: JsonValue) -> Self {
self.data.insert(key, value);
self
}
pub fn add_array(mut self, key: DataTypes, value: Array) -> Self {
self.data.insert(key, JsonValue::Array(value));
self
}
pub fn get_data(&self, key: DataTypes) -> Option<&JsonValue> {
self.data.get(&key)
}
pub(crate) fn is_type(&self, p0: CommunicationType) -> bool {
self.comm_type == p0
}
pub fn to_json(&self) -> JsonValue {
let mut jdata = object! {};
for (k, v) in &self.data {
jdata[&format!("{:?}", k)] = JsonValue::from(v.clone());
}
object! {
id: self.id.to_string(),
type: format!("{:?}", self.comm_type),
sender: self.sender.to_string(),
receiver: self.receiver.to_string(),
data: jdata
}
}
pub fn from_json(json_str: &str) -> Self {
let parsed = parse(json_str).unwrap();
let comm_type = CommunicationType::parse(parsed["type"].to_string());
let sender: Uuid = parsed["sender"]
.as_str()
.and_then(|s| Uuid::parse_str(s).ok())
.unwrap_or(Uuid::new_v4());
let receiver: Uuid = parsed["receiver"]
.as_str()
.and_then(|s| Uuid::parse_str(s).ok())
.unwrap_or(Uuid::new_v4());
let uuid = Uuid::parse_str(parsed["id"].as_str().unwrap_or("")).unwrap_or(Uuid::new_v4());
let mut data = HashMap::new();
if parsed["data"].is_object() {
for (k, v) in parsed["data"].entries() {
data.insert(DataTypes::parse(k.to_string()), v.clone());
}
}
Self {
id: uuid,
comm_type,
sender,
receiver,
data,
}
}
pub fn ack_message(message_id: Uuid, sender: Uuid) -> CommunicationValue {
let mut cv = CommunicationValue::new(CommunicationType::message).with_id(message_id);
if let s = sender {
cv = cv.add_data(DataTypes::send_time, JsonValue::String(s.to_string()));
}
cv
}
pub fn forward_to_other_iota(original: &mut CommunicationValue) -> CommunicationValue {
let receiver = Uuid::from_str(
&*original
.get_data(DataTypes::receiver_id)
.unwrap()
.to_string(),
)
.ok()
.or(Option::from(Uuid::nil()));
let now_ms = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_millis() as i64;
let cv = CommunicationValue::new(CommunicationType::message_other_iota)
.with_id(original.get_id())
.with_receiver(receiver.unwrap())
.add_data(DataTypes::send_time, JsonValue::String(now_ms.to_string()))
.add_data(
DataTypes::content,
JsonValue::String(original.get_data(DataTypes::content).unwrap().to_string()),
);
// include sender_id if the original had one
let sender = original.get_sender();
cv.add_data(DataTypes::sender_id, JsonValue::String(sender.to_string()))
}
}

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

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

49
src/data/user.rs Normal file
View file

@ -0,0 +1,49 @@
use uuid::Uuid;
#[derive(Clone)]
pub struct User {
pub iota_id: Uuid,
pub user_id: Uuid,
pub status: UserStatus,
}
impl User {
pub fn new(iota_id: Uuid, user_id: Uuid, status: UserStatus) -> Self {
User {
iota_id,
user_id,
status,
}
}
}
#[derive(Clone, Copy)]
pub enum UserStatus {
online,
do_not_disturb,
wc,
away,
user_offline,
iota_offline,
}
impl UserStatus {
pub fn to_string(self) -> String {
match self {
UserStatus::online => "online".to_string(),
UserStatus::do_not_disturb => "do_not_disturb".to_string(),
UserStatus::wc => "wc".to_string(),
UserStatus::away => "away".to_string(),
UserStatus::user_offline => "user_offline".to_string(),
UserStatus::iota_offline => "iota_offline".to_string(),
}
}
pub fn from_string(status_str: &str) -> Result<Self, String> {
match status_str.to_lowercase().as_str() {
"online" => Ok(UserStatus::online),
"do_not_disturb" => Ok(UserStatus::do_not_disturb),
"wc" => Ok(UserStatus::wc),
"away" => Ok(UserStatus::away),
"user_offline" => Ok(UserStatus::user_offline),
"iota_offline" => Ok(UserStatus::iota_offline),
_ => Err(format!("Invalid user status: {}", status_str)),
}
}
}

23
src/main.rs Normal file
View file

@ -0,0 +1,23 @@
mod auth;
mod calls;
mod data;
mod omega;
mod rho;
mod util;
use crate::rho::client_connection;
use crate::rho::iota_connection;
use crate::util::config_util::Config;
use crate::util::file_util;
use tokio_websocket_server::socket::{WebSocketMessage, WebsocketServer};
use tracing::info;
use tracing_subscriber::fmt;
#[tokio::main]
async fn main() {
fmt::init();
loop {
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
}
}

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

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

215
src/omega/omega_connection.rs Normal file
View file

@ -0,0 +1,215 @@
use std::sync::Arc;
use std::time::Duration;
use std::{collections::HashMap, net::TcpStream};
use crate::{
data::{
communication::DataTypes,
user::{User, UserStatus},
},
rho::{self, rho_manager},
util::config_util::CONFIG,
};
use axum::Json;
use dashmap::DashMap;
use futures::{SinkExt, StreamExt};
use json::JsonValue;
use once_cell::sync::Lazy;
use tokio::sync::Mutex;
use tokio::time::sleep;
use tokio_tungstenite::{
MaybeTlsStream, WebSocketStream, connect_async, tungstenite::protocol::Message,
};
use tungstenite::Utf8Bytes;
use uuid::Uuid;
use crate::data::communication::{CommunicationType, CommunicationValue};
static WAITING_TASKS: Lazy<DashMap<Uuid, Box<dyn Fn(CommunicationValue) -> bool + Send + Sync>>> =
Lazy::new(DashMap::new);
pub struct OmegaConnection {
ws_stream: Arc<Mutex<Option<tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>>>>,
}
impl OmegaConnection {
pub fn new() -> Self {
OmegaConnection {
ws_stream: Arc::new(Mutex::new(None)),
}
}
pub async fn connect(&self) {
self.connect_internal(0).await;
}
async fn connect_internal(&self, mut retry: usize) {
let uri = "wss://tensamin.methanium.net/ws/omega/";
loop {
if retry > 20 {
eprintln!("Max retry attempts reached, giving up.");
return;
}
match connect_async(uri).await {
Ok((ws_stream, _)) => {
let mut guard = Some(ws_stream);
// Send IDENTIFICATION
let identify_msg = CommunicationValue::new(CommunicationType::identification)
.add_data(
DataTypes::uuid,
JsonValue::String(CONFIG.read().unwrap().omikron_id.to_string()),
);
self.send_message(&identify_msg).await;
// Spawn reader loop
let ws_stream_clone = self.ws_stream.clone();
tokio::spawn(async move {
OmegaConnection::read_loop(ws_stream_clone).await;
});
break; // success, exit the loop
}
Err(e) => {
eprintln!("WebSocket connection failed (attempt {}): {}", retry, e);
retry += 1;
sleep(Duration::from_secs(2)).await;
continue;
}
}
}
}
async fn read_loop(
ws_stream: Arc<Mutex<Option<tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>>>>,
) {
loop {
let mut lock = ws_stream.lock().await;
let Some(ws) = lock.as_mut() else {
break;
};
match ws.next().await {
Some(Ok(Message::Text(msg))) => {
let mut cv = CommunicationValue::from_json(&msg);
let msg_id = cv.get_id();
// Handle waiting tasks
if let Some(task) = WAITING_TASKS.remove(&msg_id) {
if (task.1)(cv.clone()) {
continue;
}
}
// Handle CLIENT_CHANGED
if cv.is_type(CommunicationType::client_changed) {
let iota_id = Uuid::parse_str(
cv.get_data(DataTypes::iota_id).unwrap().as_str().unwrap(),
)
.unwrap();
let user_id = Uuid::parse_str(
cv.get_data(DataTypes::user_id).unwrap().as_str().unwrap(),
)
.unwrap();
let status_str = cv
.get_data(DataTypes::user_state)
.unwrap()
.as_str()
.unwrap();
let status = UserStatus::from_string(&status_str)
.unwrap_or(UserStatus::iota_offline);
let user = User::new(iota_id, user_id, status);
for rho_con in rho_manager::get_all_connections().await {
rho_con.are_they_interested(&user);
}
}
}
Some(Ok(Message::Close(_))) | None => {
OmegaConnection::reconnect();
break;
}
Some(Err(_)) => {
OmegaConnection::reconnect();
break;
}
_ => {}
}
}
}
async fn reconnect() {
let conn = OmegaConnection::new();
conn.connect().await;
}
pub async fn send_message(&self, cv: &CommunicationValue) {
let mut guard = self.ws_stream.lock().await;
if let Some(ws) = guard.as_mut() {
let _ = ws
.send(Message::Text(Utf8Bytes::from(cv.to_json().to_string())))
.await;
}
}
pub async fn connect_iota(iota_id: Uuid, user_ids: Vec<Uuid>) {
let user_ids_str = user_ids
.iter()
.map(|id| id.to_string())
.collect::<Vec<_>>()
.join(",");
let cv = CommunicationValue::new(CommunicationType::iota_connected)
.add_data(DataTypes::iota_id, JsonValue::from(iota_id.to_string()))
.add_data(DataTypes::user_ids, JsonValue::from(user_ids_str));
OmegaConnection::send_global(cv).await;
}
pub async fn close_iota(iota_id: Uuid) {
let cv = CommunicationValue::new(CommunicationType::iota_closed)
.add_data(DataTypes::iota_id, JsonValue::from(iota_id.to_string()));
OmegaConnection::send_global(cv).await;
}
pub async fn client_changed(iota_id: Uuid, user_id: Uuid, state: UserStatus) {
let cv = CommunicationValue::new(CommunicationType::client_changed)
.add_data(DataTypes::iota_id, JsonValue::from(iota_id.to_string()))
.add_data(DataTypes::user_id, JsonValue::from(user_id.to_string()))
.add_data(DataTypes::user_state, JsonValue::from(state.to_string()));
OmegaConnection::send_global(cv).await;
}
pub async fn user_states(user_id: Uuid, user_ids: Vec<Uuid>) {
let user_ids_str = user_ids
.iter()
.map(|id| id.to_string())
.collect::<Vec<_>>()
.join(",");
let cv = CommunicationValue::new(CommunicationType::get_states)
.add_data(DataTypes::user_ids, JsonValue::from(user_ids_str));
let msg_id = cv.get_id();
WAITING_TASKS.insert(
msg_id,
Box::new(move |response: CommunicationValue| {
Box::pin(async move |response2: CommunicationValue| {
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
});
true
}),
);
OmegaConnection::send_global(cv).await;
}
async fn send_global(cv: CommunicationValue) {
let conn = OmegaConnection::new();
conn.send_message(&cv).await;
}
}

515
src/rho/client_connection.rs Normal file
View file

@ -0,0 +1,515 @@
use futures::SinkExt;
use http::header::AUTHORIZATION;
use std::{
collections::HashMap,
sync::{Arc, Weak},
};
use tokio::sync::{Mutex, RwLock};
use tokio_tungstenite::{WebSocketStream, tungstenite::Message};
use tungstenite::Utf8Bytes;
use uuid::Uuid;
use super::{rho_connection::RhoConnection, rho_manager};
use crate::{
auth::auth_connector,
// calls::call_manager::CallManager,
data::{
communication::{CommunicationType, CommunicationValue, DataTypes},
user::{User, UserStatus},
},
omega::omega_connection::OmegaConnection,
};
/// ClientConnection represents a WebSocket connection from a client device
pub struct ClientConnection {
/// WebSocket session
session: Arc<Mutex<WebSocketStream<tokio::net::TcpStream>>>,
/// User ID associated with this client
user_id: Arc<RwLock<Option<Uuid>>>,
/// Whether this connection has been identified/authenticated
identified: Arc<RwLock<bool>>,
/// Ping latency tracking
ping: Arc<RwLock<i64>>,
/// Weak reference to RhoConnection to avoid circular references
rho_connection: Arc<RwLock<Option<Weak<RhoConnection>>>>,
/// List of user IDs this client is interested in receiving updates about
interested_users: Arc<RwLock<Vec<Uuid>>>,
}
impl ClientConnection {
/// Create a new ClientConnection
pub fn new(session: WebSocketStream<tokio::net::TcpStream>) -> Arc<Self> {
Arc::new(Self {
session: Arc::new(Mutex::new(session)),
user_id: Arc::new(RwLock::new(None)),
identified: Arc::new(RwLock::new(false)),
ping: Arc::new(RwLock::new(-1)),
rho_connection: Arc::new(RwLock::new(None)),
interested_users: Arc::new(RwLock::new(Vec::new())),
})
}
/// Get the user ID
pub async fn get_user_id(&self) -> Option<Uuid> {
*self.user_id.read().await
}
/// Check if connection is identified
pub async fn is_identified(&self) -> bool {
*self.identified.read().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: Weak<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() {
weak_ref.upgrade()
} else {
None
}
}
/// Send a string message to the client
pub async fn send_message_str(&self, message: &str) {
let mut session = self.session.lock().await;
if let Err(e) = session
.send(Message::Text(Utf8Bytes::from(message.to_string())))
.await
{
eprintln!("Failed to send message to client: {}", e);
}
}
/// Send a CommunicationValue to the client
pub async fn send_message(&self, cv: &CommunicationValue) {
self.send_message_str(&cv.to_json().to_string()).await;
}
/// Handle incoming message from client
pub async fn handle_message(self: Arc<Self>, message: String) {
let cv = CommunicationValue::from_json(&message);
// Handle identification
if cv.is_type(CommunicationType::identification) && !self.is_identified().await {
self.handle_identification(Arc::clone(&self), cv).await;
return;
}
if !self.is_identified().await {
return;
}
// 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::get_call) {
self.handle_get_call(cv).await;
return;
}
// Forward other messages to Iota
self.forward_to_iota(cv).await;
}
/// Handle identification message
async fn handle_identification(&self, sarc: Arc<ClientConnection>, mut cv: CommunicationValue) {
// Extract user ID
let user_id = match cv.get_data(DataTypes::user_id) {
Some(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;
}
},
None => {
self.send_error_response(&cv.get_id(), CommunicationType::error)
.await;
return;
}
};
// Find RhoConnection for this user
let rho_connection = match rho_manager::get_rho_con_for_user(user_id).await {
Some(rho) => rho,
None => {
self.send_error_response(&cv.get_id(), CommunicationType::error)
.await;
return;
}
};
// Validate private key
if let Some(private_key_hash) = cv.get_data(DataTypes::private_key_hash) {
let is_valid =
auth_connector::is_private_key_valid(user_id, &private_key_hash.to_string()).await;
if !is_valid {
self.send_error_response(&cv.get_id(), CommunicationType::error)
.await;
return;
}
} else {
self.send_error_response(&cv.get_id(), CommunicationType::error)
.await;
return;
}
// Set identification data
{
let mut user_id_guard = self.user_id.write().await;
*user_id_guard = Some(user_id);
}
{
let mut identified_guard = self.identified.write().await;
*identified_guard = true;
}
// Set up RhoConnection reference
self.set_rho_connection(Arc::downgrade(&rho_connection))
.await;
// Add this client to the RhoConnection
rho_connection.add_client_connection(Arc::from(sarc)).await;
// Send success response
let response = CommunicationValue::new(CommunicationType::identification_response)
.with_id(cv.get_id());
self.send_message(&response).await;
}
/// Handle ping message
async fn handle_ping(&self, mut cv: CommunicationValue) {
// Update our ping if provided
if let Some(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;
}
}
// 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_str(DataTypes::ping_iota, iota_ping.to_string());
self.send_message(&response).await;
}
/// Handle client status change
async fn handle_client_changed(&self, mut cv: CommunicationValue) {
if let Some(user_id) = self.get_user_id().await {
if let Some(status_str) = cv.get_data(DataTypes::user_state) {
// Parse user status - this would need to be implemented properly
let user_status = UserStatus::online; // placeholder
if let Some(rho_conn) = self.get_rho_connection().await {
OmegaConnection::client_changed(
rho_conn.get_iota_id().await,
user_id,
user_status,
);
}
}
}
}
/// Handle call invite
async fn handle_call_invite(&self, mut cv: CommunicationValue) {
let receiver_id = match cv.get_data(DataTypes::receiver_id) {
Some(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;
}
},
None => {
self.send_error_response(&cv.get_id(), CommunicationType::error)
.await;
return;
}
};
let call_id = match cv.get_data(DataTypes::call_id) {
Some(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;
}
},
None => {
self.send_error_response(&cv.get_id(), CommunicationType::error)
.await;
return;
}
};
let call_secret_sha = cv
.get_data(DataTypes::call_secret_sha)
.map(|s| s.to_string());
let call_secret = cv.get_data(DataTypes::call_secret).map(|s| s.to_string());
if call_secret_sha.is_none() || call_secret.is_none() {
self.send_error_response(&cv.get_id(), CommunicationType::error)
.await;
return;
}
// Get call group - placeholder implementation
// let call_group = CallManager::get_call_group(call_id, &call_secret_sha.unwrap(), true).await;
// if call_group.is_none() {
// self.send_error_response(&cv.get_id(), CommunicationType::error)
// .await;
// return;
// }
// Find target RhoConnection
let target_rho = match rho_manager::get_rho_con_for_user(receiver_id).await {
Some(rho) => rho,
None => {
self.send_error_response(&cv.get_id(), CommunicationType::error)
.await;
return;
}
};
// Get sender user ID
let sender_id = match self.get_user_id().await {
Some(id) => id,
None => return,
};
// Create and send call distribution message
let distribute = CommunicationValue::new(CommunicationType::new_call)
.with_receiver(receiver_id)
.with_sender(sender_id)
.add_data_str(DataTypes::call_id, call_id.to_string())
.add_data_str(DataTypes::receiver_id, receiver_id.to_string())
.add_data_str(DataTypes::sender_id, sender_id.to_string())
.add_data_str(DataTypes::call_secret, call_secret.unwrap());
target_rho.message_iota_to_client(distribute).await;
// Handle call group invitation logic here
// This would require implementing CallGroup::Caller and related functionality
// Send success response
let response = CommunicationValue::new(CommunicationType::call_invite).with_id(cv.get_id());
self.send_message(&response).await;
}
/// Handle get call request
async fn handle_get_call(&self, mut cv: CommunicationValue) {
let user_id = match self.get_user_id().await {
Some(id) => id,
None => return,
};
let call_id = match cv.get_data(DataTypes::call_id) {
Some(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;
}
},
None => {
self.send_error_response(&cv.get_id(), CommunicationType::error)
.await;
return;
}
};
let call_secret_sha = cv
.get_data(DataTypes::call_secret_sha)
.map(|s| s.to_string());
if call_secret_sha.is_none() {
self.send_error_response(&cv.get_id(), CommunicationType::error)
.await;
return;
}
// Get call group
let mut response = CommunicationValue::new(CommunicationType::get_call)
.with_id(cv.get_id())
.with_receiver(user_id);
// Placeholder call group logic
// if let Some(call_group) = CallManager::get_call_group(call_id, &call_secret_sha.unwrap(), true).await {
// response = response
// .add_data_str(DataTypes::call_state, call_group.call_state.to_string())
// .add_data_str(DataTypes::start_date, call_group.started_at.to_string());
//
// if call_group.ended_at != 0 {
// response = response.add_data_str(DataTypes::end_date, call_group.ended_at.to_string());
// }
// } else {
response = response.add_data_str(DataTypes::call_state, "DESTROYED".to_string());
// }
self.send_message(&response).await;
}
/// Forward message to Iota
async fn forward_to_iota(&self, mut cv: CommunicationValue) {
if let Some(user_id) = self.get_user_id().await {
if let Some(rho_conn) = self.get_rho_connection().await {
let updated_cv = cv.with_sender(user_id);
rho_conn.message_to_iota(updated_cv).await;
}
}
}
/// Send error response
async fn send_error_response(&self, message_id: &Uuid, 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 session = self.session.lock().await;
let _ = session.close(None).await;
}
/// Set interested users list
pub async fn set_interested_users(&self, interested_ids: Vec<Uuid>) {
let mut interested_guard = self.interested_users.write().await;
*interested_guard = interested_ids;
}
/// Check if interested in a user and send notification
pub async fn are_you_interested(&self, user: &User) {
let interested_guard = self.interested_users.read().await;
if interested_guard.contains(&user.user_id) {
let notification = CommunicationValue::new(CommunicationType::client_changed)
.add_data_str(DataTypes::user_id, user.user_id.to_string())
.add_data_str(
DataTypes::user_state,
format!("{:?}", user.status.to_string()),
);
self.send_message(&notification).await;
}
}
/// Handle connection close
pub async fn handle_close(&self) {
if self.is_identified().await {
if let Some(user_id) = self.get_user_id().await {
if let Some(rho_conn) = rho_manager::get_rho_con_for_user(user_id).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 {
session: Arc::clone(&self.session),
user_id: Arc::clone(&self.user_id),
identified: Arc::clone(&self.identified),
ping: Arc::clone(&self.ping),
rho_connection: Arc::clone(&self.rho_connection),
interested_users: Arc::clone(&self.interested_users),
}
}
}
impl std::fmt::Debug for ClientConnection {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ClientConnection")
.field("user_id", &"[async]")
.field("identified", &"[async]")
.field("ping", &"[async]")
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::sync::Mutex;
#[tokio::test]
async fn test_client_connection_creation() {
// Mock session - in real implementation this would be a proper WebSocket stream
let mock_session = tokio_tungstenite::WebSocketStream::from_raw_socket(
tokio::net::TcpStream::connect("127.0.0.1:0")
.await
.unwrap_or_else(|_| {
// This is just for testing, create a dummy stream
panic!("Cannot create test stream")
}),
tokio_tungstenite::tungstenite::protocol::Role::Client,
None,
)
.await;
// This test would fail in practice due to the mock stream
// but shows the intended API
// let client_conn = ClientConnection::new(mock_session);
// assert!(!client_conn.is_identified().await);
// assert_eq!(client_conn.get_ping().await, -1);
// assert!(client_conn.get_user_id().await.is_none());
}
#[tokio::test]
async fn test_interested_users() {
// This would also need a proper mock setup
// but shows the intended functionality
// let client_conn = ClientConnection::new(mock_session);
// let user_ids = vec![Uuid::new_v4(), Uuid::new_v4()];
// client_conn.set_interested_users(user_ids.clone()).await;
//
// // Test would verify that the interested users are stored correctly
}
}

345
src/rho/iota_connection.rs Normal file
View file

@ -0,0 +1,345 @@
use futures::SinkExt;
use json::{JsonValue, number::Number};
use std::{
collections::HashMap,
sync::{Arc, Weak},
};
use tokio::sync::{Mutex, RwLock};
use tokio_tungstenite::{WebSocketStream, tungstenite::Message};
use tungstenite::Utf8Bytes;
use uuid::Uuid;
use super::{rho_connection::RhoConnection, rho_manager};
use crate::{
auth::auth_connector,
// calls::call_manager::CallManager,
data::{
communication::{CommunicationType, CommunicationValue, DataTypes},
user::User,
},
omega::omega_connection::OmegaConnection,
};
/// IotaConnection represents a WebSocket connection from an Iota device
pub struct IotaConnection {
session: Arc<Mutex<WebSocketStream<tokio::net::TcpStream>>>,
iota_id: Arc<RwLock<Uuid>>,
user_ids: Arc<RwLock<Vec<Uuid>>>,
identified: Arc<RwLock<bool>>,
ping: Arc<RwLock<i64>>,
rho_connection: Arc<RwLock<Option<Weak<RhoConnection>>>>,
}
impl IotaConnection {
/// Create a new IotaConnection
pub fn new(session: WebSocketStream<tokio::net::TcpStream>) -> Arc<Self> {
Arc::new(Self {
session: Arc::new(Mutex::new(session)),
iota_id: Arc::new(RwLock::new(Uuid::nil())),
user_ids: Arc::new(RwLock::new(Vec::new())),
identified: Arc::new(RwLock::new(false)),
ping: Arc::new(RwLock::new(0)),
rho_connection: Arc::new(RwLock::new(None)),
})
}
/// Create with known IDs (for testing)
pub fn new_with_ids(
iota_id: Uuid,
user_ids: Vec<Uuid>,
session: Arc<Mutex<WebSocketStream<tokio::net::TcpStream>>>,
) -> Arc<Self> {
Arc::new(Self {
session,
iota_id: Arc::new(RwLock::new(iota_id)),
user_ids: Arc::new(RwLock::new(user_ids)),
identified: Arc::new(RwLock::new(true)),
ping: Arc::new(RwLock::new(0)),
rho_connection: Arc::new(RwLock::new(None)),
})
}
/// Get the Iota ID
pub async fn get_iota_id(&self) -> Uuid {
*self.iota_id.read().await
}
/// Get the user IDs
pub async fn get_user_ids(&self) -> Vec<Uuid> {
self.user_ids.read().await.clone()
}
/// Check if connection is identified
pub async fn is_identified(&self) -> bool {
*self.identified.read().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: Weak<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() {
weak_ref.upgrade()
} else {
None
}
}
/// Send a message to the Iota
pub async fn send_message_str(&self, message: &str) {
let mut session = self.session.lock().await;
let _ = session
.send(Message::Text(Utf8Bytes::from(message.to_string())))
.await;
}
/// Send a CommunicationValue to the Iota
pub async fn send_message(&self, cv: CommunicationValue) {
self.send_message_str(&cv.to_json().to_string()).await;
}
/// Handle incoming message from Iota
pub async fn handle_message(self: Arc<Self>, message: String) {
let cv = CommunicationValue::from_json(&message);
// Handle identification
if cv.is_type(CommunicationType::identification) && !self.is_identified().await {
self.handle_identification(cv).await;
return;
}
if !self.is_identified().await {
return;
}
// Handle ping
if cv.is_type(CommunicationType::ping) {
self.handle_ping(cv).await;
return;
}
// Handle forwarding to other Iotas or clients
let receiver_id = cv.get_receiver();
if !self.get_user_ids().await.contains(&receiver_id)
|| cv.is_type(CommunicationType::message_other_iota)
|| cv.is_type(CommunicationType::send_chat)
{
self.handle_forward_message(cv).await;
return;
}
// Handle GET_CHATS
if cv.is_type(CommunicationType::get_chats) {
self.handle_get_chats(cv).await;
return;
}
// Forward to client
self.forward_to_client(cv).await;
}
/// Handle identification message
async fn handle_identification(self: Arc<Self>, mut cv: CommunicationValue) {
// Parse Iota ID
let iota_id = match cv.get_data(DataTypes::iota_id) {
Some(id_str) => match Uuid::parse_str(&id_str.to_string()) {
Ok(id) => id,
Err(_) => {
self.send_error_response(&cv.get_id()).await;
return;
}
},
None => {
self.send_error_response(&cv.get_id()).await;
return;
}
};
// Parse user IDs
let mut validated_user_ids = Vec::new();
if let Some(user_ids_str) = cv.get_data(DataTypes::user_ids) {
for id_str in user_ids_str.to_string().split(',') {
if id_str.is_empty() {
continue;
}
if let Ok(user_id) = Uuid::parse_str(id_str.trim()) {
let auth_iota_id = auth_connector::get_iota_id(user_id).await.unwrap();
if auth_iota_id == iota_id {
validated_user_ids.push(user_id);
}
}
}
}
// Set identification data
{
let mut iota_id_guard = self.iota_id.write().await;
*iota_id_guard = iota_id;
}
{
let mut user_ids_guard = self.user_ids.write().await;
*user_ids_guard = validated_user_ids.clone();
}
{
let mut identified_guard = self.identified.write().await;
*identified_guard = true;
}
// Check for existing connection and close it
if rho_manager::contains_iota(iota_id).await {
if let Some(existing_rho) = rho_manager::get_rho_by_iota(iota_id).await {
existing_rho.close_iota_connection().await;
}
}
// Create RhoConnection
let rho_connection =
Arc::new(RhoConnection::new(self.clone(), validated_user_ids.clone()).await);
// Set up bidirectional reference
self.set_rho_connection(Arc::downgrade(&rho_connection))
.await;
// Add to manager
rho_manager::add_rho(rho_connection).await;
// Send response
let response = CommunicationValue::new(CommunicationType::identification_response)
.with_id(cv.get_id())
.add_data_str(DataTypes::accepted, validated_user_ids.len().to_string());
self.send_message(response).await;
}
/// Handle ping message
async fn handle_ping(&self, mut cv: CommunicationValue) {
// Update our ping if provided
if let Some(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;
}
}
// Collect client pings
let client_pings = if let Some(rho_conn) = self.get_rho_connection().await {
rho_conn.get_client_pings().await
} else {
HashMap::new()
};
// Send pong response
let pings = client_pings
.into_iter()
.map(|(k, v)| (k, JsonValue::String(v.to_string())))
.collect();
let response = CommunicationValue::new(CommunicationType::pong)
.with_id(cv.get_id())
.add_data(DataTypes::ping_clients, JsonValue::Object(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();
// Validate sender if present
let sender_id = cv.get_sender();
if !self.get_user_ids().await.contains(&sender_id) {
self.send_error_response(&cv.get_id()).await;
return;
}
// Find target RhoConnection
if let Some(target_rho) = rho_manager::get_rho_con_for_user(receiver_id).await {
target_rho.message_to_iota(cv).await;
} else {
// Send error if target not found
let error = CommunicationValue::new(CommunicationType::error)
.with_id(cv.get_id())
.with_sender(cv.get_sender());
self.send_message(error).await;
}
}
/// Handle GET_CHATS message
async fn handle_get_chats(&self, mut cv: CommunicationValue) {
let receiver_id = cv.get_receiver();
let mut interested_ids: Vec<Uuid> = Vec::new();
// Process contacts and add call information
if let Some(contacts_data) = cv.get_data(DataTypes::user_ids) {
// Parse contacts JSON array and enrich with call data
// This would need proper JSON parsing implementation
// For now, placeholder logic:
// Extract user IDs from contacts
// let contacts: Vec<ContactInfo> = parse_contacts(contacts_data);
// for contact in &contacts {
// interested_ids.push(contact.user_id);
// }
// Get call invites for receiver
// let invites = CallManager::get_call_invites(receiver_id).await;
// Enrich contacts with call information
// let enriched_contacts = enrich_with_calls(contacts, invites);
// cv = cv.add_data(DataTypes::user_ids, enriched_contacts.into());
}
// Notify OmegaConnection about user states
OmegaConnection::user_states(receiver_id, interested_ids.clone());
// Set interested users in RhoConnection
if let Some(rho_conn) = self.get_rho_connection().await {
rho_conn.set_interested(receiver_id, interested_ids).await;
}
// Forward to client
self.forward_to_client(cv).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);
let receiver_id = updated_cv.get_receiver();
rho_conn.message_iota_to_client(updated_cv).await;
}
}
/// Send error response
async fn send_error_response(&self, message_id: &Uuid) {
let error = CommunicationValue::new(CommunicationType::error).with_id(*message_id);
self.send_message(error).await;
}
/// Handle connection close
pub async fn handle_close(&self) {
if self.is_identified().await {
if let Some(rho_conn) = self.get_rho_connection().await {
rho_conn.close_iota_connection().await;
}
}
}
}
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()
}
}

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

@ -0,0 +1,21 @@
//! Rho module - Connection management between Iota and Client connections
//!
//! This module implements the Rho connection management system that maps
//! single Iota connections to multiple Client connections, providing
//! bidirectional communication capabilities.
pub mod client_connection;
pub mod iota_connection;
pub mod rho_connection;
pub mod rho_manager;
// Re-export commonly used types for convenience
pub use client_connection::ClientConnection;
pub use iota_connection::IotaConnection;
pub use rho_connection::RhoConnection;
// Re-export key manager functions
pub use rho_manager::{
add_rho, connection_count, contains_iota, get_all_connections, get_rho_by_iota,
get_rho_con_for_user, remove_rho,
};

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

@ -0,0 +1,251 @@
use std::collections::{self, HashMap};
use std::sync::{Arc, Weak};
use tokio::sync::RwLock;
use uuid::Uuid;
use super::{client_connection::ClientConnection, iota_connection::IotaConnection, rho_manager};
use crate::data::{
communication::{CommunicationType, CommunicationValue, DataTypes},
user::UserStatus,
};
use crate::omega::omega_connection::OmegaConnection;
pub struct RhoConnection {
iota_connection: Arc<IotaConnection>,
user_ids: Vec<Uuid>,
client_connections: Arc<RwLock<Vec<Arc<ClientConnection>>>>,
}
impl RhoConnection {
/// Create a new RhoConnection
pub async fn new(iota_connection: Arc<IotaConnection>, user_ids: Vec<Uuid>) -> Self {
let rho_connection = Self {
iota_connection,
user_ids: user_ids.clone(),
client_connections: Arc::new(RwLock::new(Vec::new())),
};
// Notify OmegaConnection about the new Iota
OmegaConnection::connect_iota(rho_connection.get_iota_id().await, user_ids);
rho_connection
}
/// Get the Iota ID
pub async fn get_iota_id(&self) -> Uuid {
self.iota_connection.get_iota_id().await
}
/// Get the user IDs associated with this Rho connection
pub fn get_user_ids(&self) -> &Vec<Uuid> {
&self.user_ids
}
/// Get reference to the IotaConnection
pub fn get_iota_connection(&self) -> &Arc<IotaConnection> {
&self.iota_connection
}
/// Get all client connections
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: Uuid,
) -> 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.unwrap() == user_id {
collections.push(con.clone());
}
}
collections
}
/// Add a client connection
pub async fn add_client_connection(&self, connection: Arc<ClientConnection>) {
// Notify Iota about new client
let notification = CommunicationValue::new(CommunicationType::client_connected)
.add_data_str(
DataTypes::user_id,
connection
.get_user_id()
.await
.unwrap_or(Uuid::nil())
.to_string(),
);
self.iota_connection.send_message(notification).await;
// Add to our list
{
let mut connections = self.client_connections.write().await;
connections.push(Arc::clone(&connection));
}
// Notify OmegaConnection
OmegaConnection::client_changed(
self.get_iota_id().await,
connection.get_user_id().await.unwrap_or(Uuid::nil()),
UserStatus::online,
);
}
/// Remove a client connection
pub async fn close_client_connection(&self, connection: Arc<ClientConnection>) {
{
let mut connections = self.client_connections.write().await;
let target_user_id = connection.get_user_id().await.unwrap();
connections.retain(|con| {
futures::executor::block_on(async {
con.get_user_id().await.unwrap() != target_user_id
})
});
// Push the new connection
connections.push(Arc::clone(&connection));
}
// Notify OmegaConnection
OmegaConnection::client_changed(
self.get_iota_id().await,
connection.get_user_id().await.unwrap_or(Uuid::nil()),
UserStatus::user_offline,
);
}
/// 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).await;
// Notify OmegaConnection
OmegaConnection::close_iota(self.get_iota_id().await);
}
/// Send message from Iota to specific client by user ID
pub async fn message_iota_to_client_by_user(&self, user_id: Uuid, message: &str) {
let connections = self.client_connections.read().await;
for connection in connections.iter() {
if let Some(conn_user_id) = connection.get_user_id().await {
if conn_user_id == user_id {
connection.send_message_str(message).await;
}
}
}
}
/// Send message from Iota to specific client
pub async fn message_iota_to_client(&self, cv: CommunicationValue) {
if let Some(receiver_id) = Some(cv.get_receiver()) {
let connections = self.client_connections.read().await;
for connection in connections.iter() {
if let Some(conn_user_id) = connection.get_user_id().await {
if conn_user_id == receiver_id {
connection.send_message(&cv).await;
}
}
}
}
}
/// Send message to Iota as string
pub async fn message_to_iota_str(&self, message: &str) {
self.iota_connection.send_message_str(message).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: Uuid, interested_ids: Vec<Uuid>) {
let connections = self.client_connections.read().await;
for connection in connections.iter() {
if let Some(conn_user_id) = connection.get_user_id().await {
if conn_user_id == user_id {
connection
.set_interested_users(interested_ids.clone())
.await;
break;
}
}
}
}
/// Check if clients are interested in a user
pub async fn are_they_interested(&self, user: &crate::data::user::User) {
let connections = self.client_connections.read().await;
for connection in connections.iter() {
connection.are_you_interested(user).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() {
if let Some(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
pub fn contains_user(&self, user_id: &Uuid) -> bool {
self.user_ids.contains(user_id)
}
/// Get count of active client connections
pub async fn client_count(&self) -> usize {
let connections = self.client_connections.read().await;
connections.len()
}
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::sync::Mutex;
#[tokio::test]
async fn test_rho_connection_creation() {
let iota_id = Uuid::new_v4();
let user_ids = vec![Uuid::new_v4(), Uuid::new_v4()];
// Mock session
// Create a mock WebSocket stream (this would fail in practice but shows the API)
// In real implementation, this would be a proper WebSocketStream
let mock_stream =
std::ptr::null_mut() as *mut tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>;
let mock_stream = unsafe { std::ptr::read(mock_stream) };
let iota_conn = IotaConnection::new_with_ids(
iota_id,
user_ids.clone(),
Arc::new(tokio::sync::Mutex::new(mock_stream)),
);
let rho_conn = RhoConnection::new(iota_conn, user_ids.clone()).await;
assert_eq!(rho_conn.get_iota_id().await, iota_id);
assert_eq!(rho_conn.get_user_ids(), &user_ids);
assert_eq!(rho_conn.client_count().await, 0);
}
}

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

@ -0,0 +1,93 @@
use std::{
collections::HashMap,
sync::{Arc, LazyLock},
};
use tokio::sync::RwLock;
use uuid::Uuid;
use super::rho_connection::RhoConnection;
// Static storage for RhoConnections, keyed by Iota ID
pub static RHO_CONNECTIONS: LazyLock<Arc<RwLock<HashMap<Uuid, Arc<RhoConnection>>>>> =
LazyLock::new(|| Arc::new(RwLock::new(HashMap::new())));
/// Get a RhoConnection for a specific user ID
pub async fn get_rho_con_for_user(user_id: Uuid) -> Option<Arc<RhoConnection>> {
let connections = RHO_CONNECTIONS.read().await;
for rho_connection in connections.values() {
if rho_connection.get_user_ids().contains(&user_id) {
return Some(Arc::clone(rho_connection));
}
}
None
}
/// Check if an Iota ID exists in the connections
pub async fn contains_iota(iota_id: Uuid) -> bool {
let connections = RHO_CONNECTIONS.read().await;
connections.contains_key(&iota_id)
}
/// Remove a RhoConnection by Iota ID
pub async fn remove_rho(iota_id: Uuid) -> 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, rho_connection);
}
/// Get a RhoConnection by Iota ID directly
pub async fn get_rho_by_iota(iota_id: Uuid) -> Option<Arc<RhoConnection>> {
let connections = RHO_CONNECTIONS.read().await;
connections.get(&iota_id).map(Arc::clone)
}
/// Get all active RhoConnections
pub async fn get_all_connections() -> Vec<Arc<RhoConnection>> {
let connections = RHO_CONNECTIONS.read().await;
connections.values().map(Arc::clone).collect()
}
/// Get the count of active connections
pub async fn connection_count() -> usize {
let connections = RHO_CONNECTIONS.read().await;
connections.len()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::rho::iota_connection::IotaConnection;
use std::sync::Arc;
use tokio::sync::Mutex;
#[tokio::test]
async fn test_add_and_get_rho() {
// Clear any existing connections
{
let mut connections = RHO_CONNECTIONS.write().await;
connections.clear();
}
let iota_id = Uuid::new_v4();
let user_ids = vec![Uuid::new_v4(), Uuid::new_v4()];
// Skip this test due to WebSocket complexity - would require proper mock setup
return;
// This test would need proper WebSocket stream mocking:
// let mock_session = create_mock_websocket_stream();
// let iota_conn = IotaConnection::new_with_ids(iota_id, user_ids.clone(), mock_session);
// let rho_conn = Arc::new(RhoConnection::new(iota_conn, user_ids.clone()));
// Test assertions would go here:
// add_rho(Arc::clone(&rho_conn)).await;
// assert!(contains_iota(iota_id).await);
// etc.
}
}

63
src/util/config_util.rs Normal file
View file

@ -0,0 +1,63 @@
use crate::util::file_util::load_file;
use json::JsonValue;
use once_cell::sync::Lazy;
use serde::{Deserialize, Serialize};
use std::sync::RwLock;
use uuid::Uuid;
const CONFIG_PATH: &str = "";
const CONFIG_FILENAME: &str = "config.json";
// Define the Config structure
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Config {
pub omega_server: String,
pub auth_server: String,
pub omikron_id: Uuid,
pub keep_people_stored_for: i32,
pub max_data: u64,
pub ip: String,
pub port: u16,
}
// Provide default values (similar to Java static defaults)
impl Default for Config {
fn default() -> Self {
Self {
omega_server: "omega.tensamin.methanium.net".into(),
auth_server: "auth.tensamin.methanium.net".into(),
omikron_id: Uuid::parse_str("a9e92dd6-08a6-4765-abf1-9fa39d0a99f9").unwrap(),
keep_people_stored_for: 90,
max_data: 1000 * 1000 * 1000 * 8,
ip: "0.0.0.0".into(),
port: 959,
}
}
}
// Global instance with thread-safe read/write access
pub static CONFIG: Lazy<RwLock<Config>> = Lazy::new(|| {
let config = Config::load().unwrap_or_default();
RwLock::new(config)
});
impl Config {
pub fn load() -> Option<Self> {
let content = load_file(CONFIG_PATH, CONFIG_FILENAME);
if content.trim().is_empty() {
return None;
}
let json = json::parse(&content).ok()?;
Some(Self {
omega_server: json["omega_server"].as_str().unwrap_or_default().into(),
auth_server: json["auth_server"].as_str().unwrap_or_default().into(),
omikron_id: Uuid::parse_str(json["omikron_id"].as_str().unwrap_or_default())
.unwrap_or_default(),
keep_people_stored_for: json["keep_people_stored_for"].as_i64().unwrap_or_default()
as i32,
max_data: json["max_data"].as_u64().unwrap_or_default(),
ip: json["ip"].as_str().unwrap_or_default().into(),
port: json["port"].as_u64().unwrap_or_default() as u16,
})
}
}

165
src/util/file_util.rs Normal file
View file

@ -0,0 +1,165 @@
use std::ffi::OsStr;
use std::fmt::Write as FmtWrite;
use std::fs::{self, File};
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use std::process;
use std::time::SystemTime;
use sysinfo::System;
use uuid::Uuid;
use walkdir::WalkDir;
pub fn delete_file(path: &str, name: &str) -> bool {
let dir = Path::new(&get_directory()).join(path);
let file = dir.join(name);
if !file.exists() {
return false;
}
fs::remove_file(file).is_ok()
}
pub fn delete_directory(path: &str) -> bool {
let dir = Path::new(&get_directory()).join(path);
delete_dir_recursive(&dir)
}
fn delete_dir_recursive(directory: &Path) -> bool {
if !directory.exists() {
return false;
}
if let Err(e) = fs::remove_dir_all(directory) {
println!(
"[IMPORTANT] Couldn't delete directory {}: {}",
directory.display(),
e
);
return false;
}
true
}
pub fn delete_user_directory(user_id: Uuid) {
let user_dir = Path::new(&get_directory())
.join("users")
.join(user_id.to_string());
let _ = delete_dir_recursive(&user_dir);
}
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) {
println!("[IMPORTANT] Couldn't create directories: {}", e);
return String::new();
}
return String::new();
}
if !file_path.exists() {
if let Err(e) = File::create(&file_path) {
println!("[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 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) {
println!("[IMPORTANT] Couldn't create directories: {}", e);
return;
}
}
if let Err(e) = fs::write(&file_path, value) {
println!(
"[IMPORTANT] Couldn't write file {}: {}",
file_path.display(),
e
);
}
}
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()
}
pub fn used_space() -> u64 {
get_directory_size(&PathBuf::from(get_directory()))
}
pub fn get_directory_size(directory: &Path) -> u64 {
let mut size = 0;
for entry in WalkDir::new(directory).into_iter().filter_map(|e| e.ok()) {
let path = entry.path();
if path.is_file() {
if let Ok(metadata) = path.metadata() {
size += path.file_name().unwrap_or(OsStr::new("")).len() as u64;
size += metadata.len();
}
}
}
size
}
pub fn get_designed_storage(user_id: Uuid) -> String {
let user_dir = Path::new(&get_directory())
.join("users")
.join(user_id.to_string());
design_byte(get_directory_size(&user_dir))
}
pub fn design_byte(bytes: u64) -> String {
let mut hr_size = format!("{:.2}B", bytes as f64);
let k = bytes as f64 / 1024.0;
let m = k / 1024.0;
let g = m / 1024.0;
let t = g / 1024.0;
if t >= 1.0 {
hr_size = format!("{:.2}TB", t);
} else if g >= 1.0 {
hr_size = format!("{:.2}GB", g);
} else if m >= 1.0 {
hr_size = format!("{:.2}MB", m);
} else if k >= 1.0 {
hr_size = format!("{:.2}KB", k);
}
hr_size
}
pub fn get_used_ram() -> String {
let mut sys = System::new_all();
sys.refresh_all();
let used = sys.used_memory() * 1024; // kB to bytes
let total = sys.total_memory() * 1024;
format!("{}/{}", design_byte(used), design_byte(total))
}

2
src/util/mod.rs Normal file
View file

@ -0,0 +1,2 @@
pub mod config_util;
pub mod file_util;