krata: implement idm support for guest <-> host messages

crates/kratad/Cargo.toml

@@ -8,6 +8,7 @@ resolver = "2"
 anyhow = { workspace = true }
 async-stream = { workspace = true }
 async-trait = { workspace = true }
+bytes = { workspace = true }
 clap = { workspace = true }
 env_logger = { workspace = true }
 futures = { workspace = true }

crates/kratad/src/event.rs

@@ -5,7 +5,10 @@ use std::{
 };
 
 use anyhow::Result;
-use krata::v1::common::{GuestExitInfo, GuestState, GuestStatus};
+use krata::{
+    idm::protocol::{idm_packet::Message, IdmPacket},
+    v1::common::{GuestExitInfo, GuestState, GuestStatus},
+};
 use log::error;
 use tokio::{
     select,
@@ -18,14 +21,15 @@ use tokio::{
 };
 use uuid::Uuid;
 
-use kratart::Runtime;
-
-use crate::db::GuestStore;
+use crate::{
+    db::GuestStore,
+    idm::{DaemonIdmHandle, DaemonIdmSubscribeHandle},
+};
 
 pub type DaemonEvent = krata::v1::control::watch_events_reply::Event;
 
 const EVENT_CHANNEL_QUEUE_LEN: usize = 1000;
-const EXIT_CODE_CHANNEL_QUEUE_LEN: usize = 1000;
+const IDM_CHANNEL_QUEUE_LEN: usize = 1000;
 
 #[derive(Clone)]
 pub struct DaemonEventContext {
@@ -44,13 +48,13 @@ impl DaemonEventContext {
 }
 
 pub struct DaemonEventGenerator {
-    runtime: Runtime,
     guests: GuestStore,
     guest_reconciler_notify: Sender<Uuid>,
     feed: broadcast::Receiver<DaemonEvent>,
-    exit_code_sender: Sender<(Uuid, i32)>,
-    exit_code_receiver: Receiver<(Uuid, i32)>,
-    exit_code_handles: HashMap<Uuid, JoinHandle<()>>,
+    idm: DaemonIdmHandle,
+    idms: HashMap<u32, (Uuid, DaemonIdmSubscribeHandle)>,
+    idm_sender: Sender<(u32, IdmPacket)>,
+    idm_receiver: Receiver<(u32, IdmPacket)>,
     _event_sender: broadcast::Sender<DaemonEvent>,
 }
 
@@ -58,18 +62,18 @@ impl DaemonEventGenerator {
     pub async fn new(
         guests: GuestStore,
         guest_reconciler_notify: Sender<Uuid>,
-        runtime: Runtime,
+        idm: DaemonIdmHandle,
     ) -> Result<(DaemonEventContext, DaemonEventGenerator)> {
         let (sender, _) = broadcast::channel(EVENT_CHANNEL_QUEUE_LEN);
-        let (exit_code_sender, exit_code_receiver) = channel(EXIT_CODE_CHANNEL_QUEUE_LEN);
+        let (idm_sender, idm_receiver) = channel(IDM_CHANNEL_QUEUE_LEN);
         let generator = DaemonEventGenerator {
-            runtime,
             guests,
             guest_reconciler_notify,
             feed: sender.subscribe(),
-            exit_code_receiver,
-            exit_code_sender,
-            exit_code_handles: HashMap::new(),
+            idm,
+            idms: HashMap::new(),
+            idm_sender,
+            idm_receiver,
             _event_sender: sender.clone(),
         };
         let context = DaemonEventContext { sender };
@@ -89,20 +93,19 @@ impl DaemonEventGenerator {
 
                 let status = state.status();
                 let id = Uuid::from_str(&guest.id)?;
+                let domid = state.domid;
                 match status {
                     GuestStatus::Started => {
-                        if let Entry::Vacant(e) = self.exit_code_handles.entry(id) {
-                            let handle = self
-                                .runtime
-                                .subscribe_exit_code(id, self.exit_code_sender.clone())
-                                .await?;
-                            e.insert(handle);
+                        if let Entry::Vacant(e) = self.idms.entry(domid) {
+                            let subscribe =
+                                self.idm.subscribe(domid, self.idm_sender.clone()).await?;
+                            e.insert((id, subscribe));
                         }
                     }
 
                     GuestStatus::Destroyed => {
-                        if let Some(handle) = self.exit_code_handles.remove(&id) {
-                            handle.abort();
+                        if let Some((_, handle)) = self.idms.remove(&domid) {
+                            handle.unsubscribe().await?;
                         }
                     }
 
@@ -113,6 +116,13 @@ impl DaemonEventGenerator {
         Ok(())
     }
 
+    async fn handle_idm_packet(&mut self, id: Uuid, packet: IdmPacket) -> Result<()> {
+        if let Some(Message::Exit(exit)) = packet.message {
+            self.handle_exit_code(id, exit.code).await?;
+        }
+        Ok(())
+    }
+
     async fn handle_exit_code(&mut self, id: Uuid, code: i32) -> Result<()> {
         if let Some(mut entry) = self.guests.read(id).await? {
             let Some(ref mut guest) = entry.guest else {
@@ -135,9 +145,12 @@ impl DaemonEventGenerator {
 
     async fn evaluate(&mut self) -> Result<()> {
         select! {
-            x = self.exit_code_receiver.recv() => match x {
-                Some((uuid, code)) => {
-                    self.handle_exit_code(uuid, code).await
+            x = self.idm_receiver.recv() => match x {
+                Some((domid, packet)) => {
+                    if let Some((id, _)) = self.idms.get(&domid) {
+                        self.handle_idm_packet(*id, packet).await?;
+                    }
+                    Ok(())
                 },
                 None => {
                     Ok(())

crates/kratad/src/idm.rs

@@ -1,9 +1,66 @@
+use std::{collections::HashMap, sync::Arc};
+
 use anyhow::Result;
+use bytes::{Buf, BytesMut};
+use krata::idm::protocol::IdmPacket;
 use kratart::channel::ChannelService;
-use log::error;
-use tokio::{sync::mpsc::Receiver, task::JoinHandle};
+use log::{error, warn};
+use prost::Message;
+use tokio::{
+    sync::{
+        mpsc::{Receiver, Sender},
+        Mutex,
+    },
+    task::JoinHandle,
+};
+
+type ListenerMap = Arc<Mutex<HashMap<u32, Sender<(u32, IdmPacket)>>>>;
+
+#[derive(Clone)]
+pub struct DaemonIdmHandle {
+    listeners: ListenerMap,
+    task: Arc<JoinHandle<()>>,
+}
+
+#[derive(Clone)]
+pub struct DaemonIdmSubscribeHandle {
+    domid: u32,
+    listeners: ListenerMap,
+}
+
+impl DaemonIdmSubscribeHandle {
+    pub async fn unsubscribe(&self) -> Result<()> {
+        let mut guard = self.listeners.lock().await;
+        let _ = guard.remove(&self.domid);
+        Ok(())
+    }
+}
+
+impl DaemonIdmHandle {
+    pub async fn subscribe(
+        &self,
+        domid: u32,
+        sender: Sender<(u32, IdmPacket)>,
+    ) -> Result<DaemonIdmSubscribeHandle> {
+        let mut guard = self.listeners.lock().await;
+        guard.insert(domid, sender);
+        Ok(DaemonIdmSubscribeHandle {
+            domid,
+            listeners: self.listeners.clone(),
+        })
+    }
+}
+
+impl Drop for DaemonIdmHandle {
+    fn drop(&mut self) {
+        if Arc::strong_count(&self.task) <= 1 {
+            self.task.abort();
+        }
+    }
+}
 
 pub struct DaemonIdm {
+    listeners: ListenerMap,
     receiver: Receiver<(u32, Vec<u8>)>,
     task: JoinHandle<()>,
 }
@@ -12,22 +69,60 @@ impl DaemonIdm {
     pub async fn new() -> Result<DaemonIdm> {
         let (service, receiver) = ChannelService::new("krata-channel".to_string()).await?;
         let task = service.launch().await?;
-        Ok(DaemonIdm { receiver, task })
+        let listeners = Arc::new(Mutex::new(HashMap::new()));
+        Ok(DaemonIdm {
+            receiver,
+            task,
+            listeners,
+        })
     }
 
-    pub async fn launch(mut self) -> Result<JoinHandle<()>> {
-        Ok(tokio::task::spawn(async move {
-            if let Err(error) = self.process().await {
+    pub async fn launch(mut self) -> Result<DaemonIdmHandle> {
+        let listeners = self.listeners.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 {
                 error!("failed to process idm: {}", error);
             }
-        }))
+        });
+        Ok(DaemonIdmHandle {
+            listeners,
+            task: Arc::new(task),
+        })
     }
 
-    async fn process(&mut self) -> Result<()> {
+    async fn process(&mut self, buffers: &mut HashMap<u32, BytesMut>) -> Result<()> {
         loop {
-            let Some(_) = self.receiver.recv().await else {
+            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);
+                        }
+                    }
+                }
+
+                Err(packet) => {
+                    warn!("received invalid packet from domain {}: {}", domid, packet);
+                }
+            }
         }
         Ok(())
     }

crates/kratad/src/lib.rs

@@ -4,7 +4,7 @@ use anyhow::Result;
 use control::RuntimeControlService;
 use db::GuestStore;
 use event::{DaemonEventContext, DaemonEventGenerator};
-use idm::DaemonIdm;
+use idm::{DaemonIdm, DaemonIdmHandle};
 use krata::{dial::ControlDialAddress, v1::control::control_service_server::ControlServiceServer};
 use kratart::Runtime;
 use log::info;
@@ -32,7 +32,7 @@ pub struct Daemon {
     guest_reconciler_task: JoinHandle<()>,
     guest_reconciler_notify: Sender<Uuid>,
     generator_task: JoinHandle<()>,
-    idm_task: JoinHandle<()>,
+    _idm: DaemonIdmHandle,
 }
 
 const GUEST_RECONCILER_QUEUE_LEN: usize = 1000;
@@ -41,22 +41,18 @@ impl Daemon {
     pub async fn new(store: String, runtime: Runtime) -> Result<Self> {
         let guests_db_path = format!("{}/guests.db", store);
         let guests = GuestStore::open(&PathBuf::from(guests_db_path))?;
-        let runtime_for_events = runtime.dupe().await?;
         let (guest_reconciler_notify, guest_reconciler_receiver) =
             channel::<Uuid>(GUEST_RECONCILER_QUEUE_LEN);
-        let (events, generator) = DaemonEventGenerator::new(
-            guests.clone(),
-            guest_reconciler_notify.clone(),
-            runtime_for_events,
-        )
-        .await?;
+        let idm = DaemonIdm::new().await?;
+        let idm = idm.launch().await?;
+        let (events, generator) =
+            DaemonEventGenerator::new(guests.clone(), guest_reconciler_notify.clone(), idm.clone())
+                .await?;
         let runtime_for_reconciler = runtime.dupe().await?;
         let guest_reconciler =
             GuestReconciler::new(guests.clone(), events.clone(), runtime_for_reconciler)?;
 
         let guest_reconciler_task = guest_reconciler.launch(guest_reconciler_receiver).await?;
-        let idm = DaemonIdm::new().await?;
-        let idm_task = idm.launch().await?;
         let generator_task = generator.launch().await?;
         Ok(Self {
             store,
@@ -66,7 +62,7 @@ impl Daemon {
             guest_reconciler_task,
             guest_reconciler_notify,
             generator_task,
-            idm_task,
+            _idm: idm,
         })
     }
 
@@ -130,6 +126,5 @@ impl Drop for Daemon {
     fn drop(&mut self) {
         self.guest_reconciler_task.abort();
         self.generator_task.abort();
-        self.idm_task.abort();
     }
 }