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

[WIP] Daemon & CLI

Browse source
This commit is contained in:
Alex-Emmet 2026-07-23 23:13:02 +02:00
commit 8b158108bb
100 changed files with 6519 additions and 1596 deletions
Download patch file Download diff file Expand all files Collapse all files

1
omikron-connector/Cargo.toml
View file

@ -4,6 +4,7 @@ version = "0.1.0"
edition = "2024"
[dependencies]
async-trait = "0.1.89"
iota-connection = { path = "../iota-connection" }
iota-logger = { path = "../iota-logger" }
iota-state = { path = "../iota-state" }

43
omikron-connector/src/client.rs Normal file
View file

@ -0,0 +1,43 @@
use async_trait::async_trait;
use mtp::codec::CommunicationValue;
use std::time::Duration;
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum OmikronError {
Disconnected(String),
Timeout(String),
Authentication(String),
Internal(String),
}
impl std::fmt::Display for OmikronError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Disconnected(v)
| Self::Timeout(v)
| Self::Authentication(v)
| Self::Internal(v) => f.write_str(v),
}
}
}
impl std::error::Error for OmikronError {}
pub enum OmikronStartupError {
Construction(String),
InitialConnectionTimeout {
connection: std::sync::Arc<crate::omikron_connection::OmikronConnection>,
},
Authentication,
}
#[async_trait]
pub trait OmikronClient: Send + Sync {
async fn send_message(&self, value: &CommunicationValue) -> Result<(), OmikronError>;
async fn await_response(
&self,
value: &CommunicationValue,
timeout: Duration,
) -> Result<CommunicationValue, OmikronError>;
async fn reconnect(&self) -> Result<(), OmikronError>;
async fn is_connected(&self) -> bool;
}

4
omikron-connector/src/lib.rs
View file

@ -1,4 +1,8 @@
pub mod client;
pub mod omega_discovery;
pub mod omikron_connection;
pub mod ping_pong_task;
pub mod user_ops;
pub use client::{OmikronClient, OmikronError, OmikronStartupError};
pub use omikron_connection::OmikronConnection;

116
omikron-connector/src/omikron_connection.rs
View file

