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

Auth Paths fixed

Browse source 
OUTLINE of calls Rho: Iota & Clients. Works Correct File handeling

TODO: Omega Connection
This commit is contained in:
Alex Emmet 2025-10-19 02:14:20 +02:00
commit ba81f3651e
18 changed files with 488 additions and 477 deletions
Download patch file Download diff file Expand all files Collapse all files

115
src/rho/client_connection.rs
View file

@ -1,10 +1,7 @@
use futures::SinkExt;
use http::header::AUTHORIZATION;
use std::{
collections::HashMap,
sync::{Arc, Weak},
};
use tokio::sync::{Mutex, RwLock};
use std::sync::{Arc, Weak};
use tokio::sync::Mutex;
use tokio::sync::RwLock;
use tokio_tungstenite::{WebSocketStream, tungstenite::Message};
use tungstenite::Utf8Bytes;
use uuid::Uuid;
@ -23,17 +20,17 @@ use crate::{
/// ClientConnection represents a WebSocket connection from a client device
pub struct ClientConnection {
/// WebSocket session
session: Arc<Mutex<WebSocketStream<tokio::net::TcpStream>>>,
pub session: Arc<Mutex<WebSocketStream<tokio::net::TcpStream>>>,
/// User ID associated with this client
user_id: Arc<RwLock<Option<Uuid>>>,
pub user_id: Arc<RwLock<Option<Uuid>>>,
/// Whether this connection has been identified/authenticated
identified: Arc<RwLock<bool>>,
pub identified: Arc<RwLock<bool>>,
/// Ping latency tracking
ping: Arc<RwLock<i64>>,
pub ping: Arc<RwLock<i64>>,
/// Weak reference to RhoConnection to avoid circular references
rho_connection: Arc<RwLock<Option<Weak<RhoConnection>>>>,
pub 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>>>,
pub interested_users: Arc<RwLock<Vec<Uuid>>>,
}
impl ClientConnection {
@ -97,7 +94,7 @@ impl ClientConnection {
}
/// Handle incoming message from client
pub async fn handle_message(self: Arc<Self>, message: String) {
pub async fn handle_message(self: Arc<Self>, message: Utf8Bytes) {
let cv = CommunicationValue::from_json(&message);
// Handle identification
@ -145,23 +142,16 @@ impl ClientConnection {
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;
self.send_error_response(
&cv.get_id(),
CommunicationType::error_invalid_user_id,
)
.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)
self.send_error_response(&cv.get_id(), CommunicationType::error_invalid_user_id)
.await;
return;
}
@ -169,20 +159,36 @@ impl ClientConnection {
// Validate private key
if let Some(private_key_hash) = cv.get_data(DataTypes::private_key_hash) {
println!("private_key_hash: {}", 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;
println!("Invalid private key");
self.send_error_response(
&cv.get_id(),
CommunicationType::error_invalid_private_key,
)
.await;
return;
}
} else {
self.send_error_response(&cv.get_id(), CommunicationType::error)
println!("Missing private key");
self.send_error_response(&cv.get_id(), CommunicationType::error_invalid_private_key)
.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_no_iota)
.await;
return;
}
};
// Set identification data
{
let mut user_id_guard = self.user_id.write().await;
@ -193,14 +199,11 @@ impl ClientConnection {
*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;
@ -242,14 +245,15 @@ impl ClientConnection {
rho_conn.get_iota_id().await,
user_id,
user_status,
);
)
.await;
}
}
}
}
/// Handle call invite
async fn handle_call_invite(&self, mut cv: CommunicationValue) {
async fn handle_call_invite(&self, 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,
@ -470,46 +474,3 @@ impl std::fmt::Debug for ClientConnection {
.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
}
}

75
src/rho/iota_connection.rs
View file

@ -1,5 +1,5 @@
use futures::SinkExt;
use json::{JsonValue, number::Number};
use json::JsonValue;
use std::{
collections::HashMap,
sync::{Arc, Weak},
@ -13,21 +13,18 @@ use super::{rho_connection::RhoConnection, rho_manager};
use crate::{
auth::auth_connector,
// calls::call_manager::CallManager,
data::{
communication::{CommunicationType, CommunicationValue, DataTypes},
user::User,
},
data::communication::{CommunicationType, CommunicationValue, DataTypes},
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>>>>,
pub session: Arc<Mutex<WebSocketStream<tokio::net::TcpStream>>>,
pub iota_id: Arc<RwLock<Uuid>>,
pub user_ids: Arc<RwLock<Vec<Uuid>>>,
pub identified: Arc<RwLock<bool>>,
pub ping: Arc<RwLock<i64>>,
pub rho_connection: Arc<RwLock<Option<Weak<RhoConnection>>>>,
}
impl IotaConnection {
@ -43,22 +40,6 @@ impl IotaConnection {
})
}
/// 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
@ -98,18 +79,24 @@ impl IotaConnection {
/// Send a message to the Iota
pub async fn send_message_str(&self, message: &str) {
let mut session = self.session.lock().await;
let _ = session
if let Err(e) = session
.send(Message::Text(Utf8Bytes::from(message.to_string())))
.await;
.await
{
eprintln!("Failed to send WebSocket message: {:?}", e);
}
}
/// Send a CommunicationValue to the Iota
pub async fn send_message(&self, cv: CommunicationValue) {
if !cv.is_type(CommunicationType::pong) {
println!("{}", cv.to_json().to_string());
}
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) {
pub async fn handle_message(self: Arc<Self>, message: Utf8Bytes) {
let cv = CommunicationValue::from_json(&message);
// Handle identification
@ -149,7 +136,7 @@ impl IotaConnection {
}
/// Handle identification message
async fn handle_identification(self: Arc<Self>, mut cv: CommunicationValue) {
async fn handle_identification(self: Arc<Self>, 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()) {
@ -166,16 +153,17 @@ impl IotaConnection {
};
// Parse user IDs
let mut validated_user_ids = Vec::new();
let mut validated_user_ids: Vec<Uuid> = 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);
if let Some(auth_iota_id) = auth_connector::get_iota_id(user_id).await {
if auth_iota_id == iota_id {
validated_user_ids.push(user_id);
}
}
}
}
@ -214,16 +202,20 @@ impl IotaConnection {
rho_manager::add_rho(rho_connection).await;
// Send response
let mut str = String::new();
for id in &validated_user_ids {
str.push_str(&format!(",{}", id));
}
let response = CommunicationValue::new(CommunicationType::identification_response)
.with_id(cv.get_id())
.add_data_str(DataTypes::accepted_ids, str)
.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
async fn handle_ping(&self, cv: CommunicationValue) {
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;
@ -231,14 +223,12 @@ impl IotaConnection {
}
}
// 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())))
@ -252,18 +242,15 @@ impl IotaConnection {
/// 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());
@ -298,7 +285,7 @@ impl IotaConnection {
}
// Notify OmegaConnection about user states
OmegaConnection::user_states(receiver_id, interested_ids.clone());
OmegaConnection::user_states(receiver_id, interested_ids.clone()).await;
// Set interested users in RhoConnection
if let Some(rho_conn) = self.get_rho_connection().await {
@ -318,13 +305,11 @@ impl IotaConnection {
}
}
/// 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 {

17
src/rho/mod.rs
View file

@ -1,21 +1,4 @@
//! 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,
};

51
src/rho/rho_connection.rs
View file

@ -1,5 +1,5 @@
use std::collections::{self, HashMap};
use std::sync::{Arc, Weak};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use uuid::Uuid;
@ -26,27 +26,23 @@ impl RhoConnection {
};
// Notify OmegaConnection about the new Iota
OmegaConnection::connect_iota(rho_connection.get_iota_id().await, user_ids);
OmegaConnection::connect_iota(rho_connection.get_iota_id().await, user_ids).await;
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()
@ -69,7 +65,6 @@ impl RhoConnection {
/// 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,
@ -82,18 +77,17 @@ impl RhoConnection {
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,
);
)
.await;
}
/// Remove a client connection
@ -109,7 +103,6 @@ impl RhoConnection {
})
});
// Push the new connection
connections.push(Arc::clone(&connection));
}
@ -118,7 +111,8 @@ impl RhoConnection {
self.get_iota_id().await,
connection.get_user_id().await.unwrap_or(Uuid::nil()),
UserStatus::user_offline,
);
)
.await;
}
/// Close the Iota connection and all associated client connections
@ -133,7 +127,7 @@ impl RhoConnection {
rho_manager::remove_rho(self.get_iota_id().await).await;
// Notify OmegaConnection
OmegaConnection::close_iota(self.get_iota_id().await);
OmegaConnection::close_iota(self.get_iota_id().await).await;
}
/// Send message from Iota to specific client by user ID
@ -220,32 +214,3 @@ impl RhoConnection {
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);
}
}

8
src/rho/rho_manager.rs
View file

@ -7,14 +7,17 @@ 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;
println!("Checking user ID: {:?}", user_id);
for rho_connection in connections.values() {
println!(
"Comparing user IDs: {:?}",
rho_connection.get_user_ids().to_vec()
);
if rho_connection.get_user_ids().contains(&user_id) {
return Some(Arc::clone(rho_connection));
}
@ -22,7 +25,6 @@ pub async fn get_rho_con_for_user(user_id: Uuid) -> Option<Arc<RhoConnection>> {
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)