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

Client auth, Encryption module

Browse source
This commit is contained in:
Alex-Emmet 2026-01-18 21:42:48 +01:00
commit 7dd3428e00
10 changed files with 582 additions and 167 deletions
Download patch file Download diff file Expand all files Collapse all files

252
src/rho/client_connection.rs
View file

@ -1,7 +1,11 @@
use async_tungstenite::tungstenite::Message;
use async_tungstenite::{WebSocketReceiver, WebSocketSender};
use json::JsonValue;
use json::number::Number;
use rand::Rng;
use rand::distributions::Alphanumeric;
use std::sync::{Arc, Weak};
use std::time::Duration;
use tokio::sync::RwLock;
use tokio_util::compat::Compat;
use tungstenite::Utf8Bytes;
@ -10,6 +14,8 @@ use uuid::Uuid;
use super::{rho_connection::RhoConnection, rho_manager};
use crate::calls::call_manager;
use crate::omega::omega_connection::{WAITING_TASKS, get_omega_connection};
use crate::util::crypto_helper::{load_public_key, public_key_to_base64};
use crate::util::crypto_util::{DataFormat, SecurePayload};
use crate::util::logger::PrintType;
use crate::{
// calls::call_manager::CallManager,
@ -19,22 +25,18 @@ use crate::{
},
omega::omega_connection::OmegaConnection,
};
use crate::{log_in, log_out};
use crate::{get_private_key, get_public_key, log_in, log_out};
/// ClientConnection represents a WebSocket connection from a client device
pub struct ClientConnection {
/// WebSocket session
pub sender: Arc<RwLock<WebSocketSender<Compat<tokio::net::TcpStream>>>>,
pub receiver: Arc<RwLock<WebSocketReceiver<Compat<tokio::net::TcpStream>>>>,
/// User ID associated with this client
pub user_id: Arc<RwLock<i64>>,
/// Whether this connection has been identified/authenticated
pub identified: Arc<RwLock<bool>>,
/// Ping latency tracking
identified: Arc<RwLock<bool>>,
challenged: Arc<RwLock<bool>>,
challenge: Arc<RwLock<String>>,
pub ping: Arc<RwLock<i64>>,
/// Weak reference to RhoConnection to avoid circular references
pub rho_connection: Arc<RwLock<Option<Weak<RhoConnection>>>>,
/// List of user IDs this client is interested in receiving updates about
pub_key: Arc<RwLock<Option<Vec<u8>>>>,
pub rho_connection: Arc<RwLock<Option<Arc<RhoConnection>>>>,
pub interested_users: Arc<RwLock<Vec<i64>>>,
}
@ -49,7 +51,10 @@ impl ClientConnection {
receiver: Arc::new(RwLock::new(receiver)),
user_id: Arc::new(RwLock::new(0)),
identified: Arc::new(RwLock::new(false)),
challenged: Arc::new(RwLock::new(false)),
challenge: Arc::new(RwLock::new(String::new())),
ping: Arc::new(RwLock::new(-1)),
pub_key: Arc::new(RwLock::new(None)),
rho_connection: Arc::new(RwLock::new(None)),
interested_users: Arc::new(RwLock::new(Vec::new())),
})
@ -70,20 +75,9 @@ impl ClientConnection {
*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
}
self.rho_connection.read().await.clone()
}
/// Send a string message to the client
@ -109,14 +103,156 @@ impl ClientConnection {
pub async fn handle_message(self: Arc<Self>, message: Utf8Bytes) {
tokio::spawn(async move {
let cv = CommunicationValue::from_json(&message);
if cv.is_type(CommunicationType::ping) {
self.handle_ping(cv).await;
return;
}
log_in!(PrintType::Client, "{}", &cv.to_json().to_string());
let identified = *self.identified.read().await;
let challenged = *self.challenged.read().await;
// Handle identification
if cv.is_type(CommunicationType::identification) && !self.is_identified().await {
self.handle_identification(Arc::clone(&self), cv).await;
if !identified && cv.is_type(CommunicationType::identification) {
let user_id = cv
.get_data(DataTypes::user_id)
.and_then(|v| v.as_i64())
.unwrap_or(0);
if user_id == 0 {
log_out!(PrintType::Client, "Invalid USER ID");
self.send_error_response(&cv.get_id(), CommunicationType::error_invalid_data)
.await;
self.close().await;
return;
}
*self.user_id.write().await = user_id;
let get_pub_key_msg = CommunicationValue::new(CommunicationType::get_user_data)
.with_id(cv.get_id())
.add_data(DataTypes::user_id, JsonValue::from(user_id));
let response_cv = get_omega_connection()
.await_response(&get_pub_key_msg, Some(Duration::from_secs(20)))
.await;
if let Ok(response_cv) = response_cv {
if !response_cv.is_type(CommunicationType::get_user_data) {
self.send_error_response(&cv.get_id(), CommunicationType::error_internal)
.await;
self.close().await;
return;
}
let base64_pub = response_cv
.get_data(DataTypes::public_key)
.and_then(|v| v.as_str())
.unwrap_or("");
let pub_key = match load_public_key(base64_pub) {
Some(pk) => pk,
None => {
self.send_error_response(
&cv.get_id(),
CommunicationType::error_invalid_public_key,
)
.await;
self.close().await;
return;
}
};
*self.pub_key.write().await = Some(pub_key.as_bytes().to_vec());
let challenge: String = rand::thread_rng()
.sample_iter(&Alphanumeric)
.take(32)
.map(char::from)
.collect();
*self.challenge.write().await = challenge.clone();
let encrypted_challenge =
SecurePayload::new(challenge, DataFormat::Raw, get_private_key())
.unwrap()
.encrypt_x448(pub_key)
.unwrap()
.export(DataFormat::Base64);
*self.identified.write().await = true;
let challenge_msg = CommunicationValue::new(CommunicationType::challenge)
.with_id(cv.get_id())
.add_data_str(
DataTypes::public_key,
public_key_to_base64(&get_public_key()),
)
.add_data_str(DataTypes::challenge, encrypted_challenge);
self.send_message(&challenge_msg).await;
} else {
self.send_error_response(&cv.get_id(), CommunicationType::error_internal)
.await;
self.close().await;
return;
}
return;
}
if identified && !challenged && cv.is_type(CommunicationType::challenge_response) {
let client_response = cv
.get_data(DataTypes::challenge)
.and_then(|v| v.as_str())
.unwrap_or("");
if client_response == *self.challenge.read().await {
*self.challenged.write().await = true;
let user_id = self.get_user_id().await;
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;
*user_id_guard = user_id;
}
{
let mut identified_guard = self.identified.write().await;
*identified_guard = true;
}
*self.rho_connection.write().await = Some(Arc::clone(&rho_connection));
let response =
CommunicationValue::new(CommunicationType::identification_response)
.with_id(cv.get_id());
self.send_message(&response).await;
} else {
self.send_error_response(
&cv.get_id(),
CommunicationType::error_not_authenticated,
)
.await;
self.close().await;
return;
}
return;
}
if !self.is_identified().await {
self.send_error_response(&cv.get_id(), CommunicationType::error_not_authenticated)
.await;
self.close().await;
return;
}
@ -174,71 +310,6 @@ impl ClientConnection {
.await;
}
/// Handle identification message
async fn handle_identification(&self, sarc: Arc<ClientConnection>, cv: CommunicationValue) {
// Extract user ID
let user_id: i64 = match cv.get_data(DataTypes::user_id) {
Some(id_str) => id_str.as_i64().unwrap_or(0),
None => {
self.send_error_response(&cv.get_id(), CommunicationType::error_invalid_user_id)
.await;
return;
}
};
// 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 = true; // NO VALIDATION,
// SWAP TO AUTH VIA CHALLENGE
// auth_connector::is_private_key_valid(user_id, &private_key_hash.to_string()).await;
if !is_valid {
println!("Invalid private key");
self.send_error_response(
&cv.get_id(),
CommunicationType::error_invalid_private_key,
)
.await;
return;
}
} else {
log_in!(PrintType::Client, "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;
*user_id_guard = user_id;
}
{
let mut identified_guard = self.identified.write().await;
*identified_guard = true;
}
self.set_rho_connection(Arc::downgrade(&rho_connection))
.await;
rho_connection.add_client_connection(Arc::from(sarc)).await;
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, cv: CommunicationValue) {
// Update our ping if provided
@ -451,7 +522,10 @@ impl Clone for ClientConnection {
receiver: Arc::clone(&self.receiver),
user_id: Arc::clone(&self.user_id),
identified: Arc::clone(&self.identified),
challenged: Arc::clone(&self.challenged),
challenge: Arc::clone(&self.challenge),
ping: Arc::clone(&self.ping),
pub_key: Arc::clone(&self.pub_key),
rho_connection: Arc::clone(&self.rho_connection),
interested_users: Arc::clone(&self.interested_users),
}

120
src/rho/iota_connection.rs Normal file → Executable file
View file

@ -10,6 +10,8 @@ use crate::omega::omega_connection::get_omega_connection;
use crate::util::crypto_helper::encrypt;
use crate::util::crypto_helper::load_public_key;
use crate::util::crypto_helper::public_key_to_base64;
use crate::util::crypto_util::DataFormat;
use crate::util::crypto_util::SecurePayload;
use crate::util::logger::PrintType;
use async_tungstenite::WebSocketReceiver;
use async_tungstenite::WebSocketSender;
@ -24,9 +26,11 @@ use std::{
time::Duration,
};
use tokio::sync::RwLock;
use tokio::sync::mpsc;
use tokio_util::compat::Compat;
use tungstenite::Utf8Bytes;
use uuid::Uuid;
use warp::filters::method::get;
use x448::PublicKey;
use super::{rho_connection::RhoConnection, rho_manager};
@ -144,6 +148,8 @@ impl IotaConnection {
return;
}
log_in!(PrintType::Iota, "{}", cv.to_json().to_string());
let identified = *self.identified.read().await;
let challenged = *self.challenged.read().await;
@ -153,27 +159,14 @@ impl IotaConnection {
.and_then(|v| v.as_i64())
.unwrap_or(0);
if iota_id == 0 {
log_out!(PrintType::Iota, "Invalid IOTA ID");
self.send_error_response(&cv.get_id(), CommunicationType::error_invalid_data)
.await;
self.close().await;
return;
}
let user_ids_json = cv
.get_data(DataTypes::user_ids)
.unwrap_or(&JsonValue::Null)
.clone();
let mut user_ids = Vec::new();
if let JsonValue::Array(ids) = user_ids_json {
for id_val in ids {
if let Some(id) = id_val.as_i64() {
user_ids.push(id);
}
}
}
*self.iota_id.write().await = iota_id;
*self.user_ids.write().await = user_ids;
let get_pub_key_msg = CommunicationValue::new(CommunicationType::get_iota_data)
.with_id(cv.get_id())
@ -220,7 +213,11 @@ impl IotaConnection {
*self.challenge.write().await = challenge.clone();
let encrypted_challenge =
encrypt(get_private_key(), pub_key, &challenge).unwrap_or_default();
SecurePayload::new(&challenge, DataFormat::Base64, get_private_key())
.unwrap()
.encrypt_x448(pub_key)
.unwrap()
.export(DataFormat::Base64);
*self.identified.write().await = true;
@ -315,12 +312,6 @@ impl IotaConnection {
*self.challenged.write().await = true;
let iota_id = self.get_iota_id().await;
let user_ids = self.get_user_ids().await;
let mut validated_user_ids: Vec<i64> = Vec::new();
for user_id in user_ids {
validated_user_ids.push(user_id);
}
if rho_manager::contains_iota(iota_id).await {
if let Some(existing_rho) = rho_manager::get_rho_by_iota(iota_id).await {
@ -328,8 +319,47 @@ impl IotaConnection {
}
}
// Inform Omega & Verify Users
let iota_users_cv = get_omega_connection()
.await_response(
&CommunicationValue::new(CommunicationType::iota_connected).add_data(
DataTypes::iota_id,
JsonValue::from(self.get_iota_id().await),
),
Some(Duration::from_secs(20)),
)
.await;
let mut user_ids: Vec<i64> = Vec::new();
if let Ok(iota_users_cv) = iota_users_cv {
if !iota_users_cv.is_type(CommunicationType::iota_user_data) {
log_err!(
PrintType::Omikron,
"Invalid communication type {:?}",
iota_users_cv.get_type()
);
return;
}
let val_user_ids = iota_users_cv.get_data(DataTypes::user_ids).unwrap().clone();
match val_user_ids {
JsonValue::Array(arr) => {
for item in arr {
if let JsonValue::Number(_) = item {
user_ids.push(item.as_i64().unwrap_or(0));
}
}
}
_ => {}
}
} else {
log_err!(PrintType::Omikron, "Failed to retrieve user IDs");
}
log_in!(PrintType::General, "User IDs: {:?}", user_ids.clone());
*self.user_ids.write().await = user_ids.clone();
let rho_connection =
Arc::new(RhoConnection::new(self.clone(), validated_user_ids.clone()).await);
Arc::new(RhoConnection::new(self.clone(), user_ids.clone()).await);
self.set_rho_connection(Arc::downgrade(&rho_connection))
.await;
@ -337,7 +367,7 @@ impl IotaConnection {
rho_manager::add_rho(rho_connection).await;
let mut str = String::new();
for id in &validated_user_ids {
for id in &user_ids {
str.push_str(&format!(",{}", id));
}
if !str.is_empty() {
@ -348,7 +378,7 @@ impl IotaConnection {
&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()),
.add_data_str(DataTypes::accepted, user_ids.len().to_string()),
)
.await;
} else {
@ -557,6 +587,48 @@ impl IotaConnection {
}
}
}
pub async fn await_response(
&self,
cv: &CommunicationValue,
timeout_duration: Option<Duration>,
) -> Result<CommunicationValue, String> {
let (tx, mut rx) = mpsc::channel(1);
let msg_id = cv.get_id();
let task_tx = tx.clone();
self.waiting_tasks.insert(
msg_id,
Box::new(move |_, response_cv| {
let inner_tx = task_tx.clone();
tokio::spawn(async move {
if let Err(e) = inner_tx.send(response_cv).await {
log_err!(
PrintType::Iota,
"Failed to send response back to awaiter: {}",
e
);
}
});
true
}),
);
self.send_message(cv).await;
let timeout = timeout_duration.unwrap_or(Duration::from_secs(10));
match tokio::time::timeout(timeout, rx.recv()).await {
Ok(Some(response_cv)) => Ok(response_cv),
Ok(None) => Err("Failed to receive response, channel was closed.".to_string()),
Err(_) => {
self.waiting_tasks.remove(&msg_id);
Err(format!(
"Request timed out after {} seconds.",
timeout.as_secs()
))
}
}
}
}
impl std::fmt::Debug for IotaConnection {

5
src/rho/rho_connection.rs
View file

@ -24,9 +24,6 @@ impl RhoConnection {
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).await;
rho_connection
}
@ -126,7 +123,7 @@ impl RhoConnection {
pub async fn message_to_client(&self, cv: CommunicationValue) {
let connections = self.client_connections.read().await;
for connection in connections.iter() {
if connection.get_user_id().await == cv.receiver {
if connection.get_user_id().await == cv.get_receiver() {
connection.send_message(&cv).await;
}
}