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

[Fix] Iota connection

Browse source
This commit is contained in:
Alex Emmet 2026-03-09 22:23:54 +01:00
commit e2d5c24ac4
5 changed files with 445 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

17
Cargo.lock generated
View file

@ -572,8 +572,9 @@ checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
[[package]]
name = "epsilon-core"
version = "0.1.0"
source = "git+https://github.com/Tensamin/Epsilon.git#05a7b22b2e05a9b023b0d976d65dddf06edccd7c"
source = "git+https://github.com/Tensamin/Epsilon.git#809a32581843e663d35d4b3222f2cc59371d2786"
dependencies = [
"base64 0.22.1",
"byteorder",
"rand 0.8.5",
"strum",
@ -583,7 +584,7 @@ dependencies = [
[[package]]
name = "epsilon-native"
version = "0.1.0"
source = "git+https://github.com/Tensamin/Epsilon.git#05a7b22b2e05a9b023b0d976d65dddf06edccd7c"
source = "git+https://github.com/Tensamin/Epsilon.git#809a32581843e663d35d4b3222f2cc59371d2786"
dependencies = [
"epsilon-core",
"quinn",
@ -1799,9 +1800,9 @@ dependencies = [
[[package]]
name = "quinn-proto"
version = "0.11.13"
version = "0.11.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f1906b49b0c3bc04b5fe5d86a77925ae6524a19b816ae38ce1e426255f1d8a31"
checksum = "434b42fec591c96ef50e21e886936e66d3cc3f737104fdb9b737c40ffb94c098"
dependencies = [
"aws-lc-rs",
"bytes",
@ -3267,18 +3268,18 @@ dependencies = [
[[package]]
name = "zerocopy"
version = "0.8.40"
version = "0.8.42"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a789c6e490b576db9f7e6b6d661bcc9799f7c0ac8352f56ea20193b2681532e5"
checksum = "f2578b716f8a7a858b7f02d5bd870c14bf4ddbbcf3a4c05414ba6503640505e3"
dependencies = [
"zerocopy-derive",
]
[[package]]
name = "zerocopy-derive"
version = "0.8.40"
version = "0.8.42"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f65c489a7071a749c849713807783f70672b28094011623e200cb86dcb835953"
checksum = "7e6cc098ea4d3bd6246687de65af3f920c430e236bee1e3bf2e441463f08a02f"
dependencies = [
"proc-macro2",
"quote",

4
src/main.rs
View file

@ -41,7 +41,9 @@ async fn main() {
get_omega_connection();
tokio::spawn(async move {
let _ = start(959).await;
if let Err(e) = start(959).await {
log_err!(0, util::logger::PrintType::General, "{}", e);
}
});
garbage_collect_calls();

8
src/omega/omega_connection.rs
View file

@ -457,7 +457,9 @@ impl OmegaConnection {
result = receiver.receive() => {
match result {
Ok(cv) => {
log_cv_in!(&cv);
if !cv.is_type(CommunicationType::pong) && !cv.is_type(CommunicationType::ping) {
log_cv_in!(PrintType::Omega, &cv);
}
if cv.is_type(CommunicationType::pong) || cv.is_type(CommunicationType::ping) {
self.handle_pong(&cv).await;
@ -545,7 +547,9 @@ impl OmegaConnection {
// -------------------------------------------------------------------------
pub async fn send_message(&self, cv: &CommunicationValue) {
log_cv_out!(cv);
if !cv.is_type(CommunicationType::pong) && !cv.is_type(CommunicationType::ping) {
log_cv_out!(PrintType::Omega, &cv);
}
let sender_guard = self.sender.read().await;
if let Some(sender) = sender_guard.as_ref() {

280
src/rho/connection.rs
View file

@ -6,12 +6,16 @@ use tokio::sync::RwLock;
use crate::{
anonymous_clients::anonymous_client_connection::AnonymousClientConnection,
get_private_key, get_public_key,
get_private_key, get_public_key, log_err, log_in, log_out,
omega::omega_connection::get_omega_connection,
rho::{client_connection::ClientConnection, iota_connection::IotaConnection},
rho::{
client_connection::ClientConnection, iota_connection::IotaConnection,
rho_connection::RhoConnection, rho_manager,
},
util::{
crypto_helper::{load_public_key, public_key_to_base64},
crypto_util::{DataFormat, SecurePayload},
logger::PrintType,
},
};
@ -52,10 +56,31 @@ impl GeneralConnection {
}
impl GeneralConnection {
pub async fn handle(self: Arc<Self>) {
log_in!(0, PrintType::General, "General connection handler started");
loop {
let cv = match self.receiver.receive().await {
Ok(v) => v,
Err(_) => break,
Ok(v) => {
log_in!(
0,
PrintType::General,
"General connection received message type={:?} id={} identified={} challenged={}",
v.get_type(),
v.get_id(),
*self.identified.read().await,
*self.challenged.read().await
);
v
}
Err(e) => {
log_err!(
0,
PrintType::General,
"General connection receive error before upgrade completion: {:?}",
e
);
break;
}
};
if !*self.identified.read().await {
@ -65,33 +90,91 @@ impl GeneralConnection {
if !*self.challenged.read().await {
self.handle_challenge_response(cv).await;
if *self.challenged.read().await {
log_out!(
0,
PrintType::General,
"General connection challenge flow completed, handler will stop after immediate migration"
);
break;
}
continue;
}
log_in!(
0,
PrintType::General,
"General connection ready to migrate for id={} kind={:?}",
*self.id.read().await,
*self.connection_kind.read().await
);
if self.migrate().await {
log_out!(
0,
PrintType::General,
"General connection migration completed, handing over to specialized connection"
);
break;
}
}
log_out!(0, PrintType::General, "General connection handler stopped");
}
async fn handle_identification(self: &Arc<Self>, cv: CommunicationValue) {
if !cv.is_type(CommunicationType::identification) {
log_in!(
0,
PrintType::General,
"Ignoring pre-identification message type={:?} id={}",
cv.get_type(),
cv.get_id()
);
return;
}
if let DataValue::Number(iota_id) = cv.get_data(DataTypes::iota_id) {
log_in!(
*iota_id,
PrintType::Iota,
"Received Iota identification request message_id={}",
cv.get_id()
);
*self.id.write().await = *iota_id as u64;
*self.connection_kind.write().await = Some(ConnectionKind::Iota);
let get_pub_key_msg = CommunicationValue::new(CommunicationType::get_iota_data)
.add_data(DataTypes::iota_id, DataValue::Number(*iota_id));
log_out!(
*iota_id,
PrintType::Iota,
"Requesting Iota public key from Omega"
);
let response_cv = get_omega_connection()
.await_response(&get_pub_key_msg, Some(Duration::from_secs(20)))
.await;
let response_cv = match response_cv {
Ok(r) => r,
Err(_) => return,
Ok(r) => {
log_in!(
*iota_id,
PrintType::Iota,
"Received Iota public key response from Omega"
);
r
}
Err(e) => {
log_err!(
*iota_id,
PrintType::Iota,
"Failed to load Iota public key from Omega: {:?}",
e
);
return;
}
};
let base64_pub = response_cv
@ -101,7 +184,14 @@ impl GeneralConnection {
let pub_key = match load_public_key(base64_pub) {
Some(pk) => pk,
None => return,
None => {
log_err!(
*iota_id,
PrintType::Iota,
"Failed to decode Iota public key from Omega response"
);
return;
}
};
*self.pub_key.write().await = Some(pub_key.as_bytes().to_vec());
@ -112,11 +202,18 @@ impl GeneralConnection {
.map(char::from)
.collect();
log_out!(
*iota_id,
PrintType::Iota,
"Generated challenge for Iota identification challenge_len={}",
challenge.len()
);
*self.challenge.write().await = challenge.clone();
*self.identified.write().await = true;
let encrypted_challenge =
SecurePayload::new(&challenge, DataFormat::Base64, get_private_key())
SecurePayload::new(challenge.as_bytes(), DataFormat::Raw, get_private_key())
.unwrap()
.encrypt_x448(pub_key)
.unwrap()
@ -129,44 +226,205 @@ impl GeneralConnection {
)
.add_data(DataTypes::challenge, DataValue::Str(encrypted_challenge));
let _ = self.sender.send(&response).await;
log_out!(
*iota_id,
PrintType::Iota,
"Sending encrypted identification challenge to Iota"
);
if let Err(e) = self.sender.send(&response).await {
log_err!(
*iota_id,
PrintType::Iota,
"Failed to send challenge to Iota: {:?}",
e
);
}
} else {
log_err!(
0,
PrintType::General,
"Identification message missing iota_id payload"
);
}
}
async fn handle_challenge_response(self: &Arc<Self>, cv: CommunicationValue) {
let id = *self.id.read().await as i64;
if !cv.is_type(CommunicationType::challenge_response) {
log_in!(
id,
PrintType::Iota,
"Ignoring pre-challenge-completion message type={:?} id={}",
cv.get_type(),
cv.get_id()
);
return;
}
if let DataValue::Str(response) = cv.get_data(DataTypes::challenge) {
if *response == *self.challenge.read().await {
let expected = self.challenge.read().await.clone();
log_in!(
id,
PrintType::Iota,
"Received challenge response message_id={} response_len={} expected_len={}",
cv.get_id(),
response.len(),
expected.len()
);
if *response == expected {
log_in!(
id,
PrintType::Iota,
"Challenge response validated successfully"
);
*self.challenged.write().await = true;
let response = CommunicationValue::new(CommunicationType::identification_response)
.with_id(cv.get_id())
.add_data(DataTypes::accepted, DataValue::Bool(true));
log_out!(
id,
PrintType::Iota,
"Sending identification_response accepted=true"
);
if let Err(e) = self.sender.send(&response).await {
log_err!(
id,
PrintType::Iota,
"Failed to send identification_response: {:?}",
e
);
return;
}
log_in!(
id,
PrintType::Iota,
"Immediately migrating upgraded connection after successful challenge validation"
);
if self.migrate().await {
log_out!(
id,
PrintType::Iota,
"Immediate migration after challenge validation completed successfully"
);
} else {
log_err!(
id,
PrintType::Iota,
"Immediate migration after challenge validation failed"
);
}
} else {
log_err!(
id,
PrintType::Iota,
"Challenge response mismatch expected={} actual={}",
expected,
response
);
}
} else {
log_err!(
id,
PrintType::Iota,
"Challenge response missing challenge payload"
);
}
}
async fn migrate(self: &Arc<Self>) -> bool {
let kind = match *self.connection_kind.read().await {
Some(kind) => kind,
None => return false,
None => {
log_err!(
0,
PrintType::General,
"Migration requested without a resolved connection kind"
);
return false;
}
};
let id = *self.id.read().await;
log_in!(
id as i64,
PrintType::General,
"Starting migration for kind={:?} id={}",
kind,
id
);
match kind {
ConnectionKind::Client => {
let client = ClientConnection::from_general(self.clone(), id).await;
client.start();
log_out!(
id as i64,
PrintType::Client,
"Migrated general connection into ClientConnection"
);
}
ConnectionKind::Iota => {
let iota = IotaConnection::from_general(self.clone(), id).await;
log_in!(
id as i64,
PrintType::Iota,
"Created upgraded IotaConnection from GeneralConnection"
);
let rho = Arc::new(RhoConnection::new(iota.clone(), Vec::new()).await);
log_in!(
id as i64,
PrintType::Iota,
"Created RhoConnection for upgraded Iota connection"
);
iota.set_rho_connection(Arc::downgrade(&rho)).await;
log_in!(
id as i64,
PrintType::Iota,
"Attached weak RhoConnection reference to IotaConnection"
);
rho_manager::add_rho(rho).await;
log_out!(
id as i64,
PrintType::Iota,
"Registered upgraded Iota connection in rho_manager"
);
iota.start();
log_out!(
id as i64,
PrintType::Iota,
"Started upgraded IotaConnection read loop"
);
}
ConnectionKind::AnonymousClient => {
let client = AnonymousClientConnection::from_general(self.clone(), id).await;
client.start();
log_out!(
id as i64,
PrintType::Client,
"Migrated general connection into AnonymousClientConnection"
);
}
ConnectionKind::Phi => {
let iota = ClientConnection::from_general(self.clone(), id).await;
iota.start();
log_out!(
id as i64,
PrintType::General,
"Migrated general connection into Phi/Client handler"
);
}
}
true

159
src/rho/iota_connection.rs
View file

@ -3,6 +3,8 @@ 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::log_out;
use crate::omega::omega_connection::get_omega_connection;
use crate::rho::connection::GeneralConnection;
use crate::util::logger::PrintType;
@ -54,9 +56,41 @@ impl IotaConnection {
pub fn start(self: Arc<Self>) {
let self_clone = self.clone();
tokio::spawn(async move {
while let Ok(cv) = self_clone.receiver.receive().await {
log_in!(
self_clone.iota_id as i64,
PrintType::Iota,
"Upgraded IotaConnection read loop started"
);
loop {
match self_clone.receiver.receive().await {
Ok(cv) => {
log_in!(
self_clone.iota_id as i64,
PrintType::Iota,
"Upgraded IotaConnection received message type={:?} id={}",
cv.get_type(),
cv.get_id()
);
self_clone.clone().handle_message(cv).await;
}
Err(e) => {
log_err!(
self_clone.iota_id as i64,
PrintType::Iota,
"Upgraded IotaConnection read loop stopped due to receive error: {:?}",
e
);
break;
}
}
}
log_out!(
self_clone.iota_id as i64,
PrintType::Iota,
"Upgraded IotaConnection read loop exited"
);
});
}
@ -116,8 +150,23 @@ impl IotaConnection {
/// Handle incoming message from Iota
pub async fn handle_message(self: Arc<Self>, cv: CommunicationValue) {
log_in!(
self.iota_id as i64,
PrintType::Iota,
"Handling upgraded Iota message type={:?} id={} sender={} receiver={}",
cv.get_type(),
cv.get_id(),
cv.get_sender(),
cv.get_receiver()
);
// Handle ping
if cv.is_type(CommunicationType::ping) || cv.is_type(CommunicationType::pong) {
log_in!(
self.iota_id as i64,
PrintType::Iota,
"Routing upgraded Iota message to ping handler"
);
self.handle_ping(cv).await;
return;
}
@ -126,6 +175,11 @@ impl IotaConnection {
// Handle GET_CHATS
if cv.is_type(CommunicationType::get_chats) {
log_in!(
self.iota_id as i64,
PrintType::Iota,
"Routing upgraded Iota message to get_chats handler"
);
self.handle_get_chats(cv).await;
return;
}
@ -136,6 +190,12 @@ impl IotaConnection {
|| cv.is_type(CommunicationType::message_other_iota)
|| cv.is_type(CommunicationType::send_chat)
{
log_in!(
self.iota_id as i64,
PrintType::Iota,
"Routing upgraded Iota message to forward_message handler receiver_id={}",
receiver_id
);
self.handle_forward_message(cv).await;
return;
}
@ -149,11 +209,22 @@ impl IotaConnection {
|| cv.is_type(CommunicationType::delete_iota)
{
let sender = self.get_iota_id().await;
log_in!(
self.iota_id as i64,
PrintType::Iota,
"Routing upgraded Iota message to Omega forwarder with sender={}",
sender
);
self.handle_omega_forward(cv.with_sender(sender as u64))
.await;
return;
}
// Forward to client
log_in!(
self.iota_id as i64,
PrintType::Iota,
"Routing upgraded Iota message to client forwarder"
);
self.forward_to_client(cv).await;
}
@ -169,16 +240,45 @@ impl IotaConnection {
async fn handle_omega_forward(self: Arc<Self>, cv: CommunicationValue) {
let iota_for_closure = self.clone();
tokio::spawn(async move {
log_out!(
self.iota_id as i64,
PrintType::Iota,
"Forwarding upgraded Iota message to Omega type={:?} id={}",
cv.get_type(),
cv.get_id()
);
let response_cv = get_omega_connection()
.await_response(&cv.with_sender(self.iota_id), Some(Duration::from_secs(20)))
.await;
if let Ok(response_cv) = response_cv {
log_in!(
self.iota_id as i64,
PrintType::Iota,
"Received Omega response for upgraded Iota type={:?} id={}",
response_cv.get_type(),
response_cv.get_id()
);
iota_for_closure.send_message(&response_cv).await;
} else if let Err(e) = response_cv {
log_err!(
self.iota_id as i64,
PrintType::Iota,
"Omega forward failed for upgraded Iota connection: {}",
e
);
}
});
}
/// Handle ping message
async fn handle_ping(&self, cv: CommunicationValue) {
log_in!(
self.iota_id as i64,
PrintType::Iota,
"Handling ping/pong for upgraded Iota connection message_id={}",
cv.get_id()
);
if let DataValue::Number(last_ping) = cv.get_data(DataTypes::last_ping) {
if let Ok(ping_val) = last_ping.to_string().parse::<i64>() {
let mut ping_guard = self.ping.write().await;
@ -199,6 +299,14 @@ impl IotaConnection {
let response = CommunicationValue::new(CommunicationType::pong)
.with_id(cv.get_id())
.add_data(DataTypes::ping_clients, DataValue::Container(pings));
log_out!(
self.iota_id as i64,
PrintType::Iota,
"Sending pong from upgraded Iota connection message_id={}",
cv.get_id()
);
self.send_message(&response).await;
}
@ -207,16 +315,43 @@ impl IotaConnection {
let receiver_id = cv.get_receiver();
let sender_id = cv.get_sender();
log_in!(
self.iota_id as i64,
PrintType::Iota,
"Handling cross-routing message sender={} receiver={} type={:?}",
sender_id,
receiver_id,
cv.get_type()
);
if self.get_user_ids().await.contains(&(sender_id as u64)) {
if let Some(target_rho) = rho_manager::get_rho_con_for_user(receiver_id as i64).await {
log_out!(
self.iota_id as i64,
PrintType::Iota,
"Forwarding upgraded Iota message to target rho receiver={}",
receiver_id
);
target_rho.message_to_iota(cv).await;
} else {
log_err!(
self.iota_id as i64,
PrintType::Iota,
"No target rho found for receiver={}",
receiver_id
);
let error = CommunicationValue::new(CommunicationType::error_no_iota)
.with_id(cv.get_id())
.with_sender(cv.get_sender());
self.send_message(&error).await;
}
} else {
log_err!(
self.iota_id as i64,
PrintType::Iota,
"Unauthorized sender attempted cross-routing sender={}",
sender_id
);
self.send_message(
&CommunicationValue::new(CommunicationType::error_invalid_user_id).add_data(
DataTypes::error_type,
@ -334,19 +469,39 @@ impl IotaConnection {
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);
log_out!(
self.get_iota_id().await as i64,
PrintType::Iota,
"Forwarding upgraded Iota message to client type={:?} id={} receiver={}",
updated_cv.get_type(),
updated_cv.get_id(),
updated_cv.get_receiver()
);
rho_conn.message_to_client(updated_cv).await;
} else {
log_err!(
self.get_iota_id().await as i64,
PrintType::General,
"Failed to forward message to client"
"Failed to forward message to client because rho connection is missing"
);
}
}
pub async fn handle_close(&self) {
log_out!(
self.get_iota_id().await as i64,
PrintType::Iota,
"Handling upgraded Iota connection close"
);
if let Some(rho_conn) = self.get_rho_connection().await {
rho_conn.close_iota_connection().await;
} else {
log_err!(
self.get_iota_id().await as i64,
PrintType::Iota,
"No rho connection available during upgraded Iota close"
);
}
}