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

[Fix] Hopefully Iota - Omikron comms work now.

Browse source 
[todo] fix Omikron - Omega comms
This commit is contained in:
Alex Emmet 2026-03-26 02:03:13 +01:00
commit 60c980ba9a
9 changed files with 493 additions and 89 deletions
Download patch file Download diff file Expand all files Collapse all files

5
src/anonymous_clients/anonymous_client_connection.rs
View file

@ -49,6 +49,7 @@ impl AnonymousClientConnection {
while let Ok(cv) = self_clone.receiver.receive().await {
self_clone.clone().handle_message(cv).await;
}
self_clone.handle_close().await;
});
}
@ -457,11 +458,10 @@ impl AnonymousClientConnection {
let error = CommunicationValue::new(error_type).with_id(*message_id);
self.send_message(&error).await;
}
#[allow(dead_code)]
/// Close the connection
pub async fn close(&self) {
let mut is_open_guard = self.is_open.write().await;
if *is_open_guard {
if !*is_open_guard {
return;
}
*is_open_guard = false;
@ -494,7 +494,6 @@ impl AnonymousClientConnection {
}
}
#[allow(dead_code)]
/// Handle connection close
pub async fn handle_close(&self) {

66
src/rho/client_connection.rs
View file

@ -5,7 +5,7 @@ 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_out};
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};
@ -31,7 +31,7 @@ impl ClientConnection {
Arc::new(Self {
ping: Arc::new(RwLock::new(0)),
pub_key: Arc::new(RwLock::new(None)),
rho_connection: 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(),
@ -45,6 +45,7 @@ impl ClientConnection {
while let Ok(cv) = self_clone.receiver.receive().await {
self_clone.clone().handle_message(cv).await;
}
self_clone.handle_close().await;
});
}
@ -393,6 +394,20 @@ impl ClientConnection {
/// 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)
@ -434,17 +449,56 @@ impl ClientConnection {
};
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(self.get_user_id().await as u64)
.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 updated_cv = cv.with_sender(self.get_user_id().await as u64);
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
);
}
}
@ -455,10 +509,9 @@ impl ClientConnection {
}
/// Close the connection
#[allow(dead_code)]
pub async fn close(&self) {
let mut is_open_guard = self.is_open.write().await;
if *is_open_guard {
if !*is_open_guard {
return;
}
*is_open_guard = false;
@ -491,7 +544,6 @@ impl ClientConnection {
}
/// Handle connection close
#[allow(dead_code)]
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 {

77
src/rho/connection.rs
View file

@ -37,6 +37,7 @@ pub struct GeneralConnection {
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>>>>,
@ -50,6 +51,7 @@ impl GeneralConnection {
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)),
})
@ -76,13 +78,13 @@ impl GeneralConnection {
if !*self.challenged.read().await {
self.handle_challenge_response(cv).await;
if *self.challenged.read().await {
break;
}
continue;
}
if self.migrate().await {
if *self.challenged.read().await {
let self_clone = self.clone();
tokio::spawn(async move {
self_clone.migrate().await;
});
break;
}
}
@ -233,10 +235,6 @@ impl GeneralConnection {
if let Err(_) = self.sender.send(&response).await {
return;
}
if self.migrate().await {
log_out!(id, PrintType::Iota, "Immediate migration");
}
} else {
log_err!(
id,
@ -270,7 +268,48 @@ impl GeneralConnection {
.add_data(DataTypes::user_id, DataValue::Number(id as i64));
get_omega_connection().send_message(&notify).await;
let user_id = id as i64;
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();
let client = ClientConnection::from_general(self.clone(), id).await;
if let Some(rho_conn) = rho {
// Make sure user is bound before the client starts forwarding
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 => {
@ -282,7 +321,25 @@ impl GeneralConnection {
let rho = Arc::new(RhoConnection::new(iota.clone(), Vec::new()).await);
iota.set_rho_connection(Arc::downgrade(&rho)).await;
iota.set_rho_connection(rho.clone()).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;
}
}
rho_manager::add_rho(rho).await;

69
src/rho/iota_connection.rs
View file

@ -3,16 +3,13 @@ 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, Weak},
time::Duration,
};
use std::{collections::HashMap, sync::Arc, time::Duration};
use tokio::sync::RwLock;
use tokio::sync::mpsc;
use ttp_core::CommunicationType;
@ -36,7 +33,7 @@ pub struct IotaConnection {
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<Weak<RhoConnection>>>>,
pub rho_connection: Arc<RwLock<Option<Arc<RhoConnection>>>>,
}
impl IotaConnection {
@ -44,7 +41,7 @@ impl IotaConnection {
Arc::new(Self {
ping: Arc::new(RwLock::new(0)),
pub_key: Arc::new(RwLock::new(None)),
rho_connection: 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(),
@ -65,6 +62,7 @@ impl IotaConnection {
}
}
}
self_clone.handle_close().await;
});
}
@ -87,13 +85,43 @@ impl IotaConnection {
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);
}
}
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: Weak<RhoConnection>) {
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);
}
@ -102,7 +130,7 @@ impl IotaConnection {
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()
Some(weak_ref.clone())
} else {
None
}
@ -217,8 +245,20 @@ impl IotaConnection {
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;
if self.get_user_ids().await.contains(&(sender_id as u64)) {
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 {
@ -228,6 +268,14 @@ impl IotaConnection {
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,
@ -350,7 +398,6 @@ impl IotaConnection {
}
}
#[allow(dead_code)]
pub async fn handle_close(&self) {
if let Some(rho_conn) = self.get_rho_connection().await {
rho_conn.close_iota_connection().await;

29
src/rho/rho_connection.rs
View file

@ -9,7 +9,7 @@ use ttp_core::{CommunicationType, CommunicationValue, DataTypes, DataValue};
pub struct RhoConnection {
iota_connection: Arc<IotaConnection>,
user_ids: Vec<i64>,
user_ids: Arc<RwLock<Vec<i64>>>,
client_connections: Arc<RwLock<Vec<Arc<ClientConnection>>>>,
}
@ -18,7 +18,7 @@ impl RhoConnection {
pub async fn new(iota_connection: Arc<IotaConnection>, user_ids: Vec<i64>) -> Self {
let rho_connection = Self {
iota_connection,
user_ids: user_ids.clone(),
user_ids: Arc::new(RwLock::new(user_ids.clone())),
client_connections: Arc::new(RwLock::new(Vec::new())),
};
@ -29,8 +29,25 @@ impl RhoConnection {
self.iota_connection.iota_id
}
pub fn get_user_ids(&self) -> &Vec<i64> {
&self.user_ids
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> {
@ -169,8 +186,8 @@ impl RhoConnection {
/// Check if this RhoConnection contains a specific user ID
#[allow(dead_code)]
pub fn contains_user(&self, user_id: &i64) -> bool {
self.user_ids.contains(user_id)
pub async fn contains_user(&self, user_id: &i64) -> bool {
self.user_ids.read().await.contains(user_id)
}
/// Get count of active client connections

28
src/rho/rho_manager.rs
View file

@ -13,13 +13,14 @@ pub static RHO_CONNECTIONS: LazyLock<Arc<RwLock<HashMap<i64, Arc<RhoConnection>>
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_connection.get_user_ids().to_vec()
rho_user_ids
);
if rho_connection.get_user_ids().contains(&user_id) {
if rho_user_ids.contains(&user_id) {
return Some(Arc::clone(rho_connection));
}
}
@ -32,6 +33,29 @@ pub async fn contains_iota(iota_id: i64) -> bool {
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;

44
src/util/crypto_util.rs
View file

@ -4,12 +4,12 @@ use aes_gcm::{
};
use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64_STD};
use hkdf::Hkdf;
use sha2::{Digest, Sha256};
type HkdfSha256 = sha2::Sha256;
use sha2::{Digest, Sha256 as HashSha256};
use x448::{PublicKey, Secret};
// --- Custom Errors ---
#[allow(dead_code)]
#[derive(Debug)]
#[allow(dead_code)]
pub enum SecurePayloadError {
InvalidBase64,
InvalidHex,
@ -18,9 +18,8 @@ pub enum SecurePayloadError {
InvalidKeyLength,
}
// --- Data Format Enum ---
#[allow(dead_code)]
#[derive(Clone, Copy, Debug)]
#[allow(dead_code)]
pub enum DataFormat {
Raw,
Base64,
@ -41,8 +40,8 @@ impl Clone for SecurePayload {
}
}
#[allow(dead_code)]
impl SecurePayload {
/// Clear Constructor: Takes data in any format and the user's private key.
pub fn new<S, T: AsRef<[u8]>>(
data: T,
format: DataFormat,
@ -67,13 +66,10 @@ impl SecurePayload {
})
}
/// Helper to get the public key associated with this instance's private key.
#[allow(dead_code)]
pub fn get_public_key(&self) -> [u8; 56] {
*PublicKey::from(&self.private_key).as_bytes()
}
/// Exports the internal data to the requested format
pub fn export(&self, format: DataFormat) -> String {
match format.into() {
DataFormat::Raw => String::from_utf8_lossy(&self.inner_data).to_string(),
@ -82,16 +78,12 @@ impl SecurePayload {
}
}
/// Access raw bytes directly
#[allow(dead_code)]
pub fn get_bytes(&self) -> &[u8] {
&self.inner_data
}
/// Returns the SHA-256 Hash of the data in the requested format
#[allow(dead_code)]
pub fn get_hash(&self, format: DataFormat) -> String {
let mut hasher = Sha256::new();
let mut hasher = HashSha256::new();
hasher.update(&self.inner_data);
let result = hasher.finalize();
@ -102,8 +94,6 @@ impl SecurePayload {
}
}
/// Encrypts the held data for a specific recipient using AES-256-GCM.
/// The message will contain ONLY the ciphertext.
pub fn encrypt_x448<S>(&self, public_key: S) -> Result<SecurePayload, SecurePayloadError>
where
S: Into<PublicKey>,
@ -111,8 +101,9 @@ impl SecurePayload {
let peer_pub = public_key.into();
let shared_secret = self.private_key.as_diffie_hellman(&peer_pub).unwrap();
let hkdf = Hkdf::<Sha256>::new(None, shared_secret.as_bytes());
let hkdf = Hkdf::<HkdfSha256>::new(None, shared_secret.as_bytes());
let mut okm = [0u8; 44];
hkdf.expand(b"x448-aes-gcm-no-overhead", &mut okm)
.map_err(|_| SecurePayloadError::EncryptionError)?;
@ -138,26 +129,27 @@ impl SecurePayload {
})
}
/// Decrypts the held data providing the sender's public key manually.
#[allow(dead_code)]
pub fn decrypt_to_format(
&self,
peer_public_key_bytes: &[u8; 56],
output_format: DataFormat,
) -> Result<String, SecurePayloadError> {
let decrypted_instance = self.decrypt_x448(peer_public_key_bytes)?;
let decrypted_instance =
self.decrypt_x448(PublicKey::from_bytes(peer_public_key_bytes).unwrap())?;
Ok(decrypted_instance.export(output_format))
}
/// Decrypts the held data using the internal Private Key and the provided Peer Public Key.
pub fn decrypt_x448(
pub fn decrypt_x448<S>(
&self,
peer_public_key_bytes: &[u8; 56],
) -> Result<SecurePayload, SecurePayloadError> {
let peer_pub = PublicKey::from_bytes(peer_public_key_bytes).unwrap();
peer_public_key_bytes: S,
) -> Result<SecurePayload, SecurePayloadError>
where
S: Into<PublicKey>,
{
let peer_pub = peer_public_key_bytes.into();
let shared_secret = self.private_key.as_diffie_hellman(&peer_pub).unwrap();
let hkdf = Hkdf::<Sha256>::new(None, shared_secret.as_bytes());
let hkdf = Hkdf::<HkdfSha256>::new(None, shared_secret.as_bytes());
let mut okm = [0u8; 44];
hkdf.expand(b"x448-aes-gcm-no-overhead", &mut okm)
.map_err(|_| SecurePayloadError::DecryptionError)?;