feat: initial support for idm send in daemon

crates/daemon/src/idm.rs

@@ -7,8 +7,9 @@ use kratart::channel::ChannelService;
 use log::{error, warn};
 use prost::Message;
 use tokio::{
+    select,
     sync::{
-        mpsc::{Receiver, Sender},
+        mpsc::{channel, Receiver, Sender},
         Mutex,
     },
     task::JoinHandle,
@@ -19,16 +20,23 @@ type ListenerMap = Arc<Mutex<HashMap<u32, Sender<(u32, IdmPacket)>>>>;
 #[derive(Clone)]
 pub struct DaemonIdmHandle {
     listeners: ListenerMap,
+    tx_sender: Sender<(u32, IdmPacket)>,
     task: Arc<JoinHandle<()>>,
 }
 
 #[derive(Clone)]
 pub struct DaemonIdmSubscribeHandle {
     domid: u32,
+    tx_sender: Sender<(u32, IdmPacket)>,
     listeners: ListenerMap,
 }
 
 impl DaemonIdmSubscribeHandle {
+    pub async fn send(&self, packet: IdmPacket) -> Result<()> {
+        self.tx_sender.send((self.domid, packet)).await?;
+        Ok(())
+    }
+
     pub async fn unsubscribe(&self) -> Result<()> {
         let mut guard = self.listeners.lock().await;
         let _ = guard.remove(&self.domid);
@@ -37,6 +45,11 @@ impl DaemonIdmSubscribeHandle {
 }
 
 impl DaemonIdmHandle {
+    pub async fn send(&self, domid: u32, packet: IdmPacket) -> Result<()> {
+        self.tx_sender.send((domid, packet)).await?;
+        Ok(())
+    }
+
     pub async fn subscribe(
         &self,
         domid: u32,
@@ -46,6 +59,7 @@ impl DaemonIdmHandle {
         guard.insert(domid, sender);
         Ok(DaemonIdmSubscribeHandle {
             domid,
+            tx_sender: self.tx_sender.clone(),
             listeners: self.listeners.clone(),
         })
     }
@@ -61,17 +75,25 @@ impl Drop for DaemonIdmHandle {
 
 pub struct DaemonIdm {
     listeners: ListenerMap,
-    receiver: Receiver<(u32, Vec<u8>)>,
+    tx_sender: Sender<(u32, IdmPacket)>,
+    tx_raw_sender: Sender<(u32, Vec<u8>)>,
+    tx_receiver: Receiver<(u32, IdmPacket)>,
+    rx_receiver: Receiver<(u32, Vec<u8>)>,
     task: JoinHandle<()>,
 }
 
 impl DaemonIdm {
     pub async fn new() -> Result<DaemonIdm> {
-        let (service, _, receiver) = ChannelService::new("krata-channel".to_string(), None).await?;
+        let (service, tx_raw_sender, rx_receiver) =
+            ChannelService::new("krata-channel".to_string(), None).await?;
+        let (tx_sender, tx_receiver) = channel(100);
         let task = service.launch().await?;
         let listeners = Arc::new(Mutex::new(HashMap::new()));
         Ok(DaemonIdm {
-            receiver,
+            rx_receiver,
+            tx_receiver,
+            tx_sender,
+            tx_raw_sender,
             task,
             listeners,
         })
@@ -79,6 +101,7 @@ impl DaemonIdm {
 
     pub async fn launch(mut self) -> Result<DaemonIdmHandle> {
         let listeners = self.listeners.clone();
+        let tx_sender = self.tx_sender.clone();
         let task = tokio::task::spawn(async move {
             let mut buffers: HashMap<u32, BytesMut> = HashMap::new();
             if let Err(error) = self.process(&mut buffers).await {
@@ -87,42 +110,59 @@ impl DaemonIdm {
         });
         Ok(DaemonIdmHandle {
             listeners,
+            tx_sender,
             task: Arc::new(task),
         })
     }
 
     async fn process(&mut self, buffers: &mut HashMap<u32, BytesMut>) -> Result<()> {
         loop {
-            let Some((domid, data)) = self.receiver.recv().await else {
-                break;
-            };
-
-            let buffer = buffers.entry(domid).or_insert_with_key(|_| BytesMut::new());
-            buffer.extend_from_slice(&data);
-            if buffer.len() < 2 {
-                continue;
-            }
-            let size = (buffer[0] as u16 | (buffer[1] as u16) << 8) as usize;
-            let needed = size + 2;
-            if buffer.len() < needed {
-                continue;
-            }
-            let mut packet = buffer.split_to(needed);
-            packet.advance(2);
-            match IdmPacket::decode(packet) {
-                Ok(packet) => {
-                    let guard = self.listeners.lock().await;
-                    if let Some(sender) = guard.get(&domid) {
-                        if let Err(error) = sender.try_send((domid, packet)) {
-                            warn!("dropped idm packet from domain {}: {}", domid, error);
+            select! {
+                x = self.rx_receiver.recv() => match x {
+                    Some((domid, data)) => {
+                        let buffer = buffers.entry(domid).or_insert_with_key(|_| BytesMut::new());
+                        buffer.extend_from_slice(&data);
+                        if buffer.len() < 2 {
+                            continue;
                         }
+                        let size = (buffer[0] as u16 | (buffer[1] as u16) << 8) as usize;
+                        let needed = size + 2;
+                        if buffer.len() < needed {
+                            continue;
+                        }
+                        let mut packet = buffer.split_to(needed);
+                        packet.advance(2);
+                        match IdmPacket::decode(packet) {
+                            Ok(packet) => {
+                                let guard = self.listeners.lock().await;
+                                if let Some(sender) = guard.get(&domid) {
+                                    if let Err(error) = sender.try_send((domid, packet)) {
+                                        warn!("dropped idm packet from domain {}: {}", domid, error);
+                                    }
+                                }
+                            }
+
+                            Err(packet) => {
+                                warn!("received invalid packet from domain {}: {}", domid, packet);
+                            }
+                        }
+                    },
+
+                    None => {
+                        break;
+                    }
+                },
+                x = self.tx_receiver.recv() => match x {
+                    Some((domid, packet)) => {
+                        let data = packet.encode_to_vec();
+                        self.tx_raw_sender.send((domid, data)).await?;
+                    },
+
+                    None => {
+                        break;
                     }
                 }
-
-                Err(packet) => {
-                    warn!("received invalid packet from domain {}: {}", domid, packet);
-                }
-            }
+            };
         }
         Ok(())
     }

crates/guest/bin/init.rs

@@ -23,6 +23,8 @@ async fn main() -> Result<()> {
     if let Err(error) = guest.init().await {
         error!("failed to initialize guest: {}", error);
         death(127).await?;
+        return Ok(());
     }
+    death(1).await?;
     Ok(())
 }

crates/guest/src/background.rs

@@ -8,7 +8,7 @@ use krata::idm::{
     client::IdmClient,
     protocol::{idm_event::Event, IdmEvent, IdmExitEvent, IdmPacket},
 };
-use log::error;
+use log::debug;
 use nix::unistd::Pid;
 use tokio::select;
 
@@ -38,7 +38,7 @@ impl GuestBackground {
                     },
 
                     None => {
-                        error!("idm packet channel closed");
+                        debug!("idm packet channel closed");
                         break;
                     }
                 },

crates/runtime/src/console.rs

@@ -1,18 +0,0 @@
-use anyhow::Result;
-use tokio::fs::File;
-
-pub struct XenConsole {
-    pub read_handle: File,
-    pub write_handle: File,
-}
-
-impl XenConsole {
-    pub async fn new(tty: &str) -> Result<XenConsole> {
-        let read_handle = File::options().read(true).write(false).open(tty).await?;
-        let write_handle = File::options().read(false).write(true).open(tty).await?;
-        Ok(XenConsole {
-            read_handle,
-            write_handle,
-        })
-    }
-}

crates/runtime/src/lib.rs

@@ -15,7 +15,6 @@ use xenstore::{XsdClient, XsdInterface};
 
 use self::{
     autoloop::AutoLoop,
-    console::XenConsole,
     launch::{GuestLaunchRequest, GuestLauncher},
 };
 use krataoci::cache::ImageCache;
@@ -23,7 +22,6 @@ use krataoci::cache::ImageCache;
 pub mod autoloop;
 pub mod cfgblk;
 pub mod channel;
-pub mod console;
 pub mod launch;
 
 pub struct GuestLoopInfo {
@@ -321,17 +319,6 @@ impl Runtime {
         Ok(uuid)
     }
 
-    pub async fn console(&self, uuid: Uuid) -> Result<XenConsole> {
-        let info = self
-            .context
-            .resolve(uuid)
-            .await?
-            .ok_or_else(|| anyhow!("unable to resolve guest: {}", uuid))?;
-        let domid = info.domid;
-        let tty = self.context.xen.get_console_path(domid).await?;
-        XenConsole::new(&tty).await
-    }
-
     pub async fn list(&self) -> Result<Vec<GuestInfo>> {
         self.context.list().await
     }