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

Connection With Omikron

Browse source
This commit is contained in:
Alex Emmet 2026-01-05 02:15:55 +01:00
commit ebd5822dd6
10 changed files with 426 additions and 217 deletions
Download patch file Download diff file Expand all files Collapse all files

237
src/server/omikron_connection.rs
View file

@ -1,8 +1,9 @@
use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes};
use crate::get_public_key;
use crate::sql::sql::get_omikron_by_id;
use crate::util::crypto_helper::encrypt;
use crate::{get_private_key, log_in_from, log_out_from};
use crate::util::logger::PrintType;
use crate::{get_private_key, log_out};
use crate::{get_public_key, log_in};
use base64::{Engine as _, engine::general_purpose::STANDARD};
use dashmap::DashMap;
use futures::SinkExt;
@ -53,10 +54,17 @@ impl OmikronConnection {
waiting_tasks: DashMap::new(),
})
}
pub async fn send_message(&self, message: &CommunicationValue) {
pub async fn send_message(&self, cv: &CommunicationValue) {
let mut sender = self.sender.write().await;
let message_text = Message::Text(Utf8Bytes::from(message.to_json().to_string()));
log_out_from!(*self.omikron_id.read().await, "{}", message_text);
let message_text = Message::Text(Utf8Bytes::from(cv.to_json().to_string()));
if !cv.is_type(CommunicationType::pong) {
log_out!(
*self.omikron_id.read().await,
PrintType::Omikron,
"{}",
message_text
);
}
sender.send(message_text).await.unwrap();
}
pub async fn get_user_id(&self) -> i64 {
@ -76,155 +84,124 @@ impl OmikronConnection {
return;
}
log_in_from!(*self.omikron_id.read().await, "{}", message);
log_in!(
*self.omikron_id.read().await,
PrintType::Omikron,
"{}",
message
);
if let Some((_, task)) = self.waiting_tasks.remove(&cv.get_id()) {
let _ = task(self.clone(), cv.clone());
return;
}
// Handle identification
if !*self.identified.read().await && cv.is_type(CommunicationType::identification) {
let omikron_id = cv
.get_data(DataTypes::omikron)
.unwrap_or(&JsonValue::Null)
.as_i64()
.unwrap_or(0);
if let Ok((_, public_key, _)) = get_omikron_by_id(omikron_id).await {
// Generate Challenge, encrypt it and send it to the omikron
*self.omikron_id.write().await = omikron_id;
*self.identified.write().await = true;
match get_omikron_by_id(omikron_id).await {
Ok((public_key, _)) => {
// Generate Challenge, encrypt it and send it to the omikron
*self.omikron_id.write().await = omikron_id;
let challenge_str: String = rand::thread_rng()
.sample_iter(&Alphanumeric)
.take(32)
.map(char::from)
.collect();
let challenge_str: String = rand::thread_rng()
.sample_iter(&Alphanumeric)
.take(32)
.map(char::from)
.collect();
*self.challenge.write().await = challenge_str.clone();
*self.challenge.write().await = challenge_str.clone();
let user_public_key_bytes = match STANDARD.decode(&public_key) {
Ok(bytes) => bytes,
Err(_) => {
self.send_error_response(
&cv.get_id(),
CommunicationType::error_invalid_user_id,
let user_public_key_bytes = match STANDARD.decode(&public_key) {
Ok(bytes) => bytes,
Err(_) => {
self.send_error_response(
&cv.get_id(),
CommunicationType::error_invalid_omikron_id,
)
.await;
return;
}
};
*self.pub_key.write().await = Some(user_public_key_bytes.clone());
let omikron_pub_key: PublicKey =
match PublicKey::from_bytes(&user_public_key_bytes) {
Some(key) => key,
None => {
self.send_error_response(
&cv.get_id(),
CommunicationType::error_invalid_public_key,
)
.await;
return;
}
};
let encrypted_challenge =
encrypt(get_private_key(), omikron_pub_key, &challenge_str)
.unwrap_or("".to_string());
let response = CommunicationValue::new(CommunicationType::challenge)
.add_data_str(
DataTypes::public_key,
STANDARD.encode(get_public_key().as_bytes()),
)
.await;
return;
}
};
*self.pub_key.write().await = Some(user_public_key_bytes.clone());
.add_data_str(DataTypes::challenge, encrypted_challenge)
.with_id(cv.get_id());
let omikron_pub_key: PublicKey = match PublicKey::from_bytes(&user_public_key_bytes)
{
Some(key) => key,
None => {
self.send_error_response(
&cv.get_id(),
CommunicationType::error_invalid_user_id,
)
.await;
return;
}
};
let encrypted_challenge =
encrypt(get_private_key(), omikron_pub_key, &challenge_str)
.unwrap_or("".to_string());
let response = CommunicationValue::new(CommunicationType::challenge)
.add_data_str(
DataTypes::public_key,
STANDARD.encode(get_public_key().as_bytes()),
self.send_message(&response).await;
*self.identified.write().await = true;
return;
}
Err(e) => {
self.send_message(
&CommunicationValue::new(CommunicationType::error_not_authenticated)
.with_id(cv.get_id())
.add_data_str(DataTypes::error_type, e.to_string()),
)
.add_data_str(DataTypes::challenge, encrypted_challenge)
.await;
return;
}
}
}
// Handle challenge response
if *self.identified.read().await
&& !*self.challenged.read().await
&& cv.is_type(CommunicationType::challenge_response)
{
let client_response = cv
.get_data(DataTypes::challenge)
.unwrap_or(&JsonValue::Null)
.as_str()
.unwrap_or("");
let expected_challenge = self.challenge.read().await.clone();
if client_response == expected_challenge {
*self.challenged.write().await = true;
let response = CommunicationValue::new(CommunicationType::identification_response)
.with_id(cv.get_id());
self.send_message(&response).await;
// prepare Challenge Response handling
self.waiting_tasks.insert(
cv.get_id(),
Box::new(
|selfc: Arc<OmikronConnection>, cv: CommunicationValue| -> bool {
tokio::spawn(async move {
let client_challenge_response_b64 =
match cv.get_data(DataTypes::challenge) {
Some(data) => data.to_string(),
None => {
selfc
.send_error_response(
&cv.get_id(),
CommunicationType::error,
)
.await;
return;
}
};
let challenge_response_bytes =
match STANDARD.decode(&client_challenge_response_b64) {
Ok(bytes) => bytes,
Err(_) => {
selfc
.send_error_response(
&cv.get_id(),
CommunicationType::error,
)
.await;
return;
}
};
if challenge_response_bytes.len() < 12 {
selfc
.send_error_response(&cv.get_id(), CommunicationType::error)
.await;
return;
}
let client_response = cv
.get_data(DataTypes::challenge)
.unwrap_or(&JsonValue::Null)
.as_str()
.unwrap_or("");
let expected_challenge = selfc.challenge.read().await.clone();
if client_response != expected_challenge {
selfc
.send_error_response(
&cv.get_id(),
CommunicationType::error_invalid_challenge,
)
.await;
selfc.close().await;
return;
}
*selfc.challenged.write().await = true;
let response = CommunicationValue::new(
CommunicationType::identification_response,
)
.with_id(cv.get_id());
selfc.send_message(&response).await;
return;
});
return true;
},
),
);
} else {
self.send_error_response(&cv.get_id(), CommunicationType::error_not_found)
self.send_error_response(&cv.get_id(), CommunicationType::error_invalid_challenge)
.await;
self.close().await;
}
return;
}
if self.waiting_tasks.contains_key(&cv.get_id()) {
let (_, task) = self.waiting_tasks.remove(&cv.get_id()).unwrap();
let _ = task(self.clone(), cv.clone());
}
if !self.is_identified().await {
self.send_error_response(&cv.get_id(), CommunicationType::error_not_found)
self.send_error_response(&cv.get_id(), CommunicationType::error_not_authenticated)
.await;
self.close().await;
return;
}
}