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

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