@ -1,6 +1,6 @@
use dashmap::DashMap;
use dashmap::{DashMap, DashSet};
use iota_logger::{log, log_cv_in, log_cv_out, log_t};
use iota_state::ACTIVE_TASKS;
use iota_state::AppState;
use iota_storage::util::chat_files::{self, MessageState, change_message_state};
use iota_storage::util::config_util::{CONFIG, modify_config};
use iota_storage::util::e2ee_storage::{self, PendingChatSecretForward, StoredChatSecret};
@ -19,6 +19,7 @@ use tokio::time::sleep;
use tokio_util::sync::CancellationToken;
use uuid::Uuid;
use crate::client::{OmikronClient, OmikronError};
use crate::omega_discovery;
use iota_connection::message_common::*;
@ -106,7 +107,7 @@ const MAX_CONCURRENT_HANDLERS: usize = 20;
// ============================================================================
pub struct WaitingTask {
pub task: Box<dyn FnOnce(Arc<OmikronConnection>, CommunicationValue) -> bool + Send + Sync>,
pub task: Box<dyn FnOnce(CommunicationValue) -> bool + Send + Sync>,
pub inserted_at: Instant,
}
@ -163,14 +164,20 @@ pub struct OmikronConnection {
pub(crate) missed_pongs: Arc<AtomicU32>,
handler_semaphore: Arc<Semaphore>,
cancellation: CancellationToken,
pub(crate) active_tasks: Arc<DashSet<String>>,
pub(crate) app: Arc<std::sync::Mutex<AppState>>,
}
impl OmikronConnection {
pub fn new() -> Self {
Self::with_cancellation(CancellationToken::new())
pub fn new(active_tasks: Arc<DashSet<String>>, app: Arc<std::sync::Mutex<AppState>>) -> Self {
Self::with_cancellation(CancellationToken::new(), active_tasks, app)
}
pub fn with_cancellation(cancellation: CancellationToken) -> Self {
pub fn with_cancellation(
cancellation: CancellationToken,
active_tasks: Arc<DashSet<String>>,
app: Arc<std::sync::Mutex<AppState>>,
) -> Self {
let (shutdown_tx, _) = watch::channel(false);
let (state_watch_tx, _) = watch::channel(ConnectionState::Disconnected);
@ -190,6 +197,8 @@ impl OmikronConnection {
missed_pongs: Arc::new(AtomicU32::new(0)),
handler_semaphore: Arc::new(Semaphore::new(MAX_CONCURRENT_HANDLERS)),
cancellation,
active_tasks,
app,
}
}
@ -388,7 +397,7 @@ impl OmikronConnection {
*self.heartbeat_handle.lock().await = Some(heartbeat_handle);
{
ACTIVE_TASKS.insert("Omikron Listener".to_string());
self.active_tasks.insert("Omikron Listener".to_string());
}
// Wait for read loop to complete
@ -396,7 +405,7 @@ impl OmikronConnection {
*self.sender.write().await = None;
self.set_state(ConnectionState::Disconnected).await;
{
ACTIVE_TASKS.remove("Omikron Listener");
self.active_tasks.remove("Omikron Listener");
}
if let Some(handle) = self.heartbeat_handle.lock().await.take() {
@ -570,7 +579,7 @@ impl OmikronConnection {
Ok(cv) => {
let msg_id = cv.get_id();
if let Some((_, task)) = WAITING_TASKS.remove(&msg_id) {
if (task.task)(self.clone(), cv.clone()) {
if (task.task)(cv.clone()) {
continue;
}
}
@ -832,7 +841,7 @@ impl OmikronConnection {
let msg_id = cv.get_id();
if let Some((_, task)) = WAITING_TASKS.remove(&msg_id) {
if (task.task)(self.clone(), cv.clone()) {
if (task.task)(cv.clone()) {
return;
}
}
@ -1797,7 +1806,7 @@ impl OmikronConnection {
let response = CommunicationValue::new(CommunicationType::ErrorInternal)
.with_id(key)
.add_typed_default(DataType::Message, DataValue::Str(reason.clone()));
let _ = (waiting_task.task)(OMIKRON_CONNECTION.clone(), response);
let _ = (waiting_task.task)(response);
}
}
}
@ -1821,7 +1830,7 @@ impl OmikronConnection {
WAITING_TASKS.insert(
msg_id,
WaitingTask {
task: Box::new(move |_, response_cv| {
task: Box::new(move |response_cv| {
let _ = tx.send(response_cv);
true
}),
@ -1932,20 +1941,26 @@ impl OmikronConnection {
// Global Instance
// ============================================================================
pub static OMIKRON_CONNECTION: LazyLock<Arc<OmikronConnection>> = LazyLock::new(|| {
let conn = Arc::new(OmikronConnection::new());
start_task_cleanup_loop();
conn
});
pub async fn get_omikron_connection(
pub async fn connect_initial(
cancellation: CancellationToken,
) -> Option<Arc<OmikronConnection>> {
let conn = Arc::new(OmikronConnection::with_cancellation(cancellation));
active_tasks: Arc<DashSet<String>>,
app: Arc<std::sync::Mutex<AppState>>,
) -> Result<Arc<OmikronConnection>, crate::client::OmikronStartupError> {
let conn = Arc::new(OmikronConnection::with_cancellation(
cancellation,
active_tasks,
app,
));
conn.connect().await;
Some(conn)
match conn.await_connection(Some(CONNECTION_TIMEOUT)).await {
Ok(()) => Ok(conn),
Err(_) if conn.has_auth_failure().await => {
Err(crate::client::OmikronStartupError::Authentication)
}
Err(_) => {
Err(crate::client::OmikronStartupError::InitialConnectionTimeout { connection: conn })
}
}
}
impl iota_connection::connection_handler::ConnectionHandler for OmikronConnection {
@ -1973,3 +1988,56 @@ impl iota_connection::connection_handler::ConnectionHandler for OmikronConnectio
OmikronConnection::stop(self).await
}
}
#[async_trait::async_trait]
impl OmikronClient for OmikronConnection {
async fn send_message(&self, value: &CommunicationValue) -> Result<(), OmikronError> {
Self::send_message(self, value)
.await
.map_err(OmikronError::Disconnected)
}
async fn await_response(
&self,
value: &CommunicationValue,
timeout: Duration,
) -> Result<CommunicationValue, OmikronError> {
Self::await_response(self, value, Some(timeout))
.await
.map_err(|error| {
if error.contains("timed out") {
OmikronError::Timeout(error)
} else {
OmikronError::Disconnected(error)
}
})
}
async fn reconnect(&self) -> Result<(), OmikronError> {
let this = Arc::new(Self {
state: self.state.clone(),
state_watch_tx: self.state_watch_tx.clone(),
sender: self.sender.clone(),
connection_loop_handle: self.connection_loop_handle.clone(),
last_ping: self.last_ping.clone(),
heartbeat_handle: self.heartbeat_handle.clone(),
connection_id: self.connection_id,
shutdown_tx: self.shutdown_tx.clone(),
reconnect_on_close: self.reconnect_on_close.clone(),
auth_failure: self.auth_failure.clone(),
app_challenges: self.app_challenges.clone(),
app_sessions: self.app_sessions.clone(),
missed_pongs: self.missed_pongs.clone(),
handler_semaphore: self.handler_semaphore.clone(),
cancellation: self.cancellation.clone(),
active_tasks: self.active_tasks.clone(),
app: self.app.clone(),
});
Self::reconnect(&this).await;
Ok(())
}
async fn is_connected(&self) -> bool {
Self::is_connected(self).await
}
}

9
omikron-connector/src/ping_pong_task.rs
View file

@ -1,9 +1,8 @@
use crate::omikron_connection::OmikronConnection;
use dashmap::DashMap;
use iota_state::APP_STATE;
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
use std::sync::atomic::Ordering;
use std::sync::LazyLock;
use std::sync::atomic::Ordering;
use std::time::Instant;
use tokio::time::Duration;
@ -23,7 +22,9 @@ impl OmikronConnection {
.with_id(id)
.add_typed_default(
DataType::LastPing,
DataValue::Array(vec![DataValue::SignedNumber(*self.last_ping.lock().await as i128)]),
DataValue::Array(vec![DataValue::SignedNumber(
*self.last_ping.lock().await as i128,
)]),
);
let _ = self.send_message(&ping_message).await;
@ -37,7 +38,7 @@ impl OmikronConnection {
if let Some((_, send_time)) = PING_TIMES.remove(&id) {
let ping_ms = Instant::now().duration_since(send_time).as_millis() as i64;
*self.last_ping.lock().await = ping_ms;
APP_STATE.lock().unwrap().push_ping_val(ping_ms as f64);
self.app.lock().unwrap().push_ping_val(ping_ms as f64);
}
}
}

27
omikron-connector/src/user_ops.rs
View file

@ -8,21 +8,22 @@ use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
use rand_core::{OsRng, RngCore};
use std::time::Duration;
use crate::OmikronClient;
use crate::omega_discovery;
use crate::omikron_connection::OMIKRON_CONNECTION;
pub async fn create_user(username: &str) -> (Option<UserProfile>, Option<String>) {
pub async fn create_user(
connection: &dyn OmikronClient,
username: &str,
) -> (Option<UserProfile>, Option<String>) {
let register_communication_value = CommunicationValue::new(CommunicationType::GetRegister);
let connection = OMIKRON_CONNECTION.clone();
let response_communication_value = match connection
.await_response(&register_communication_value, Some(Duration::from_secs(20)))
.await_response(&register_communication_value, Duration::from_secs(20))
.await
{
Ok(communication_value) => communication_value,
Err(e) => {
log_t!("User creation: {}", e);
log_t!("User creation: {}", e.to_string());
return (None, None);
}
};
@ -68,7 +69,7 @@ pub async fn create_user(username: &str) -> (Option<UserProfile>, Option<String>
.add_typed_default(DataType::ResetToken, DataValue::Str(reset_token));
let response_communication_value = connection
.await_response(&communication_value, Some(Duration::from_secs(20)))
.await_response(&communication_value, Duration::from_secs(20))
.await;
if let Ok(response) = response_communication_value {
@ -84,13 +85,15 @@ pub async fn create_user(username: &str) -> (Option<UserProfile>, Option<String>
save_file(
"",
&format!("{}.tu", username),
&format!("{}@{}::{}", user_id, omega_discovery::omega_host(), keyring_b64),
&format!(
"{}@{}::{}",
user_id,
omega_discovery::omega_host(),
keyring_b64
),
);
add_user(user_profile.clone());
save_users();
(
Some(user_profile),
Some(keyring_b64),
)
(Some(user_profile), Some(keyring_b64))
}