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