1231 lines
46 KiB
Rust
1231 lines
46 KiB
Rust
use crate::{
|
|
get_private_key, get_public_key, log, log_cv_in, log_cv_out, log_err, log_in,
|
|
server::short_link::add_short_link,
|
|
sql::{
|
|
connection_status::UserStatus,
|
|
sql::{self, get_by_user_id, get_by_username, get_iota_by_id, get_omikron_by_id},
|
|
user_online_tracker::{self},
|
|
},
|
|
transport::omikron_manager,
|
|
util::{crypto_helper::encrypt, file_util::load_file_vec, logger::PrintType},
|
|
};
|
|
use base64::{Engine as _, engine::general_purpose::STANDARD};
|
|
use dashmap::DashMap;
|
|
use rand::{Rng, distributions::Alphanumeric};
|
|
use std::{
|
|
sync::Arc,
|
|
time::{Duration, Instant},
|
|
};
|
|
use tokio::{
|
|
sync::{Mutex, RwLock},
|
|
time::interval,
|
|
};
|
|
use ttp_core::{CommunicationType, CommunicationValue, DataTypes, DataValue};
|
|
use ttp_native::{Host, Policy, Receiver, SendMode, Sender};
|
|
use x448::PublicKey;
|
|
|
|
// ============================================================================
|
|
// Configuration
|
|
// ============================================================================
|
|
|
|
const CLEANUP_INTERVAL: Duration = Duration::from_secs(30);
|
|
const MAX_WAITING_AGE: Duration = Duration::from_secs(60);
|
|
|
|
// ============================================================================
|
|
// Error Types
|
|
// ============================================================================
|
|
|
|
#[derive(Debug, thiserror::Error)]
|
|
pub enum OmikronError {
|
|
#[error("Not connected")]
|
|
NotConnected,
|
|
#[error("Not authenticated")]
|
|
NotAuthenticated,
|
|
#[error("Invalid response")]
|
|
InvalidResponse,
|
|
#[error("Authentication failed")]
|
|
AuthenticationFailed,
|
|
#[error("SQL error: {0}")]
|
|
Sql(String),
|
|
#[error("Send error: {0}")]
|
|
Send(String),
|
|
}
|
|
|
|
pub type OmikronResult<T> = Result<T, OmikronError>;
|
|
|
|
// ============================================================================
|
|
// Waiting Task System (Preserved from original)
|
|
// ============================================================================
|
|
|
|
pub struct WaitingTask {
|
|
pub task: Box<dyn Fn(Arc<OmikronConnection>, CommunicationValue) -> bool + Send + Sync>,
|
|
pub inserted_at: Instant,
|
|
}
|
|
|
|
// ============================================================================
|
|
// Connection State
|
|
// ============================================================================
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
enum AuthState {
|
|
Unauthenticated,
|
|
Identified { omikron_id: i64 },
|
|
Authenticated { omikron_id: i64 },
|
|
}
|
|
|
|
impl AuthState {
|
|
fn is_authenticated(&self) -> bool {
|
|
match self {
|
|
AuthState::Authenticated { omikron_id } => true,
|
|
_ => false,
|
|
}
|
|
}
|
|
|
|
fn omikron_id(&self) -> Option<i64> {
|
|
match self {
|
|
AuthState::Identified { omikron_id } | AuthState::Authenticated { omikron_id } => {
|
|
Some(*omikron_id)
|
|
}
|
|
_ => None,
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Omikron Connection (ttp/QUIC-based)
|
|
// ============================================================================
|
|
|
|
pub struct OmikronConnection {
|
|
id: u64,
|
|
sender: Mutex<Option<Sender>>,
|
|
state: RwLock<AuthState>,
|
|
challenge: RwLock<String>,
|
|
pub_key: RwLock<Option<Vec<u8>>>,
|
|
pub ping: RwLock<i64>,
|
|
waiting_tasks: DashMap<u32, WaitingTask>,
|
|
cleanup_handle: std::sync::Mutex<Option<tokio::task::JoinHandle<()>>>,
|
|
}
|
|
|
|
impl Drop for OmikronConnection {
|
|
fn drop(&mut self) {
|
|
if let Some(handle) = self.cleanup_handle.lock().unwrap().take() {
|
|
handle.abort();
|
|
}
|
|
}
|
|
}
|
|
|
|
impl OmikronConnection {
|
|
// -------------------------------------------------------------------------
|
|
// Construction
|
|
// -------------------------------------------------------------------------
|
|
|
|
pub fn new(sender: Sender) -> Arc<Self> {
|
|
let conn = Arc::new(Self {
|
|
id: rand::random(),
|
|
sender: Mutex::new(Some(sender)),
|
|
state: RwLock::new(AuthState::Unauthenticated),
|
|
challenge: RwLock::new(String::new()),
|
|
pub_key: RwLock::new(None),
|
|
ping: RwLock::new(-1),
|
|
waiting_tasks: DashMap::new(),
|
|
cleanup_handle: std::sync::Mutex::new(None),
|
|
});
|
|
|
|
conn
|
|
}
|
|
|
|
// -------------------------------------------------------------------------
|
|
// Main Handler Loop
|
|
// -------------------------------------------------------------------------
|
|
|
|
pub async fn handle(self: Arc<Self>, receiver: &mut Receiver) {
|
|
log_in!(
|
|
self.id as i64,
|
|
PrintType::Omega,
|
|
"Omikron connection started"
|
|
);
|
|
|
|
// Start cleanup task
|
|
let cleanup_conn = self.clone();
|
|
let cleanup_handle = tokio::spawn(async move {
|
|
let mut ticker = interval(CLEANUP_INTERVAL);
|
|
loop {
|
|
ticker.tick().await;
|
|
cleanup_conn
|
|
.waiting_tasks
|
|
.retain(|_, v| v.inserted_at.elapsed() < MAX_WAITING_AGE);
|
|
}
|
|
});
|
|
*self.cleanup_handle.lock().unwrap() = Some(cleanup_handle);
|
|
|
|
while let Ok(cv) = receiver.receive().await {
|
|
if let Err(e) = self.clone().process_message(cv).await {
|
|
log_err!(0, PrintType::Omega, "Error processing message: {}", e);
|
|
if matches!(e, OmikronError::NotConnected) {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
self.clone().cleanup().await;
|
|
log_in!(
|
|
self.id as i64,
|
|
PrintType::Omega,
|
|
"Omikron connection closed"
|
|
);
|
|
}
|
|
|
|
// -------------------------------------------------------------------------
|
|
// Message Processing
|
|
// -------------------------------------------------------------------------
|
|
|
|
async fn process_message(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
|
|
if !cv.is_type(CommunicationType::pong) && !cv.is_type(CommunicationType::ping) {
|
|
log_cv_in!(PrintType::Omikron, &cv);
|
|
}
|
|
|
|
let msg_id = cv.get_id();
|
|
|
|
// Check waiting tasks first (response to previous request)
|
|
if let Some((_, task)) = self.waiting_tasks.remove(&msg_id) {
|
|
let _ = (task.task)(self.clone(), cv);
|
|
return Ok(());
|
|
}
|
|
|
|
// Handle ping regardless of auth state
|
|
if cv.is_type(CommunicationType::ping) {
|
|
return self.handle_ping(cv).await;
|
|
}
|
|
|
|
let current_state = *self.state.read().await;
|
|
// Route based on authentication state
|
|
match current_state {
|
|
AuthState::Unauthenticated => self.clone().handle_unauthenticated(cv).await,
|
|
AuthState::Identified { .. } => self.clone().handle_identified(cv).await,
|
|
AuthState::Authenticated { omikron_id } => {
|
|
self.clone().handle_authenticated(cv, omikron_id).await
|
|
}
|
|
}
|
|
}
|
|
|
|
// -------------------------------------------------------------------------
|
|
// Authentication Handlers
|
|
// -------------------------------------------------------------------------
|
|
|
|
async fn handle_unauthenticated(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
|
|
if !cv.is_type(CommunicationType::identification) {
|
|
let _ = self
|
|
.send_error_response(cv.get_id(), CommunicationType::error_not_authenticated)
|
|
.await;
|
|
return Err(OmikronError::NotAuthenticated);
|
|
}
|
|
|
|
// Extract omikron ID
|
|
let omikron_id = cv
|
|
.get_data(DataTypes::omikron_id)
|
|
.as_number()
|
|
.ok_or(OmikronError::InvalidResponse)?;
|
|
log!("Omikron {:?} connected", omikron_id);
|
|
|
|
// Lookup omikron in database
|
|
let (public_key, _) = get_omikron_by_id(omikron_id)
|
|
.await
|
|
.map_err(|e| OmikronError::Sql(e.to_string()))?;
|
|
|
|
log!("Got public Key");
|
|
|
|
let pub_key_bytes = STANDARD
|
|
.decode(&public_key)
|
|
.map_err(|_| OmikronError::AuthenticationFailed)?;
|
|
|
|
let pub_key_bytes_clone = pub_key_bytes.clone();
|
|
let omikron_pub_key = PublicKey::from_bytes(&pub_key_bytes_clone)
|
|
.ok_or(OmikronError::AuthenticationFailed)?;
|
|
|
|
log!("Decoded public Key");
|
|
|
|
// Generate challenge
|
|
let challenge: String = rand::thread_rng()
|
|
.sample_iter(&Alphanumeric)
|
|
.take(32)
|
|
.map(char::from)
|
|
.collect();
|
|
|
|
log!("Generated Challenge");
|
|
|
|
*self.challenge.write().await = challenge.clone();
|
|
|
|
log!("Stored Challenge");
|
|
|
|
*self.pub_key.write().await = Some(pub_key_bytes);
|
|
|
|
log!("Stored Pubkey");
|
|
|
|
*self.state.write().await = AuthState::Identified { omikron_id };
|
|
|
|
log!("Stored State");
|
|
|
|
let challenge_clone = challenge.clone();
|
|
let private_key = get_private_key();
|
|
let public_key_for_encrypt = omikron_pub_key;
|
|
|
|
let encrypted = tokio::task::spawn_blocking(move || {
|
|
encrypt(private_key, public_key_for_encrypt, &challenge_clone)
|
|
.map_err(|_| OmikronError::AuthenticationFailed)
|
|
})
|
|
.await
|
|
.map_err(|_| OmikronError::AuthenticationFailed)??;
|
|
|
|
log!("Encrypted Challenge");
|
|
|
|
// Send challenge response
|
|
let response = CommunicationValue::new(CommunicationType::challenge)
|
|
.with_id(cv.get_id())
|
|
.add_data(
|
|
DataTypes::public_key,
|
|
DataValue::Str(STANDARD.encode(get_public_key().as_bytes())),
|
|
)
|
|
.add_data(DataTypes::challenge, DataValue::Str(encrypted));
|
|
|
|
log!("Sending Challenge");
|
|
self.send(&response).await
|
|
}
|
|
|
|
async fn handle_identified(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
|
|
if !cv.is_type(CommunicationType::challenge_response) {
|
|
let _ = self
|
|
.send_error_response(cv.get_id(), CommunicationType::error_not_authenticated)
|
|
.await;
|
|
return Err(OmikronError::NotAuthenticated);
|
|
}
|
|
|
|
let client_response = cv
|
|
.get_data(DataTypes::challenge)
|
|
.as_str()
|
|
.ok_or(OmikronError::InvalidResponse)?;
|
|
|
|
let expected_challenge = self.challenge.read().await.clone();
|
|
|
|
if client_response == expected_challenge {
|
|
let omikron_id = self.state.read().await.omikron_id().unwrap_or(0);
|
|
*self.state.write().await = AuthState::Authenticated { omikron_id };
|
|
|
|
omikron_manager::add_omikron(self.clone()).await;
|
|
|
|
let response = CommunicationValue::new(CommunicationType::identification_response)
|
|
.with_id(cv.get_id())
|
|
.add_data(DataTypes::accepted, DataValue::Bool(true));
|
|
|
|
self.clone().send(&response).await?;
|
|
log_in!(omikron_id, PrintType::Omega, "Omikron authenticated");
|
|
Ok(())
|
|
} else {
|
|
let _ = self
|
|
.send_error_response(cv.get_id(), CommunicationType::error_invalid_challenge)
|
|
.await;
|
|
Err(OmikronError::AuthenticationFailed)
|
|
}
|
|
}
|
|
|
|
// -------------------------------------------------------------------------
|
|
// Authenticated Message Handlers
|
|
// -------------------------------------------------------------------------
|
|
|
|
async fn handle_authenticated(
|
|
self: Arc<Self>,
|
|
cv: CommunicationValue,
|
|
omikron_id: i64,
|
|
) -> OmikronResult<()> {
|
|
match cv.get_type() {
|
|
// Link shortening
|
|
CommunicationType::shorten_link => self.handle_shorten_link(cv).await,
|
|
|
|
// Online status tracking
|
|
CommunicationType::user_connected => {
|
|
self.handle_user_connected(cv, omikron_id).await;
|
|
Ok(())
|
|
}
|
|
CommunicationType::user_disconnected => {
|
|
self.handle_user_disconnected(cv, omikron_id).await;
|
|
Ok(())
|
|
}
|
|
CommunicationType::iota_connected => {
|
|
self.handle_iota_connected(cv, omikron_id).await;
|
|
Ok(())
|
|
}
|
|
CommunicationType::iota_disconnected => {
|
|
self.handle_iota_disconnected(cv, omikron_id).await;
|
|
Ok(())
|
|
}
|
|
CommunicationType::sync_client_iota_status => {
|
|
self.handle_sync_status(cv, omikron_id).await;
|
|
Ok(())
|
|
}
|
|
|
|
CommunicationType::get_user_data => self.handle_get_user_data(cv).await,
|
|
CommunicationType::get_iota_data => self.handle_get_iota_data(cv).await,
|
|
|
|
CommunicationType::get_register => self.handle_get_register(cv).await,
|
|
CommunicationType::complete_register_iota => {
|
|
self.handle_complete_register_iota(cv).await
|
|
}
|
|
CommunicationType::complete_register_user => {
|
|
self.handle_complete_register_user(cv).await
|
|
}
|
|
|
|
CommunicationType::change_user_data => self.handle_change_user_data(cv).await,
|
|
CommunicationType::change_iota_data => self.handle_change_iota_data(cv).await,
|
|
CommunicationType::delete_user => self.handle_delete_user(cv).await,
|
|
CommunicationType::delete_iota => self.handle_delete_iota(cv).await,
|
|
|
|
CommunicationType::get_notifications => self.handle_get_notifications(cv).await,
|
|
CommunicationType::read_notification => self.handle_read_notification(cv).await,
|
|
CommunicationType::push_notification => self.handle_push_notification(cv).await,
|
|
CommunicationType::get_states => self.handle_get_states(cv).await,
|
|
|
|
_ => {
|
|
log_err!(
|
|
0,
|
|
PrintType::Omega,
|
|
"Unknown message type: {:?}",
|
|
cv.get_type()
|
|
);
|
|
Ok(())
|
|
}
|
|
}
|
|
}
|
|
|
|
// -------------------------------------------------------------------------
|
|
// Specific Handlers (ported from original WebSocket implementation)
|
|
// -------------------------------------------------------------------------
|
|
|
|
async fn handle_shorten_link(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
|
|
let link = cv
|
|
.get_data(DataTypes::link)
|
|
.as_str()
|
|
.ok_or(OmikronError::InvalidResponse)?;
|
|
|
|
let short = add_short_link(link)
|
|
.await
|
|
.map_err(|_| OmikronError::Sql("Shortend link Error".to_string()))?;
|
|
|
|
let response = CommunicationValue::new(CommunicationType::shorten_link)
|
|
.with_id(cv.get_id())
|
|
.add_data(DataTypes::link, DataValue::Str(short));
|
|
|
|
self.send(&response).await
|
|
}
|
|
|
|
async fn handle_user_connected(self: Arc<Self>, cv: CommunicationValue, omikron_id: i64) {
|
|
log_in!(PrintType::Omega, "User connected");
|
|
if let Some(user_id) = cv.get_data(DataTypes::user_id).as_number() {
|
|
let status = cv
|
|
.get_data(DataTypes::user_state)
|
|
.as_str()
|
|
.and_then(|s| UserStatus::from_str(s))
|
|
.unwrap_or(UserStatus::user_online);
|
|
user_online_tracker::track_user_status(
|
|
user_id.try_into().unwrap(),
|
|
status,
|
|
omikron_id,
|
|
);
|
|
}
|
|
}
|
|
|
|
async fn handle_user_disconnected(self: Arc<Self>, cv: CommunicationValue, _omikron_id: i64) {
|
|
log_in!(PrintType::Omega, "User disconnected");
|
|
if let Some(user_id) = cv.get_data(DataTypes::user_id).as_number() {
|
|
if let Some(status) = user_online_tracker::get_user_status(user_id as i64) {
|
|
user_online_tracker::track_user_status(
|
|
user_id as i64,
|
|
UserStatus::user_offline,
|
|
status.omikron_id,
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn handle_iota_connected(self: Arc<Self>, cv: CommunicationValue, omikron_id: i64) {
|
|
log_in!(PrintType::Omega, "IOTA connected");
|
|
let iota_id = match cv.get_data(DataTypes::iota_id).as_number() {
|
|
Some(id) => id as i64,
|
|
None => return,
|
|
};
|
|
user_online_tracker::track_iota_connection(iota_id, omikron_id, true);
|
|
|
|
let mut user_ids = Vec::new();
|
|
if let Ok(users) = sql::get_users_by_iota_id(iota_id.try_into().unwrap()).await {
|
|
for (user_id, _, _, _, _, _, _, _, _, _, _, _) in users {
|
|
user_ids.push(DataValue::Number(user_id.try_into().unwrap()));
|
|
user_online_tracker::track_user_status(
|
|
user_id.try_into().unwrap(),
|
|
UserStatus::user_offline,
|
|
omikron_id,
|
|
);
|
|
}
|
|
} else {
|
|
log_in!(PrintType::General, "SQL error loading users for IOTA");
|
|
}
|
|
|
|
let response = CommunicationValue::new(CommunicationType::iota_user_data)
|
|
.with_id(cv.get_id())
|
|
.add_data(DataTypes::user_ids, DataValue::Array(user_ids));
|
|
|
|
let _ = self.send(&response).await;
|
|
}
|
|
|
|
async fn handle_iota_disconnected(self: Arc<Self>, cv: CommunicationValue, omikron_id: i64) {
|
|
log_in!(PrintType::Omega, "IOTA disconnected");
|
|
let iota_id = match cv.get_data(DataTypes::iota_id).as_number() {
|
|
Some(id) => id as i64,
|
|
None => return,
|
|
};
|
|
let iota_offline = user_online_tracker::untrack_iota_connection(iota_id, omikron_id);
|
|
if iota_offline {
|
|
if let Ok(users) = sql::get_users_by_iota_id(iota_id.try_into().unwrap()).await {
|
|
let user_ids: Vec<i64> = users.iter().map(|u| u.0.try_into().unwrap()).collect();
|
|
user_online_tracker::untrack_many_users(&user_ids);
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn handle_sync_status(self: Arc<Self>, cv: CommunicationValue, omikron_id: i64) {
|
|
if let DataValue::Array(user_ids) = cv.get_data(DataTypes::user_ids) {
|
|
for user_id_val in user_ids {
|
|
if let DataValue::Number(user_id) = user_id_val {
|
|
user_online_tracker::track_user_status(
|
|
*user_id,
|
|
UserStatus::user_offline,
|
|
omikron_id,
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
if let DataValue::Array(iota_ids) = cv.get_data(DataTypes::iota_ids) {
|
|
for iota_id_val in iota_ids {
|
|
if let DataValue::Number(iota_id) = iota_id_val {
|
|
user_online_tracker::track_iota_connection(*iota_id, omikron_id, true);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn handle_get_user_data(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
|
|
// Try by user_id first
|
|
if let Some(user_id) = cv.get_data(DataTypes::user_id).as_number() {
|
|
if let Ok(user_data) = get_by_user_id(user_id as i64).await {
|
|
let response = self
|
|
.clone()
|
|
.build_user_data_response(cv.get_id(), user_data)
|
|
.await;
|
|
return self.send(&response).await;
|
|
}
|
|
}
|
|
|
|
// Try by username
|
|
if let Some(username) = cv.get_data(DataTypes::username).as_str() {
|
|
if let Ok(user_data) = get_by_username(username).await {
|
|
let response = self
|
|
.clone()
|
|
.build_user_data_response(cv.get_id(), user_data)
|
|
.await;
|
|
return self.send(&response).await;
|
|
}
|
|
}
|
|
|
|
// Not found
|
|
let response =
|
|
CommunicationValue::new(CommunicationType::error_not_found).with_id(cv.get_id());
|
|
self.send(&response).await
|
|
}
|
|
|
|
async fn build_user_data_response(
|
|
self: Arc<Self>,
|
|
msg_id: u32,
|
|
user: (
|
|
i64,
|
|
i64,
|
|
String,
|
|
Option<String>,
|
|
Option<String>,
|
|
Option<String>,
|
|
Option<Vec<u8>>,
|
|
i32,
|
|
i64,
|
|
String,
|
|
String,
|
|
String,
|
|
),
|
|
) -> CommunicationValue {
|
|
let (
|
|
id,
|
|
iota_id,
|
|
username,
|
|
display,
|
|
status,
|
|
about,
|
|
avatar,
|
|
sub_level,
|
|
sub_end,
|
|
public_key,
|
|
_,
|
|
_,
|
|
) = user;
|
|
|
|
let mut response = CommunicationValue::new(CommunicationType::get_user_data)
|
|
.with_id(msg_id)
|
|
.add_data(DataTypes::username, DataValue::Str(username.clone()))
|
|
.add_data(DataTypes::public_key, DataValue::Str(public_key))
|
|
.add_data(DataTypes::user_id, DataValue::Number(id))
|
|
.add_data(DataTypes::iota_id, DataValue::Number(iota_id))
|
|
.add_data(DataTypes::sub_level, DataValue::Number(sub_level as i64))
|
|
.add_data(DataTypes::sub_end, DataValue::Number(sub_end));
|
|
|
|
// Display name (fallback to username)
|
|
let display_name = display.filter(|d| !d.is_empty()).unwrap_or(username);
|
|
response = response.add_data(DataTypes::display, DataValue::Str(display_name));
|
|
|
|
// Optional fields
|
|
if let Some(s) = status.filter(|s| !s.is_empty()) {
|
|
response = response.add_data(DataTypes::status, DataValue::Str(s));
|
|
}
|
|
if let Some(a) = about.filter(|a| !a.is_empty()) {
|
|
response = response.add_data(DataTypes::about, DataValue::Str(a));
|
|
}
|
|
if let Some(av) = avatar {
|
|
response = response.add_data(DataTypes::avatar, DataValue::Str(STANDARD.encode(av)));
|
|
}
|
|
|
|
// Online status
|
|
let user_status = user_online_tracker::get_user_status(id);
|
|
let iota_connections =
|
|
user_online_tracker::get_iota_omikron_connections(iota_id).unwrap_or_default();
|
|
|
|
if let Some(us) = user_status {
|
|
let display_status = if us.connection_type == UserStatus::user_invisible {
|
|
UserStatus::user_offline
|
|
} else {
|
|
us.connection_type.clone()
|
|
};
|
|
response = response.add_data(
|
|
DataTypes::online_status,
|
|
DataValue::Str(display_status.to_string()),
|
|
);
|
|
response = response.add_data(DataTypes::omikron_id, DataValue::Number(us.omikron_id));
|
|
} else {
|
|
response = response.add_data(
|
|
DataTypes::online_status,
|
|
DataValue::Str(UserStatus::iota_offline.to_string()),
|
|
);
|
|
}
|
|
|
|
response = response.add_data(
|
|
DataTypes::omikron_connections,
|
|
DataValue::Array(
|
|
iota_connections
|
|
.into_iter()
|
|
.map(DataValue::Number)
|
|
.collect(),
|
|
),
|
|
);
|
|
|
|
response
|
|
}
|
|
|
|
async fn handle_get_iota_data(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
|
|
// Try by iota_id
|
|
if let Some(iota_id) = cv.get_data(DataTypes::iota_id).as_number() {
|
|
if let Ok((iota_id, public_key)) = get_iota_by_id(iota_id as i64).await {
|
|
let response = self
|
|
.clone()
|
|
.build_iota_data_response(cv.get_id(), iota_id, public_key, None, None)
|
|
.await;
|
|
return self.send(&response).await;
|
|
}
|
|
}
|
|
|
|
// Try by user_id
|
|
if let Some(user_id) = cv.get_data(DataTypes::user_id).as_number() {
|
|
if let Ok((_, iota_id, _, _, _, _, _, _, _, _, _, _)) =
|
|
get_by_user_id(user_id as i64).await
|
|
{
|
|
if let Ok((iota_id, public_key)) = get_iota_by_id(iota_id).await {
|
|
let response = self
|
|
.clone()
|
|
.build_iota_data_response(
|
|
cv.get_id(),
|
|
iota_id,
|
|
public_key,
|
|
Some(user_id as i64),
|
|
None,
|
|
)
|
|
.await;
|
|
return self.send(&response).await;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Try by username
|
|
if let Some(username) = cv.get_data(DataTypes::username).as_str() {
|
|
if let Ok((user_id, iota_id, _, _, _, _, _, _, _, _, _, _)) =
|
|
get_by_username(username).await
|
|
{
|
|
if let Ok((iota_id, public_key)) = get_iota_by_id(iota_id).await {
|
|
let response = self
|
|
.clone()
|
|
.build_iota_data_response(
|
|
cv.get_id(),
|
|
iota_id,
|
|
public_key,
|
|
Some(user_id),
|
|
Some(username.to_string()),
|
|
)
|
|
.await;
|
|
return self.send(&response).await;
|
|
}
|
|
}
|
|
}
|
|
|
|
let response =
|
|
CommunicationValue::new(CommunicationType::error_not_found).with_id(cv.get_id());
|
|
self.send(&response).await
|
|
}
|
|
|
|
async fn build_iota_data_response(
|
|
self: Arc<Self>,
|
|
msg_id: u32,
|
|
iota_id: i64,
|
|
public_key: String,
|
|
user_id: Option<i64>,
|
|
username: Option<String>,
|
|
) -> CommunicationValue {
|
|
let mut response = CommunicationValue::new(CommunicationType::get_iota_data)
|
|
.with_id(msg_id)
|
|
.add_data(DataTypes::public_key, DataValue::Str(public_key))
|
|
.add_data(DataTypes::iota_id, DataValue::Number(iota_id));
|
|
|
|
if let Some(uid) = user_id {
|
|
response = response.add_data(DataTypes::user_id, DataValue::Number(uid));
|
|
}
|
|
if let Some(uname) = username {
|
|
response = response.add_data(DataTypes::username, DataValue::Str(uname));
|
|
}
|
|
|
|
let iota_connections =
|
|
user_online_tracker::get_iota_omikron_connections(iota_id).unwrap_or_default();
|
|
|
|
response.add_data(
|
|
DataTypes::omikron_connections,
|
|
DataValue::Array(
|
|
iota_connections
|
|
.into_iter()
|
|
.map(DataValue::Number)
|
|
.collect(),
|
|
),
|
|
)
|
|
}
|
|
|
|
async fn handle_get_register(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
|
|
let register_id = sql::get_register_id().await;
|
|
let response = CommunicationValue::new(CommunicationType::get_register)
|
|
.with_id(cv.get_id())
|
|
.add_data(DataTypes::user_id, DataValue::Number(register_id as i64));
|
|
self.send(&response).await
|
|
}
|
|
|
|
async fn handle_complete_register_iota(
|
|
self: Arc<Self>,
|
|
cv: CommunicationValue,
|
|
) -> OmikronResult<()> {
|
|
let iota_id_opt = cv
|
|
.get_data(DataTypes::iota_id)
|
|
.as_number()
|
|
.map(|n| n as i64);
|
|
|
|
if let Some(public_key) = cv.get_data(DataTypes::public_key).as_str() {
|
|
if let Some(iota_id) = iota_id_opt {
|
|
// Register existing IOTA
|
|
match sql::register_complete_iota(iota_id, public_key.to_string()).await {
|
|
Ok(_) => {
|
|
let response = CommunicationValue::new(CommunicationType::success)
|
|
.with_id(cv.get_id());
|
|
self.send(&response).await
|
|
}
|
|
Err(e) => {
|
|
let response = CommunicationValue::new(CommunicationType::error_internal)
|
|
.with_id(cv.get_id())
|
|
.add_data(DataTypes::error_type, DataValue::Str(e.to_string()));
|
|
self.send(&response).await
|
|
}
|
|
}
|
|
} else {
|
|
// Create new IOTA
|
|
match sql::create_new_iota(public_key.to_string()).await {
|
|
Ok(new_iota_id) => {
|
|
let response =
|
|
CommunicationValue::new(CommunicationType::complete_register_iota)
|
|
.with_id(cv.get_id())
|
|
.add_data(DataTypes::iota_id, DataValue::Number(new_iota_id));
|
|
self.send(&response).await
|
|
}
|
|
Err(e) => {
|
|
let response = CommunicationValue::new(CommunicationType::error_internal)
|
|
.with_id(cv.get_id())
|
|
.add_data(DataTypes::error_type, DataValue::Str(e.to_string()));
|
|
self.send(&response).await
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
self.send_error_response(cv.get_id(), CommunicationType::error_invalid_data)
|
|
.await
|
|
}
|
|
}
|
|
|
|
async fn handle_complete_register_user(
|
|
self: Arc<Self>,
|
|
cv: CommunicationValue,
|
|
) -> OmikronResult<()> {
|
|
let user_id = cv
|
|
.get_data(DataTypes::user_id)
|
|
.as_number()
|
|
.map(|n| n as i64);
|
|
let username = cv
|
|
.get_data(DataTypes::username)
|
|
.as_str()
|
|
.map(|s| s.to_string());
|
|
let public_key = cv
|
|
.get_data(DataTypes::public_key)
|
|
.as_str()
|
|
.map(|s| s.to_string());
|
|
let iota_id = cv.get_sender();
|
|
let reset_token = cv
|
|
.get_data(DataTypes::reset_token)
|
|
.as_str()
|
|
.map(|s| s.to_string());
|
|
|
|
if let (Some(uid), Some(uname), Some(pk), Some(rt)) =
|
|
(user_id, username, public_key, reset_token)
|
|
{
|
|
match sql::register_complete_user(uid, uname, pk, iota_id as i64, rt).await {
|
|
Ok(_) => {
|
|
let response =
|
|
CommunicationValue::new(CommunicationType::success).with_id(cv.get_id());
|
|
self.send(&response).await
|
|
}
|
|
Err(e) => {
|
|
let response = CommunicationValue::new(CommunicationType::error_internal)
|
|
.with_id(cv.get_id())
|
|
.add_data(DataTypes::error_type, DataValue::Str(e.to_string()));
|
|
self.send(&response).await
|
|
}
|
|
}
|
|
} else {
|
|
self.send_error_response(cv.get_id(), CommunicationType::error_invalid_data)
|
|
.await
|
|
}
|
|
}
|
|
|
|
async fn handle_change_user_data(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
|
|
let user_id = cv.get_sender() as i64;
|
|
let mut success = true;
|
|
let mut error_message = String::new();
|
|
|
|
// Process each field
|
|
if let Some(username) = cv.get_data(DataTypes::username).as_str() {
|
|
if let Err(e) = sql::change_username(user_id, username.to_string()).await {
|
|
success = false;
|
|
error_message = e.to_string();
|
|
}
|
|
}
|
|
if let Some(display) = cv.get_data(DataTypes::display).as_str() {
|
|
if let Err(e) = sql::change_display_name(user_id, display.to_string()).await {
|
|
success = false;
|
|
error_message = e.to_string();
|
|
}
|
|
}
|
|
if let Some(avatar) = cv.get_data(DataTypes::avatar).as_str() {
|
|
if let Err(e) = sql::change_avatar(user_id, avatar.to_string()).await {
|
|
success = false;
|
|
error_message = e.to_string();
|
|
}
|
|
}
|
|
if let Some(about) = cv.get_data(DataTypes::about).as_str() {
|
|
if let Err(e) = sql::change_about(user_id, about.to_string()).await {
|
|
success = false;
|
|
error_message = e.to_string();
|
|
}
|
|
}
|
|
if let Some(status) = cv.get_data(DataTypes::status).as_str() {
|
|
if let Err(e) = sql::change_status(user_id, status.to_string()).await {
|
|
success = false;
|
|
error_message = e.to_string();
|
|
}
|
|
}
|
|
if let (Some(public_key), Some(private_key_hash)) = (
|
|
cv.get_data(DataTypes::public_key).as_str(),
|
|
cv.get_data(DataTypes::private_key_hash).as_str(),
|
|
) {
|
|
if let Err(e) = sql::change_keys(
|
|
user_id,
|
|
public_key.to_string(),
|
|
private_key_hash.to_string(),
|
|
)
|
|
.await
|
|
{
|
|
success = false;
|
|
error_message = e.to_string();
|
|
}
|
|
}
|
|
|
|
if success {
|
|
let response = CommunicationValue::new(CommunicationType::success).with_id(cv.get_id());
|
|
self.send(&response).await
|
|
} else {
|
|
let response = CommunicationValue::new(CommunicationType::error_internal)
|
|
.with_id(cv.get_id())
|
|
.add_data(DataTypes::error_type, DataValue::Str(error_message));
|
|
self.send(&response).await
|
|
}
|
|
}
|
|
|
|
async fn handle_change_iota_data(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
|
|
let user_id = cv.get_sender() as i64;
|
|
|
|
if let (iota_id, Some(reset_token), Some(new_token)) = (
|
|
cv.get_sender(),
|
|
cv.get_data(DataTypes::reset_token).as_str(),
|
|
cv.get_data(DataTypes::new_token).as_str(),
|
|
) {
|
|
match sql::get_by_user_id(user_id).await {
|
|
Ok(user) => {
|
|
let current_token = user.11; // reset_token field
|
|
if current_token == reset_token {
|
|
let mut success = true;
|
|
let mut error_message = String::new();
|
|
|
|
if let Err(e) = sql::change_iota_id(user_id, iota_id as i64).await {
|
|
success = false;
|
|
error_message = e.to_string();
|
|
}
|
|
if success {
|
|
if let Err(e) = sql::change_token(user_id, new_token.to_string()).await
|
|
{
|
|
success = false;
|
|
error_message = e.to_string();
|
|
}
|
|
}
|
|
|
|
if success {
|
|
let response = CommunicationValue::new(CommunicationType::success)
|
|
.with_id(cv.get_id());
|
|
self.send(&response).await
|
|
} else {
|
|
let response =
|
|
CommunicationValue::new(CommunicationType::error_internal)
|
|
.with_id(cv.get_id())
|
|
.add_data(DataTypes::error_type, DataValue::Str(error_message));
|
|
self.send(&response).await
|
|
}
|
|
} else {
|
|
self.send_error_response(
|
|
cv.get_id(),
|
|
CommunicationType::error_invalid_challenge,
|
|
)
|
|
.await
|
|
}
|
|
}
|
|
Err(_) => {
|
|
self.send_error_response(cv.get_id(), CommunicationType::error_not_found)
|
|
.await
|
|
}
|
|
}
|
|
} else {
|
|
self.send_error_response(cv.get_id(), CommunicationType::error_invalid_data)
|
|
.await
|
|
}
|
|
}
|
|
|
|
async fn handle_delete_user(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
|
|
let user_id = cv.get_sender() as i64;
|
|
match sql::delete_user(user_id).await {
|
|
Ok(_) => {
|
|
let response =
|
|
CommunicationValue::new(CommunicationType::success).with_id(cv.get_id());
|
|
self.send(&response).await
|
|
}
|
|
Err(e) => {
|
|
let response = CommunicationValue::new(CommunicationType::error_internal)
|
|
.with_id(cv.get_id())
|
|
.add_data(DataTypes::error_type, DataValue::Str(e.to_string()));
|
|
self.send(&response).await
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn handle_delete_iota(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
|
|
let iota_id = cv.get_sender();
|
|
match sql::delete_iota(iota_id as i64).await {
|
|
Ok(_) => {
|
|
let response =
|
|
CommunicationValue::new(CommunicationType::success).with_id(cv.get_id());
|
|
self.send(&response).await
|
|
}
|
|
Err(e) => {
|
|
let response = CommunicationValue::new(CommunicationType::error_internal)
|
|
.with_id(cv.get_id())
|
|
.add_data(DataTypes::error_type, DataValue::Str(e.to_string()));
|
|
self.send(&response).await
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn handle_get_notifications(
|
|
self: Arc<Self>,
|
|
cv: CommunicationValue,
|
|
) -> OmikronResult<()> {
|
|
let user_id = cv.get_sender() as i64;
|
|
let response_array = match sql::get_notifications(user_id).await {
|
|
Ok(notifications) => notifications
|
|
.into_iter()
|
|
.map(|(sender, amount)| {
|
|
DataValue::Container(vec![
|
|
(DataTypes::sender_id, DataValue::Number(sender)),
|
|
(DataTypes::amount, DataValue::Number(amount)),
|
|
])
|
|
})
|
|
.collect(),
|
|
Err(e) => {
|
|
log!(PrintType::General, "SQL get_notifications error: {}", e);
|
|
vec![]
|
|
}
|
|
};
|
|
|
|
let response = CommunicationValue::new(CommunicationType::get_notifications)
|
|
.with_id(cv.get_id())
|
|
.add_data(DataTypes::notifications, DataValue::Array(response_array));
|
|
self.send(&response).await
|
|
}
|
|
|
|
async fn handle_read_notification(
|
|
self: Arc<Self>,
|
|
cv: CommunicationValue,
|
|
) -> OmikronResult<()> {
|
|
let receiver_id = match cv.get_sender() {
|
|
s if s > 0 => s as i64,
|
|
_ => match cv.get_data(DataTypes::receiver_id).as_number() {
|
|
Some(id) => id as i64,
|
|
None => return Ok(()),
|
|
},
|
|
};
|
|
|
|
if let Some(other_id) = cv
|
|
.get_data(DataTypes::sender_id)
|
|
.as_number()
|
|
.map(|n| n as i64)
|
|
{
|
|
if let Err(e) = sql::read_notification(receiver_id, other_id).await {
|
|
log!(PrintType::General, "SQL read_notification error: {}", e);
|
|
} else {
|
|
let response = CommunicationValue::new(CommunicationType::read_notification)
|
|
.with_id(cv.get_id());
|
|
let _ = self.send(&response).await;
|
|
|
|
// Sync with Tauri
|
|
crate::notifications::tauri::remove_notification(receiver_id, other_id).await;
|
|
|
|
// Sync with other Omikron clients
|
|
let sync_cv = CommunicationValue::new(CommunicationType::read_notification)
|
|
.with_receiver(receiver_id as u64)
|
|
.add_data(DataTypes::sender_id, DataValue::Number(other_id));
|
|
crate::transport::omikron_manager::send_to_user(receiver_id, &sync_cv).await;
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn handle_push_notification(
|
|
self: Arc<Self>,
|
|
cv: CommunicationValue,
|
|
) -> OmikronResult<()> {
|
|
let receiver_id = match cv.get_receiver() {
|
|
r if r > 0 => r as i64,
|
|
_ => match cv.get_data(DataTypes::receiver_id).as_number() {
|
|
Some(id) => id as i64,
|
|
None => return Ok(()),
|
|
},
|
|
};
|
|
|
|
let sender_id = match cv.get_data(DataTypes::sender_id).as_number() {
|
|
Some(id) => id as i64,
|
|
None => cv.get_sender() as i64,
|
|
};
|
|
|
|
if let Err(e) = sql::add_notification(receiver_id, sender_id).await {
|
|
log!(PrintType::General, "SQL add_notification error: {}", e);
|
|
} else {
|
|
let response =
|
|
CommunicationValue::new(CommunicationType::push_notification).with_id(cv.get_id());
|
|
let _ = self.send(&response).await;
|
|
|
|
// Sync with Tauri
|
|
crate::notifications::tauri::send_notification(receiver_id, sender_id).await;
|
|
|
|
// Sync with other Omikron clients
|
|
let push_cv = CommunicationValue::new(CommunicationType::push_notification)
|
|
.with_receiver(receiver_id as u64)
|
|
.add_data(DataTypes::sender_id, DataValue::Number(sender_id));
|
|
crate::transport::omikron_manager::send_to_user(receiver_id, &push_cv).await;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn handle_get_states(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
|
|
let user_ids = match cv.get_data(DataTypes::user_ids) {
|
|
DataValue::Array(ids) => ids,
|
|
_ => return Ok(()),
|
|
};
|
|
|
|
let mut states = Vec::new();
|
|
for id_val in user_ids {
|
|
if let DataValue::Number(user_id) = id_val {
|
|
let user_id = *user_id;
|
|
let status = user_online_tracker::get_user_status(user_id);
|
|
let status_str = match status {
|
|
Some(ref us) => {
|
|
if us.connection_type == UserStatus::user_invisible {
|
|
"user_offline".to_string()
|
|
} else {
|
|
us.connection_type.to_string()
|
|
}
|
|
}
|
|
None => UserStatus::iota_offline.to_string(),
|
|
};
|
|
let mut map = Vec::new();
|
|
map.push((DataTypes::user_id, DataValue::Number(user_id)));
|
|
map.push((DataTypes::user_state, DataValue::Str(status_str)));
|
|
states.push(DataValue::Container(map));
|
|
}
|
|
}
|
|
|
|
let response = CommunicationValue::new(CommunicationType::get_states)
|
|
.with_id(cv.get_id())
|
|
.add_data(DataTypes::user_states, DataValue::Array(states));
|
|
self.send(&response).await
|
|
}
|
|
|
|
async fn handle_ping(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
|
|
if let DataValue::Number(last_ping) = cv.get_data(DataTypes::last_ping) {
|
|
*self.ping.write().await = *last_ping;
|
|
}
|
|
|
|
let response = CommunicationValue::new(CommunicationType::pong).with_id(cv.get_id());
|
|
self.send(&response).await
|
|
}
|
|
|
|
// -------------------------------------------------------------------------
|
|
// Utilities
|
|
// -------------------------------------------------------------------------
|
|
|
|
async fn send(self: Arc<Self>, cv: &CommunicationValue) -> OmikronResult<()> {
|
|
if !cv.is_type(CommunicationType::pong) && !cv.is_type(CommunicationType::ping) {
|
|
log_cv_out!(PrintType::Omikron, cv);
|
|
}
|
|
|
|
let guard = self.sender.lock().await;
|
|
let sender = guard.as_ref().ok_or(OmikronError::NotConnected)?;
|
|
|
|
sender
|
|
.send(cv)
|
|
.await
|
|
.map_err(|e| OmikronError::Send(e.to_string()))
|
|
}
|
|
|
|
async fn send_error_response(
|
|
self: Arc<Self>,
|
|
message_id: u32,
|
|
error_type: CommunicationType,
|
|
) -> OmikronResult<()> {
|
|
let error = CommunicationValue::new(error_type).with_id(message_id);
|
|
self.send(&error).await
|
|
}
|
|
|
|
pub async fn close(self: Arc<Self>) {
|
|
log_in!(
|
|
self.get_omikron_id().await.unwrap_or(0),
|
|
PrintType::Omega,
|
|
"Omikron connection Closed"
|
|
);
|
|
}
|
|
|
|
async fn cleanup(self: Arc<Self>) {
|
|
if let Some(omikron_id) = self.state.read().await.omikron_id() {
|
|
if omikron_id != 0 {
|
|
log_in!(omikron_id, PrintType::Omega, "Omikron disconnected");
|
|
omikron_manager::remove_omikron(omikron_id).await;
|
|
user_online_tracker::untrack_omikron(omikron_id).await;
|
|
}
|
|
}
|
|
|
|
if let Some(handle) = self.cleanup_handle.lock().unwrap().take() {
|
|
handle.abort();
|
|
}
|
|
}
|
|
|
|
// Public API for external use
|
|
pub async fn is_authenticated(self: Arc<Self>) -> bool {
|
|
self.state.read().await.is_authenticated()
|
|
}
|
|
|
|
pub async fn get_omikron_id(self: Arc<Self>) -> Option<i64> {
|
|
self.state.read().await.omikron_id()
|
|
}
|
|
|
|
pub async fn send_message(self: Arc<Self>, cv: &CommunicationValue) -> OmikronResult<()> {
|
|
self.send(cv).await
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Server Startup
|
|
// ============================================================================
|
|
|
|
pub async fn start(port: u16) -> Result<(), Box<dyn std::error::Error>> {
|
|
let cert_pem = load_file_vec("certs", "transport_cert.pem").expect("Error loading Pemfile");
|
|
|
|
let key_pem = load_file_vec("certs", "transport_key.pem").expect("Error loading Keyfile");
|
|
|
|
let mut host: Host = ttp_native::host(
|
|
port,
|
|
cert_pem,
|
|
key_pem,
|
|
Policy {
|
|
send_mode: SendMode::SingleStreamPerMessage,
|
|
max_message_size: 1_000_000_000,
|
|
close_frame_len: u32::MAX,
|
|
application_close_code: 0,
|
|
open_stream_timeout: Duration::from_millis(2_000),
|
|
write_timeout: Duration::from_millis(2_000),
|
|
accept_stream_timeout: Duration::from_millis(10_000),
|
|
read_timeout: Duration::from_millis(30_000),
|
|
keep_alive_interval: Some(Duration::from_secs(6)),
|
|
max_idle_timeout: Some(Duration::from_secs(30)),
|
|
force_close_delay: Duration::from_millis(300),
|
|
max_transient_recv_errors: 20,
|
|
transient_recv_backoff: Duration::from_millis(100),
|
|
receiver_queue_capacity: 1000,
|
|
},
|
|
)
|
|
.await?;
|
|
log!("OmikronServer listening on port {}", port);
|
|
|
|
while let Some((sender, mut receiver)) = host.next().await {
|
|
tokio::spawn(async move {
|
|
let conn = OmikronConnection::new(sender);
|
|
conn.handle(&mut receiver).await;
|
|
});
|
|
}
|
|
|
|
Ok(())
|
|
}
|