krata: utilize gRPC for control service

shared/Cargo.toml

@@ -10,6 +10,12 @@ serde = { workspace = true }
 libc = { workspace = true }
 log = { workspace = true }
 tokio = { workspace = true }
+url = { workspace = true }
+tonic = { workspace = true }
+prost = { workspace = true }
+
+[build-dependencies]
+tonic-build = { workspace = true }
 
 [dependencies.nix]
 workspace = true

shared/build.rs

@@ -0,0 +1,5 @@
+fn main() {
+    tonic_build::configure()
+        .compile(&["proto/krata/control.proto"], &["proto"])
+        .unwrap();
+}

shared/proto/krata/control.proto

@@ -0,0 +1,56 @@
+syntax = "proto3";
+
+option java_multiple_files = true;
+option java_package = "dev.krata.proto.control";
+option java_outer_classname = "ControlProto";
+
+package krata.control;
+
+message GuestInfo {
+    string id = 1;
+    string image = 2;
+    string ipv4 = 3;
+    string ipv6 = 4;
+}
+
+message LaunchGuestRequest {
+    string image = 1;
+    uint32 vcpus = 2;
+    uint64 mem = 3;
+    repeated string env = 4;
+    repeated string run = 5;
+}
+
+message LaunchGuestReply {
+    GuestInfo guest = 1;
+}
+
+message ListGuestsRequest {}
+
+message ListGuestsReply {
+    repeated GuestInfo guests = 1;
+}
+
+message DestroyGuestRequest {
+    string guest_id = 1;
+}
+
+message DestroyGuestReply {}
+
+message ConsoleDataRequest {
+    string guest = 1;
+    bytes data = 2;
+}
+
+message ConsoleDataReply {
+    bytes data = 1;
+}
+
+service ControlService {
+    rpc LaunchGuest(LaunchGuestRequest) returns (LaunchGuestReply);
+    rpc DestroyGuest(DestroyGuestRequest) returns (DestroyGuestReply);
+
+    rpc ListGuests(ListGuestsRequest) returns (ListGuestsReply);
+
+    rpc ConsoleData(stream ConsoleDataRequest) returns (stream ConsoleDataReply);
+}

shared/src/control.rs

@@ -1,115 +1 @@
-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),
-}
+tonic::include_proto!("krata.control");

shared/src/dial.rs

@@ -0,0 +1,100 @@
+use std::{fmt::Display, str::FromStr};
+
+use anyhow::anyhow;
+use url::{Host, Url};
+
+pub const KRATA_DEFAULT_TCP_PORT: u16 = 4350;
+pub const KRATA_DEFAULT_TLS_PORT: u16 = 4353;
+
+#[derive(Clone)]
+pub enum ControlDialAddress {
+    UnixSocket {
+        path: String,
+    },
+    Tcp {
+        host: String,
+        port: u16,
+    },
+    Tls {
+        host: String,
+        port: u16,
+        insecure: bool,
+    },
+}
+
+impl FromStr for ControlDialAddress {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let url: Url = s.parse()?;
+
+        let host = url.host().unwrap_or(Host::Domain("localhost")).to_string();
+
+        match url.scheme() {
+            "unix" => Ok(ControlDialAddress::UnixSocket {
+                path: url.path().to_string(),
+            }),
+
+            "tcp" => {
+                let port = url.port().unwrap_or(KRATA_DEFAULT_TCP_PORT);
+                Ok(ControlDialAddress::Tcp { host, port })
+            }
+
+            "tls" | "tls-insecure" => {
+                let insecure = url.scheme() == "tls-insecure";
+                let port = url.port().unwrap_or(KRATA_DEFAULT_TLS_PORT);
+                Ok(ControlDialAddress::Tls {
+                    host,
+                    port,
+                    insecure,
+                })
+            }
+
+            _ => Err(anyhow!("unknown control address scheme: {}", url.scheme())),
+        }
+    }
+}
+
+impl From<ControlDialAddress> for Url {
+    fn from(val: ControlDialAddress) -> Self {
+        match val {
+            ControlDialAddress::UnixSocket { path } => {
+                let mut url = Url::parse("unix:///").unwrap();
+                url.set_path(&path);
+                url
+            }
+
+            ControlDialAddress::Tcp { host, port } => {
+                let mut url = Url::parse("tcp://").unwrap();
+                url.set_host(Some(&host)).unwrap();
+                if port != KRATA_DEFAULT_TCP_PORT {
+                    url.set_port(Some(port)).unwrap();
+                }
+                url
+            }
+
+            ControlDialAddress::Tls {
+                host,
+                port,
+                insecure,
+            } => {
+                let mut url = Url::parse("tls://").unwrap();
+                if insecure {
+                    url.set_scheme("tls-insecure").unwrap();
+                }
+                url.set_host(Some(&host)).unwrap();
+                if port != KRATA_DEFAULT_TLS_PORT {
+                    url.set_port(Some(port)).unwrap();
+                }
+                url
+            }
+        }
+    }
+}
+
+impl Display for ControlDialAddress {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let url: Url = self.clone().into();
+        write!(f, "{}", url)
+    }
+}

shared/src/lib.rs

@@ -1,7 +1,4 @@
 pub mod control;
+pub mod dial;
 pub mod ethtool;
 pub mod launchcfg;
-pub mod stream;
-
-pub const KRATA_DEFAULT_TCP_PORT: u16 = 4350;
-pub const KRATA_DEFAULT_TLS_PORT: u16 = 4353;

shared/src/stream.rs

@@ -1,152 +0,0 @@
-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(),
-        })
-    }
-}