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

omikron connection handling

Browse source
This commit is contained in:
Alex Emmet 2026-01-10 23:40:36 +01:00
commit 84841601fb
1 changed files with 76 additions and 79 deletions
Download patch file Download diff file Expand all files Collapse all files

117
src/server/omikron_connection.rs
View file

@ -4,8 +4,7 @@ use crate::sql::sql::{self, get_by_user_id, get_by_username, get_iota_by_id, get
use crate::sql::user_online_tracker::{self}; use crate::sql::user_online_tracker::{self};
use crate::util::crypto_helper::encrypt; use crate::util::crypto_helper::encrypt;
use crate::util::logger::PrintType; use crate::util::logger::PrintType;
use crate::{get_private_key, log_out}; use crate::{get_private_key, get_public_key, log_in, log_out};
use crate::{get_public_key, log_in};
use base64::{Engine as _, engine::general_purpose::STANDARD}; use base64::{Engine as _, engine::general_purpose::STANDARD};
use dashmap::DashMap; use dashmap::DashMap;
use futures::SinkExt; use futures::SinkExt;
@ -99,30 +98,31 @@ impl OmikronConnection {
return; return;
} }
// If not yet identified
if !self.is_identified().await { if !self.is_identified().await {
// handle identification let identified = *self.identified.read().await;
if !*self.identified.read().await && cv.is_type(CommunicationType::identification) { let challenged = *self.challenged.read().await;
if !identified && cv.is_type(CommunicationType::identification) {
let omikron_id = cv let omikron_id = cv
.get_data(DataTypes::omikron) .get_data(DataTypes::omikron)
.unwrap_or(&JsonValue::Null) .and_then(|v| v.as_i64())
.as_i64()
.unwrap_or(0); .unwrap_or(0);
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() let (public_key, _) = match get_omikron_by_id(omikron_id).await {
.sample_iter(&Alphanumeric) Ok(v) => v,
.take(32) Err(e) => {
.map(char::from) self.send_message(
.collect(); &CommunicationValue::new(CommunicationType::error_not_authenticated)
.with_id(cv.get_id())
.add_data_str(DataTypes::error_type, e.to_string()),
)
.await;
return;
}
};
*self.challenge.write().await = challenge_str.clone(); let pub_key_bytes = match STANDARD.decode(&public_key) {
Ok(b) => b,
let user_public_key_bytes = match STANDARD.decode(&public_key) {
Ok(bytes) => bytes,
Err(_) => { Err(_) => {
self.send_error_response( self.send_error_response(
&cv.get_id(), &cv.get_id(),
@ -132,12 +132,10 @@ impl OmikronConnection {
return; return;
} }
}; };
*self.pub_key.write().await = Some(user_public_key_bytes.clone());
let omikron_pub_key: PublicKey = let omikron_pub_key = match PublicKey::from_bytes(&pub_key_bytes) {
match PublicKey::from_bytes(&user_public_key_bytes) { Some(k) => k,
Some(key) => key, None => {
_ => {
self.send_error_response( self.send_error_response(
&cv.get_id(), &cv.get_id(),
CommunicationType::error_invalid_public_key, CommunicationType::error_invalid_public_key,
@ -147,55 +145,50 @@ impl OmikronConnection {
} }
}; };
let encrypted_challenge = let challenge: String = rand::thread_rng()
encrypt(get_private_key(), omikron_pub_key, &challenge_str) .sample_iter(&Alphanumeric)
.unwrap_or("".to_string()); .take(32)
.map(char::from)
.collect();
*self.omikron_id.write().await = omikron_id;
*self.challenge.write().await = challenge.clone();
*self.pub_key.write().await = Some(pub_key_bytes);
*self.identified.write().await = true;
let encrypted =
encrypt(get_private_key(), omikron_pub_key, &challenge).unwrap_or_default();
let response = CommunicationValue::new(CommunicationType::challenge) let response = CommunicationValue::new(CommunicationType::challenge)
.with_id(cv.get_id())
.add_data_str( .add_data_str(
DataTypes::public_key, DataTypes::public_key,
STANDARD.encode(get_public_key().as_bytes()), STANDARD.encode(get_public_key().as_bytes()),
) )
.add_data_str(DataTypes::challenge, encrypted_challenge) .add_data_str(DataTypes::challenge, encrypted);
.with_id(cv.get_id());
self.send_message(&response).await; self.send_message(&response).await;
*self.identified.write().await = true;
return; 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()),
)
.await;
return; // ──────────────────────────────
} // Challenge response
} // ──────────────────────────────
} if identified && !challenged && cv.is_type(CommunicationType::challenge_response) {
// Handle challenge response
if *self.identified.read().await
&& !*self.challenged.read().await
&& cv.is_type(CommunicationType::challenge_response)
{
let client_response = cv let client_response = cv
.get_data(DataTypes::challenge) .get_data(DataTypes::challenge)
.unwrap_or(&JsonValue::Null) .and_then(|v| v.as_str())
.as_str()
.unwrap_or(""); .unwrap_or("");
let expected_challenge = self.challenge.read().await.clone();
if client_response == expected_challenge { if client_response == *self.challenge.read().await {
*self.challenged.write().await = true; *self.challenged.write().await = true;
let response =
CommunicationValue::new(CommunicationType::identification_response)
.with_id(cv.get_id());
let _ = sql::set_omikron_active(self.get_omikron_id().await, true); let _ = sql::set_omikron_active(self.get_omikron_id().await, true);
self.send_message(&response).await;
self.send_message(
&CommunicationValue::new(CommunicationType::identification_response)
.with_id(cv.get_id()),
)
.await;
} else { } else {
self.send_error_response( self.send_error_response(
&cv.get_id(), &cv.get_id(),
@ -207,14 +200,13 @@ impl OmikronConnection {
return; return;
} }
// if not identified && not identifying
self.send_error_response(&cv.get_id(), CommunicationType::error_not_authenticated) self.send_error_response(&cv.get_id(), CommunicationType::error_not_authenticated)
.await; .await;
self.close().await; self.close().await;
return;
} }
let omikron_id = self.get_omikron_id().await; let omikron_id = self.get_omikron_id().await;
// ONLINE STATUS TRACKING
if cv.is_type(CommunicationType::user_connected) { if cv.is_type(CommunicationType::user_connected) {
if let Some(user_id) = cv.get_data(DataTypes::user_id).and_then(|v| v.as_i64()) { if let Some(user_id) = cv.get_data(DataTypes::user_id).and_then(|v| v.as_i64()) {
user_online_tracker::track_user_status(user_id, ConnectionType::Online, omikron_id) user_online_tracker::track_user_status(user_id, ConnectionType::Online, omikron_id)
@ -290,6 +282,8 @@ impl OmikronConnection {
} }
return; return;
} }
// DATA RETURN
if cv.is_type(CommunicationType::get_user_data) { if cv.is_type(CommunicationType::get_user_data) {
if let Some(user_id) = cv.get_data(DataTypes::user_id).cloned() { if let Some(user_id) = cv.get_data(DataTypes::user_id).cloned() {
if let Some(user_id) = user_id.as_i64() { if let Some(user_id) = user_id.as_i64() {
@ -563,6 +557,7 @@ impl OmikronConnection {
return; return;
} }
// REGISTERING
if cv.is_type(CommunicationType::get_register) { if cv.is_type(CommunicationType::get_register) {
let register_id = sql::get_register_id().await; let register_id = sql::get_register_id().await;
let response = CommunicationValue::new(CommunicationType::get_register) let response = CommunicationValue::new(CommunicationType::get_register)
@ -645,6 +640,8 @@ impl OmikronConnection {
} }
return; return;
} }
// CHANGING DATA
if cv.is_type(CommunicationType::change_user_data) { if cv.is_type(CommunicationType::change_user_data) {
if let Some(user_id) = cv.get_data(DataTypes::user_id).and_then(|v| v.as_i64()) { if let Some(user_id) = cv.get_data(DataTypes::user_id).and_then(|v| v.as_i64()) {
let mut success = true; let mut success = true;