krata: rework into daemon / controller structure

shared/Cargo.toml

@@ -8,6 +8,8 @@ resolver = "2"
 anyhow = { workspace = true }
 serde = { workspace = true }
 libc = { workspace = true }
+log = { workspace = true }
+tokio = { workspace = true }
 
 [dependencies.nix]
 workspace = true

shared/src/control.rs

@@ -0,0 +1,115 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct GuestInfo {
+    pub id: String,
+    pub image: String,
+    pub ipv4: Option<String>,
+    pub ipv6: Option<String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct LaunchRequest {
+    pub image: String,
+    pub vcpus: u32,
+    pub mem: u64,
+    pub env: Option<Vec<String>>,
+    pub run: Option<Vec<String>>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct LaunchResponse {
+    pub guest: GuestInfo,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ListRequest {}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ListResponse {
+    pub guests: Vec<GuestInfo>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct DestroyRequest {
+    pub guest: String,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct DestroyResponse {
+    pub guest: String,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ConsoleStreamRequest {
+    pub guest: String,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ConsoleStreamResponse {
+    pub stream: u64,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ConsoleStreamUpdate {
+    pub data: Vec<u8>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ErrorResponse {
+    pub message: String,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum Request {
+    Launch(LaunchRequest),
+    Destroy(DestroyRequest),
+    List(ListRequest),
+    ConsoleStream(ConsoleStreamRequest),
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum Response {
+    Error(ErrorResponse),
+    Launch(LaunchResponse),
+    Destroy(DestroyResponse),
+    List(ListResponse),
+    ConsoleStream(ConsoleStreamResponse),
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct RequestBox {
+    pub id: u64,
+    pub request: Request,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ResponseBox {
+    pub id: u64,
+    pub response: Response,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+pub enum StreamStatus {
+    Open,
+    Closed,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum StreamUpdate {
+    ConsoleStream(ConsoleStreamUpdate),
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct StreamUpdated {
+    pub id: u64,
+    pub update: Option<StreamUpdate>,
+    pub status: StreamStatus,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum Message {
+    Request(RequestBox),
+    Response(ResponseBox),
+    StreamUpdated(StreamUpdated),
+}

shared/src/launchcfg.rs

@@ -0,0 +1,33 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct LaunchNetworkIpv4 {
+    pub address: String,
+    pub gateway: String,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct LaunchNetworkIpv6 {
+    pub address: String,
+    pub gateway: String,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct LaunchNetworkResolver {
+    pub nameservers: Vec<String>,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct LaunchNetwork {
+    pub link: String,
+    pub ipv4: LaunchNetworkIpv4,
+    pub ipv6: LaunchNetworkIpv6,
+    pub resolver: LaunchNetworkResolver,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct LaunchInfo {
+    pub network: Option<LaunchNetwork>,
+    pub env: Option<Vec<String>>,
+    pub run: Option<Vec<String>>,
+}

shared/src/lib.rs

@@ -1,35 +1,4 @@
+pub mod control;
 pub mod ethtool;
-
-use serde::{Deserialize, Serialize};
-
-#[derive(Serialize, Deserialize, Debug)]
-pub struct LaunchNetworkIpv4 {
-    pub address: String,
-    pub gateway: String,
-}
-
-#[derive(Serialize, Deserialize, Debug)]
-pub struct LaunchNetworkIpv6 {
-    pub address: String,
-    pub gateway: String,
-}
-
-#[derive(Serialize, Deserialize, Debug)]
-pub struct LaunchNetworkResolver {
-    pub nameservers: Vec<String>,
-}
-
-#[derive(Serialize, Deserialize, Debug)]
-pub struct LaunchNetwork {
-    pub link: String,
-    pub ipv4: LaunchNetworkIpv4,
-    pub ipv6: LaunchNetworkIpv6,
-    pub resolver: LaunchNetworkResolver,
-}
-
-#[derive(Serialize, Deserialize, Debug)]
-pub struct LaunchInfo {
-    pub network: Option<LaunchNetwork>,
-    pub env: Option<Vec<String>>,
-    pub run: Option<Vec<String>>,
-}
+pub mod launchcfg;
+pub mod stream;

shared/src/stream.rs

@@ -0,0 +1,152 @@
+use crate::control::{Message, StreamStatus, StreamUpdate, StreamUpdated};
+use anyhow::{anyhow, Result};
+use log::warn;
+use std::{collections::HashMap, sync::Arc};
+use tokio::sync::{
+    mpsc::{channel, Receiver, Sender},
+    Mutex,
+};
+
+pub struct StreamContext {
+    pub id: u64,
+    pub receiver: Receiver<StreamUpdate>,
+    sender: Sender<Message>,
+}
+
+impl StreamContext {
+    pub async fn send(&self, update: StreamUpdate) -> Result<()> {
+        self.sender
+            .send(Message::StreamUpdated(StreamUpdated {
+                id: self.id,
+                update: Some(update),
+                status: StreamStatus::Open,
+            }))
+            .await?;
+        Ok(())
+    }
+}
+
+impl Drop for StreamContext {
+    fn drop(&mut self) {
+        if self.sender.is_closed() {
+            return;
+        }
+        let result = self.sender.try_send(Message::StreamUpdated(StreamUpdated {
+            id: self.id,
+            update: None,
+            status: StreamStatus::Closed,
+        }));
+
+        if let Err(error) = result {
+            warn!(
+                "failed to send close message for stream {}: {}",
+                self.id, error
+            );
+        }
+    }
+}
+
+struct StreamStorage {
+    rx_sender: Sender<StreamUpdate>,
+    rx_receiver: Option<Receiver<StreamUpdate>>,
+}
+
+#[derive(Clone)]
+pub struct ConnectionStreams {
+    next: Arc<Mutex<u64>>,
+    streams: Arc<Mutex<HashMap<u64, StreamStorage>>>,
+    tx_sender: Sender<Message>,
+}
+
+const QUEUE_MAX_LEN: usize = 100;
+
+impl ConnectionStreams {
+    pub fn new(tx_sender: Sender<Message>) -> Self {
+        Self {
+            next: Arc::new(Mutex::new(0)),
+            streams: Arc::new(Mutex::new(HashMap::new())),
+            tx_sender,
+        }
+    }
+
+    pub async fn open(&self) -> Result<StreamContext> {
+        let id = {
+            let mut next = self.next.lock().await;
+            let id = *next;
+            *next = id + 1;
+            id
+        };
+
+        let (rx_sender, rx_receiver) = channel(QUEUE_MAX_LEN);
+        let store = StreamStorage {
+            rx_sender,
+            rx_receiver: None,
+        };
+
+        self.streams.lock().await.insert(id, store);
+
+        let open = Message::StreamUpdated(StreamUpdated {
+            id,
+            update: None,
+            status: StreamStatus::Open,
+        });
+        self.tx_sender.send(open).await?;
+
+        Ok(StreamContext {
+            id,
+            sender: self.tx_sender.clone(),
+            receiver: rx_receiver,
+        })
+    }
+
+    pub async fn incoming(&self, updated: StreamUpdated) -> Result<()> {
+        let mut streams = self.streams.lock().await;
+        if updated.update.is_none() && updated.status == StreamStatus::Open {
+            let (rx_sender, rx_receiver) = channel(QUEUE_MAX_LEN);
+            let store = StreamStorage {
+                rx_sender,
+                rx_receiver: Some(rx_receiver),
+            };
+            streams.insert(updated.id, store);
+        }
+
+        let Some(storage) = streams.get(&updated.id) else {
+            return Ok(());
+        };
+
+        if let Some(update) = updated.update {
+            storage.rx_sender.send(update).await?;
+        }
+
+        if updated.status == StreamStatus::Closed {
+            streams.remove(&updated.id);
+        }
+
+        Ok(())
+    }
+
+    pub async fn outgoing(&self, updated: &StreamUpdated) -> Result<()> {
+        if updated.status == StreamStatus::Closed {
+            let mut streams = self.streams.lock().await;
+            streams.remove(&updated.id);
+        }
+        Ok(())
+    }
+
+    pub async fn acquire(&self, id: u64) -> Result<StreamContext> {
+        let mut streams = self.streams.lock().await;
+        let Some(storage) = streams.get_mut(&id) else {
+            return Err(anyhow!("stream {} has not been opened", id));
+        };
+
+        let Some(receiver) = storage.rx_receiver.take() else {
+            return Err(anyhow!("stream has already been acquired"));
+        };
+
+        Ok(StreamContext {
+            id,
+            receiver,
+            sender: self.tx_sender.clone(),
+        })
+    }
+}