[Fix] IPC, cli, daemon

iota-daemon-lib/Cargo.toml

@@ -11,5 +11,9 @@ iota-storage = { path = "../iota-storage" }
 iota-util = { path = "../iota-util" }
 omikron-connector = { path = "../omikron-connector" }
 mtp = { git = "https://git.methanium.net/Methanium/mtp.git" }
+dashmap = "6.1.0"
+libc = "0.2"
 sysinfo = "0.38.3"
 tokio = { version = "1.50.0", features = ["full"] }
+tokio-util = { version = "0.7", features = ["rt"] }
+uuid = { version = "*", features = ["v4"] }

iota-daemon-lib/src/command_router.rs

@@ -1,5 +1,7 @@
 use crate::DaemonRuntime;
-use iota_ipc::DaemonMessage;
+use iota_ipc::{
+    IpcErrorCode, LocalRequest, ResponseEnvelope, ResponseResult,
+};
 use iota_logger::{log, log_command};
 use iota_storage::users::user_manager;
 use iota_storage::util::config_util::modify_config;
@@ -8,6 +10,8 @@ use omikron_connector::omikron_connection::OMIKRON_CONNECTION;
 use std::sync::Arc;
 use std::time::Duration;
 
+use crate::daemon_state::ShutdownReason;
+
 #[derive(Clone)]
 pub struct CommandRouter {
     runtime: Arc<DaemonRuntime>,
@@ -18,85 +22,124 @@ impl CommandRouter {
         Self { runtime }
     }
 
-    pub async fn route(&self, seq: u64, line: String) -> DaemonMessage {
-        log_command!("{}", line);
-        let result = self.execute(&line).await;
-        DaemonMessage::CommandResult {
-            seq,
-            success: result.is_ok(),
-            message: result.unwrap_or_else(|error| error),
+    pub async fn route(&self, request_id: u64, request: LocalRequest) -> ResponseEnvelope {
+        log_command!("{:?}", request);
+        let result = self.execute(request).await;
+        ResponseEnvelope {
+            request_id,
+            result,
         }
     }
 
-    async fn execute(&self, line: &str) -> Result<String, String> {
-        let parts = line
+    /// Parse a legacy console command string into a typed request.
+    pub fn parse_console_command(line: &str) -> Option<LocalRequest> {
+        let parts: Vec<&str> = line
             .trim_start_matches('/')
             .split_whitespace()
-            .collect::<Vec<_>>();
+            .collect();
         match parts.as_slice() {
-            ["tasks"] => Ok(self
-                .runtime
-                .state
-                .active_tasks
-                .iter()
-                .map(|task| task.to_string())
-                .collect::<Vec<_>>()
-                .join(", ")),
-            ["help"] => Ok(
-                "Available commands: tasks, ping, user, reconnect, regenerate, reload, shutdown"
-                    .into(),
-            ),
-            ["ping"] => self.ping(20).await,
-            ["ping", seconds] => self.ping(seconds.parse::<u64>().unwrap_or(20)).await,
-            ["user", "add", username] => {
-                let (user, _) = omikron_connector::user_ops::create_user(username).await;
-                user.map(|user| format!("Created user {}", user.user_id))
-                    .ok_or_else(|| "User creation failed".into())
-            }
+            ["help"] => None,
+            ["tasks"] => Some(LocalRequest::ListTasks),
+            ["ping", _] | ["ping"] => None,
+            ["user", "add", username] => Some(LocalRequest::CreateUser {
+                username: username.to_string(),
+            }),
             ["user", "remove", username] => {
-                let user = user_manager::get_user_by_username(username)
-                    .ok_or_else(|| "Username does not exist".to_string())?;
+                let user = user_manager::get_user_by_username(username)?;
+                Some(LocalRequest::RemoveUser {
+                    user_id: user.user_id,
+                })
+            }
+            ["user", "list"] => Some(LocalRequest::ListUsers),
+            ["reconnect"] => Some(LocalRequest::ReconnectOmikron),
+            ["regenerate", "keys"] => Some(LocalRequest::RotateIotaIdentity),
+            ["reload"] | ["restart"] => Some(LocalRequest::RestartDaemon),
+            ["shutdown"] | ["stop"] => Some(LocalRequest::StopDaemon),
+            _ => None,
+        }
+    }
+
+    async fn execute(&self, request: LocalRequest) -> ResponseResult {
+        match request {
+            LocalRequest::GetStatus => {
+                let phase = self.runtime.current_startup_phase();
+                let degraded = self.runtime.degraded_reason.borrow().clone();
+                let tasks: Vec<String> = self
+                    .runtime
+                    .state
+                    .active_tasks
+                    .iter()
+                    .map(|task| task.to_string())
+                    .collect();
+                let mut info = format!("Phase: {:?}, Tasks: {}", phase, tasks.join(", "));
+                if let Some(reason) = degraded {
+                    info.push_str(&format!(", Degraded: {}", reason));
+                }
+                ResponseResult::Ok(info)
+            }
+            LocalRequest::ListTasks => {
+                let tasks: Vec<String> = self
+                    .runtime
+                    .state
+                    .active_tasks
+                    .iter()
+                    .map(|task| task.to_string())
+                    .collect();
+                ResponseResult::Ok(tasks.join(", "))
+            }
+            LocalRequest::ListUsers => {
+                let users: Vec<String> = user_manager::get_users()
+                    .into_iter()
+                    .map(|user| format!("{} ({})", user.username, user.user_id))
+                    .collect();
+                ResponseResult::Ok(users.join("\n"))
+            }
+            LocalRequest::CreateUser { username } => {
+                match omikron_connector::user_ops::create_user(&username).await {
+                    (Some(user), _) => {
+                        ResponseResult::Ok(format!("Created user {}", user.user_id))
+                    }
+                    _ => ResponseResult::Error(IpcErrorCode::StorageFailure),
+                }
+            }
+            LocalRequest::RemoveUser { user_id } => {
+                let user = match user_manager::get_user(user_id) {
+                    Some(user) => user,
+                    None => return ResponseResult::Error(IpcErrorCode::NotFound),
+                };
                 let message = CommunicationValue::new(CommunicationType::DeleteUser)
                     .with_sender(user.user_id as u64);
-                OMIKRON_CONNECTION
-                    .send_message(&message)
-                    .await
-                    .map_err(|error| error.to_string())?;
+                if let Err(_e) = OMIKRON_CONNECTION.send_message(&message).await {
+                    return ResponseResult::Error(IpcErrorCode::OmikronUnavailable);
+                }
                 user_manager::remove_user(user.user_id);
-                Ok(format!("Removed user {}", user.user_id))
+                ResponseResult::Ok(format!("Removed user {}", user.user_id))
             }
-            ["user", "list"] => Ok(user_manager::get_users()
-                .into_iter()
-                .map(|user| format!("{} ({})", user.username, user.user_id))
-                .collect::<Vec<_>>()
-                .join("\n")),
-            ["reconnect"] => {
+            LocalRequest::ReconnectOmikron => {
                 OMIKRON_CONNECTION.reconnect().await;
-                Ok("Reconnected to Omikron server".into())
+                ResponseResult::Ok("Reconnected to Omikron server".into())
             }
-            ["regenerate", "keys"] => {
+            LocalRequest::RotateIotaIdentity => {
                 modify_config(|config| {
                     config.public_key = None;
                     config.private_key = None;
                     config.iota_id = None;
                 });
                 OMIKRON_CONNECTION.reconnect().await;
-                Ok("Key pair regenerated and Omikron reconnection requested".into())
+                ResponseResult::Ok("Key pair regenerated and Omikron reconnection requested".into())
             }
-            ["reload"] | ["restart"] => {
-                *self.runtime.state.reload.write().await = true;
-                *self.runtime.state.shutdown.write().await = true;
-                Ok("Daemon restart requested".into())
+            LocalRequest::RestartDaemon => {
+                self.runtime.shutdown(ShutdownReason::Restart);
+                ResponseResult::Ok("Daemon restart requested".into())
             }
-            ["shutdown"] | ["stop"] => {
-                *self.runtime.state.shutdown.write().await = true;
-                Ok("Daemon shutdown requested".into())
+            LocalRequest::StopDaemon => {
+                self.runtime.shutdown(ShutdownReason::Stop);
+                ResponseResult::Ok("Daemon shutdown requested".into())
             }
-            _ => Err("Unknown command".into()),
         }
     }
 
-    async fn ping(&self, seconds: u64) -> Result<String, String> {
+    pub async fn ping(&self, seconds: u64) -> Result<String, String> {
         let response = OMIKRON_CONNECTION
             .await_response(
                 &CommunicationValue::new(CommunicationType::Ping),

iota-daemon-lib/src/daemon_state.rs

@@ -3,21 +3,123 @@ use iota_state::DaemonState;
 use std::sync::Arc;
 use std::time::Duration;
 use sysinfo::{RefreshKind, System};
+use tokio::sync::watch;
+use tokio_util::sync::CancellationToken;
+
+/// Reason the daemon is shutting down.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum ShutdownReason {
+    Stop,
+    Restart,
+    Fatal(String),
+}
+
+impl ShutdownReason {
+    pub fn exit_code(&self) -> i32 {
+        match self {
+            ShutdownReason::Stop => 0,
+            ShutdownReason::Restart => 75,
+            ShutdownReason::Fatal(_) => 1,
+        }
+    }
+}
+
+/// Tracks the lifecycle phase of the daemon for IPC visibility.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum StartupPhase {
+    Starting,
+    MigratingStorage,
+    LoadingUsers,
+    StartingServices,
+    Ready,
+    Degraded,
+    Stopping,
+}
+
+impl From<StartupPhase> for iota_ipc::StartupPhase {
+    fn from(phase: StartupPhase) -> Self {
+        match phase {
+            StartupPhase::Starting => iota_ipc::StartupPhase::Starting,
+            StartupPhase::MigratingStorage => iota_ipc::StartupPhase::MigratingStorage,
+            StartupPhase::LoadingUsers => iota_ipc::StartupPhase::LoadingUsers,
+            StartupPhase::StartingServices => iota_ipc::StartupPhase::StartingServices,
+            StartupPhase::Ready => iota_ipc::StartupPhase::Ready,
+            StartupPhase::Degraded => iota_ipc::StartupPhase::Degraded,
+            StartupPhase::Stopping => iota_ipc::StartupPhase::Stopping,
+        }
+    }
+}
 
 /* This wrapper exposes daemon state as IPC-safe snapshots while preserving a
- * single owned state instance for all daemon subsystems. */
-#[derive(Clone, Default)]
+ * single owned state instance for all daemon subsystems. The cancellation token
+ * is the single lifecycle signal — all subsystems check it instead of a
+ * separate boolean. */
 pub struct DaemonRuntime {
     pub state: Arc<DaemonState>,
+    pub cancellation: CancellationToken,
+    pub shutdown_tx: watch::Sender<Option<ShutdownReason>>,
+    pub startup_phase: watch::Sender<StartupPhase>,
+    pub degraded_reason: watch::Sender<Option<String>>,
+}
+
+impl Clone for DaemonRuntime {
+    fn clone(&self) -> Self {
+        Self {
+            state: self.state.clone(),
+            cancellation: self.cancellation.clone(),
+            shutdown_tx: self.shutdown_tx.clone(),
+            startup_phase: self.startup_phase.clone(),
+            degraded_reason: self.degraded_reason.clone(),
+        }
+    }
+}
+
+impl Default for DaemonRuntime {
+    fn default() -> Self {
+        Self::new()
+    }
 }
 
 impl DaemonRuntime {
     pub fn new() -> Self {
+        let (shutdown_tx, _) = watch::channel(None);
+        let (startup_phase, _) = watch::channel(StartupPhase::Starting);
+        let (degraded_reason, _) = watch::channel(None);
         Self {
             state: Arc::new(DaemonState::new()),
+            cancellation: CancellationToken::new(),
+            shutdown_tx,
+            startup_phase,
+            degraded_reason,
         }
     }
 
+    pub fn shutdown(&self, reason: ShutdownReason) {
+        self.cancellation.cancel();
+        let _ = self.shutdown_tx.send(Some(reason));
+    }
+
+    pub fn shutdown_reason(&self) -> Option<ShutdownReason> {
+        self.shutdown_tx.borrow().clone()
+    }
+
+    pub fn is_shutting_down(&self) -> bool {
+        self.cancellation.is_cancelled()
+    }
+
+    pub fn set_startup_phase(&self, phase: StartupPhase) {
+        let _ = self.startup_phase.send(phase);
+    }
+
+    pub fn current_startup_phase(&self) -> StartupPhase {
+        *self.startup_phase.borrow()
+    }
+
+    pub fn mark_degraded(&self, reason: String) {
+        let _ = self.degraded_reason.send(Some(reason.clone()));
+        let _ = self.startup_phase.send(StartupPhase::Degraded);
+    }
+
     pub fn snapshot(&self) -> StateSnapshot {
         let state = self
             .state
@@ -41,7 +143,7 @@ impl DaemonRuntime {
             let mut system = System::new_with_specifics(RefreshKind::everything());
             let mut counter = 0.0;
             loop {
-                if *runtime.state.shutdown.read().await {
+                if runtime.is_shutting_down() {
                     break;
                 }
                 system.refresh_cpu_all();

iota-daemon-lib/src/ipc_server.rs

@@ -1,28 +1,42 @@
 use crate::{CommandRouter, DaemonRuntime};
-use iota_ipc::{ClientMessage, DaemonMessage, read_msg, write_msg};
+use iota_ipc::{
+    ClientMessage, DaemonMessage, HelloAck, MIN_PROTOCOL_VERSION, PROTOCOL_VERSION, read_msg,
+    write_msg,
+};
+use iota_logger::log;
 use std::io::Result;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use std::{env, os::fd::FromRawFd};
 use tokio::net::{UnixListener, UnixStream};
-use tokio::sync::broadcast;
+use tokio::sync::{broadcast, mpsc, watch};
+use uuid::Uuid;
+
+/// Per-client outbound queue capacity.
+const CLIENT_CHANNEL_SIZE: usize = 256;
+
+/// Maximum handshake retries before giving up.
+const MAX_HANDSHAKE_RETRIES: u32 = 10;
 
 pub struct IpcServer {
     path: PathBuf,
     runtime: Arc<DaemonRuntime>,
-    messages: broadcast::Sender<DaemonMessage>,
+    log_tx: broadcast::Sender<DaemonMessage>,
+    state_rx: watch::Sender<iota_ipc::StateSnapshot>,
 }
 
 impl IpcServer {
     pub fn new(
         path: impl Into<PathBuf>,
         runtime: Arc<DaemonRuntime>,
-        messages: broadcast::Sender<DaemonMessage>,
+        log_tx: broadcast::Sender<DaemonMessage>,
+        state_rx: watch::Sender<iota_ipc::StateSnapshot>,
     ) -> Self {
         Self {
             path: path.into(),
             runtime,
-            messages,
+            log_tx,
+            state_rx,
         }
     }
 
@@ -38,11 +52,14 @@ impl IpcServer {
             }
         };
         loop {
-            let (stream, _) = listener.accept().await?;
+            let (stream, _addr) = listener.accept().await?;
             let runtime = self.runtime.clone();
-            let messages = self.messages.clone();
+            let log_tx = self.log_tx.clone();
+            let state_rx = self.state_rx.clone();
             tokio::spawn(async move {
-                let _ = handle_client(stream, runtime, messages).await;
+                if let Err(error) = handle_client(stream, runtime, log_tx, state_rx).await {
+                    eprintln!("IPC client error: {error}");
+                }
             });
         }
     }
@@ -72,34 +89,159 @@ async fn remove_stale_socket(path: &Path) -> Result<()> {
     }
 }
 
+#[derive(Clone, Debug)]
+struct PeerIdentity {
+    pid: i32,
+    uid: u32,
+    gid: u32,
+}
+
+fn peer_credentials(stream: &UnixStream) -> PeerIdentity {
+    #[cfg(target_os = "linux")]
+    {
+        use std::os::unix::io::AsRawFd;
+        unsafe {
+            let mut cred: libc::ucred = std::mem::zeroed();
+            let mut len = std::mem::size_of::<libc::ucred>() as libc::socklen_t;
+            let fd = stream.as_raw_fd();
+            libc::getsockopt(
+                fd,
+                libc::SOL_SOCKET,
+                libc::SO_PEERCRED,
+                &mut cred as *mut _ as *mut libc::c_void,
+                &mut len,
+            );
+            PeerIdentity {
+                pid: cred.pid,
+                uid: cred.uid,
+                gid: cred.gid,
+            }
+        }
+    }
+    #[cfg(not(target_os = "linux"))]
+    {
+        PeerIdentity {
+            pid: 0,
+            uid: 0,
+            gid: 0,
+        }
+    }
+}
+
 async fn handle_client(
     stream: UnixStream,
     runtime: Arc<DaemonRuntime>,
-    messages: broadcast::Sender<DaemonMessage>,
+    log_tx: broadcast::Sender<DaemonMessage>,
+    _state_rx: watch::Sender<iota_ipc::StateSnapshot>,
 ) -> Result<()> {
+    let peer = peer_credentials(&stream);
     let (mut reader, mut writer) = stream.into_split();
-    let mut outgoing = messages.subscribe();
-    let initial = DaemonMessage::StateUpdate(runtime.snapshot());
-    write_msg(&mut writer, &initial).await?;
-    let writer_task = tokio::spawn(async move {
-        while let Ok(message) = outgoing.recv().await {
-            if write_msg(&mut writer, &message).await.is_err() {
+    let (directed_tx, directed_rx) = mpsc::channel::<DaemonMessage>(CLIENT_CHANNEL_SIZE);
+
+    // --- Handshake ---
+    let mut negotiated_version: Option<u16> = None;
+    for _ in 0..MAX_HANDSHAKE_RETRIES {
+        match read_msg::<_, ClientMessage>(&mut reader).await {
+            Ok(ClientMessage::Hello { supported_versions }) => {
+                let version = supported_versions
+                    .iter()
+                    .copied()
+                    .find(|v| *v >= MIN_PROTOCOL_VERSION && *v <= PROTOCOL_VERSION)
+                    .unwrap_or(PROTOCOL_VERSION);
+                negotiated_version = Some(version);
+                let instance_id = Uuid::new_v4().to_string();
+                let ack = DaemonMessage::HelloAck(HelloAck {
+                    protocol_version: version,
+                    daemon_version: env!("CARGO_PKG_VERSION").to_string(),
+                    instance_id,
+                    startup_phase: runtime.current_startup_phase().into(),
+                    capabilities: vec!["commands".into(), "metrics".into(), "logs".into()],
+                });
+                write_msg(&mut writer, &ack).await?;
                 break;
             }
+            Ok(_) => {
+                // Unexpected first message — send error and close.
+                return Err(std::io::Error::new(
+                    std::io::ErrorKind::InvalidData,
+                    "Expected Hello as first message",
+                ));
+            }
+            Err(e) => return Err(e),
         }
-    });
+    }
+    let _version = negotiated_version.ok_or_else(|| {
+        std::io::Error::new(std::io::ErrorKind::Other, "Handshake failed after retries")
+    })?;
+
+    log!("IPC client connected (pid={}, uid={})", peer.pid, peer.uid);
+
+    // --- Send initial state snapshot ---
+    let initial = DaemonMessage::StateUpdate(runtime.snapshot());
+    let _ = directed_tx.send(initial).await;
+
+    // --- Writer task: merge directed responses + shared log events ---
+    let mut log_rx = log_tx.subscribe();
+    let directed_for_writer = directed_tx.clone();
+    let writer_task = {
+        let runtime = runtime.clone();
+        tokio::spawn(async move {
+            let mut directed_rx = directed_rx;
+            loop {
+                tokio::select! {
+                    // Directed messages (responses to this client's requests)
+                    msg = directed_rx.recv() => {
+                        match msg {
+                            Some(message) => {
+                                if write_msg(&mut writer, &message).await.is_err() {
+                                    break;
+                                }
+                            }
+                            None => break,
+                        }
+                    }
+                    // Shared log events
+                    result = log_rx.recv() => {
+                        match result {
+                            Ok(message) => {
+                                if write_msg(&mut writer, &message).await.is_err() {
+                                    break;
+                                }
+                            }
+                            Err(broadcast::error::RecvError::Lagged(skipped)) => {
+                                let _ = directed_for_writer.send(DaemonMessage::Gap { skipped }).await;
+                                // Then send current snapshot for resync
+                                let _ = directed_for_writer.send(
+                                    DaemonMessage::StateUpdate(runtime.snapshot())
+                                ).await;
+                            }
+                            Err(broadcast::error::RecvError::Closed) => break,
+                        }
+                    }
+                }
+            }
+        })
+    };
+
+    // --- Reader loop ---
     let router = CommandRouter::new(runtime.clone());
     loop {
         match read_msg::<_, ClientMessage>(&mut reader).await {
-            Ok(ClientMessage::Command { seq, line }) => {
-                let result = router.route(seq, line).await;
-                let _ = messages.send(result);
+            Ok(ClientMessage::Request(envelope)) => {
+                let response = router.route(envelope.request_id, envelope.request).await;
+                let _ = directed_tx.send(DaemonMessage::Response(response)).await;
             }
-            Ok(ClientMessage::Subscribe) => {
-                let _ = messages.send(DaemonMessage::StateUpdate(runtime.snapshot()));
+            Ok(ClientMessage::Subscribe { .. }) => {
+                let snapshot = DaemonMessage::StateUpdate(runtime.snapshot());
+                let _ = directed_tx.send(snapshot).await;
             }
             Ok(ClientMessage::Ping { seq }) => {
-                let _ = messages.send(DaemonMessage::Pong { seq });
+                let _ = directed_tx.send(DaemonMessage::Pong { seq }).await;
+            }
+            Ok(ClientMessage::Hello { .. }) => {
+                // Re-handshake on existing connection: treat as resubscribe
+                let snapshot = DaemonMessage::StateUpdate(runtime.snapshot());
+                let _ = directed_tx.send(snapshot).await;
             }
             Err(error) if error.kind() == std::io::ErrorKind::UnexpectedEof => break,
             Err(error) => {
@@ -109,5 +251,10 @@ async fn handle_client(
         }
     }
     writer_task.abort();
+    log!(
+        "IPC client disconnected (pid={}, uid={})",
+        peer.pid,
+        peer.uid
+    );
     Ok(())
 }

iota-daemon-lib/src/lib.rs

@@ -4,5 +4,5 @@ pub mod ipc_server;
 pub mod log_broadcaster;
 
 pub use command_router::CommandRouter;
-pub use daemon_state::DaemonRuntime;
+pub use daemon_state::{DaemonRuntime, ShutdownReason, StartupPhase};
 pub use ipc_server::IpcServer;