feat: idm v2 (#102)

* feat: rebuild idm to separate transport from content

* feat: fast guest lookup table and host identification

crates/ctl/Cargo.toml

@@ -11,6 +11,7 @@ resolver = "2"
 [dependencies]
 anyhow = { workspace = true }
 async-stream = { workspace = true }
+base64 = { workspace = true }
 clap = { workspace = true }
 comfy-table = { workspace = true }
 crossterm = { workspace = true, features = ["event-stream"] }
@@ -24,6 +25,7 @@ log = { workspace = true }
 prost-reflect = { workspace = true, features = ["serde"] }
 prost-types = { workspace = true }
 ratatui = { workspace = true }
+serde = { workspace = true }
 serde_json = { workspace = true }
 serde_yaml = { workspace = true }
 termtree = { workspace = true }

crates/ctl/src/cli/identify_host.rs

@@ -0,0 +1,22 @@
+use anyhow::Result;
+use clap::Parser;
+use krata::v1::control::{control_service_client::ControlServiceClient, IdentifyHostRequest};
+
+use tonic::{transport::Channel, Request};
+
+#[derive(Parser)]
+#[command(about = "Identify information about the host")]
+pub struct IdentifyHostCommand {}
+
+impl IdentifyHostCommand {
+    pub async fn run(self, mut client: ControlServiceClient<Channel>) -> Result<()> {
+        let response = client
+            .identify_host(Request::new(IdentifyHostRequest {}))
+            .await?
+            .into_inner();
+        println!("Host UUID: {}", response.host_uuid);
+        println!("Host Domain: {}", response.host_domid);
+        println!("Krata Version: {}", response.krata_version);
+        Ok(())
+    }
+}

crates/ctl/src/cli/idm_snoop.rs

@@ -1,14 +1,18 @@
 use anyhow::Result;
+use base64::Engine;
 use clap::{Parser, ValueEnum};
 use krata::{
     events::EventStream,
+    idm::{internal, serialize::IdmSerializable, transport::IdmTransportPacketForm},
     v1::control::{control_service_client::ControlServiceClient, SnoopIdmReply, SnoopIdmRequest},
 };
 
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
 use tokio_stream::StreamExt;
 use tonic::transport::Channel;
 
-use crate::format::{kv2line, proto2dynamic, proto2kv};
+use crate::format::{kv2line, proto2dynamic, value2kv};
 
 #[derive(ValueEnum, Clone, Debug, PartialEq, Eq)]
 enum IdmSnoopFormat {
@@ -34,19 +38,22 @@ impl IdmSnoopCommand {
 
         while let Some(reply) = stream.next().await {
             let reply = reply?;
+            let Some(line) = convert_idm_snoop(reply) else {
+                continue;
+            };
+
             match self.format {
                 IdmSnoopFormat::Simple => {
-                    self.print_simple(reply)?;
+                    self.print_simple(line)?;
                 }
 
                 IdmSnoopFormat::Jsonl => {
-                    let value = serde_json::to_value(proto2dynamic(reply)?)?;
-                    let encoded = serde_json::to_string(&value)?;
+                    let encoded = serde_json::to_string(&line)?;
                     println!("{}", encoded.trim());
                 }
 
                 IdmSnoopFormat::KeyValue => {
-                    self.print_key_value(reply)?;
+                    self.print_key_value(line)?;
                 }
             }
         }
@@ -54,21 +61,86 @@ impl IdmSnoopCommand {
         Ok(())
     }
 
-    fn print_simple(&self, reply: SnoopIdmReply) -> Result<()> {
-        let from = reply.from;
-        let to = reply.to;
-        let Some(packet) = reply.packet else {
-            return Ok(());
+    fn print_simple(&self, line: IdmSnoopLine) -> Result<()> {
+        let encoded = if !line.packet.decoded.is_null() {
+            serde_json::to_string(&line.packet.decoded)?
+        } else {
+            base64::prelude::BASE64_STANDARD.encode(&line.packet.data)
         };
-        let value = serde_json::to_value(proto2dynamic(packet)?)?;
-        let encoded = serde_json::to_string(&value)?;
-        println!("({} -> {}) {}", from, to, encoded);
+        println!(
+            "({} -> {}) {} {} {}",
+            line.from, line.to, line.packet.id, line.packet.form, encoded
+        );
         Ok(())
     }
 
-    fn print_key_value(&self, reply: SnoopIdmReply) -> Result<()> {
-        let kvs = proto2kv(reply)?;
+    fn print_key_value(&self, line: IdmSnoopLine) -> Result<()> {
+        let kvs = value2kv(serde_json::to_value(line)?)?;
         println!("{}", kv2line(kvs));
         Ok(())
     }
 }
+
+#[derive(Serialize, Deserialize)]
+pub struct IdmSnoopLine {
+    pub from: String,
+    pub to: String,
+    pub packet: IdmSnoopData,
+}
+
+#[derive(Serialize, Deserialize)]
+pub struct IdmSnoopData {
+    pub id: u64,
+    pub channel: u64,
+    pub form: String,
+    pub data: String,
+    pub decoded: Value,
+}
+
+pub fn convert_idm_snoop(reply: SnoopIdmReply) -> Option<IdmSnoopLine> {
+    let packet = &(reply.packet?);
+
+    let decoded = if packet.channel == 0 {
+        match packet.form() {
+            IdmTransportPacketForm::Event => internal::Event::decode(&packet.data)
+                .ok()
+                .and_then(|event| proto2dynamic(event).ok()),
+
+            IdmTransportPacketForm::Request => internal::Request::decode(&packet.data)
+                .ok()
+                .and_then(|event| proto2dynamic(event).ok()),
+
+            IdmTransportPacketForm::Response => internal::Response::decode(&packet.data)
+                .ok()
+                .and_then(|event| proto2dynamic(event).ok()),
+
+            _ => None,
+        }
+    } else {
+        None
+    };
+
+    let decoded = decoded
+        .and_then(|message| serde_json::to_value(message).ok())
+        .unwrap_or(Value::Null);
+
+    let data = IdmSnoopData {
+        id: packet.id,
+        channel: packet.channel,
+        form: match packet.form() {
+            IdmTransportPacketForm::Raw => "raw".to_string(),
+            IdmTransportPacketForm::Event => "event".to_string(),
+            IdmTransportPacketForm::Request => "request".to_string(),
+            IdmTransportPacketForm::Response => "response".to_string(),
+            _ => format!("unknown-{}", packet.form),
+        },
+        data: base64::prelude::BASE64_STANDARD.encode(&packet.data),
+        decoded,
+    };
+
+    Some(IdmSnoopLine {
+        from: reply.from,
+        to: reply.to,
+        packet: data,
+    })
+}

crates/ctl/src/cli/mod.rs

@@ -1,5 +1,6 @@
 pub mod attach;
 pub mod destroy;
+pub mod identify_host;
 pub mod idm_snoop;
 pub mod launch;
 pub mod list;
@@ -20,9 +21,10 @@ use krata::{
 use tonic::{transport::Channel, Request};
 
 use self::{
-    attach::AttachCommand, destroy::DestroyCommand, idm_snoop::IdmSnoopCommand,
-    launch::LauchCommand, list::ListCommand, logs::LogsCommand, metrics::MetricsCommand,
-    pull::PullCommand, resolve::ResolveCommand, top::TopCommand, watch::WatchCommand,
+    attach::AttachCommand, destroy::DestroyCommand, identify_host::IdentifyHostCommand,
+    idm_snoop::IdmSnoopCommand, launch::LauchCommand, list::ListCommand, logs::LogsCommand,
+    metrics::MetricsCommand, pull::PullCommand, resolve::ResolveCommand, top::TopCommand,
+    watch::WatchCommand,
 };
 
 #[derive(Parser)]
@@ -56,6 +58,7 @@ pub enum Commands {
     Metrics(MetricsCommand),
     IdmSnoop(IdmSnoopCommand),
     Top(TopCommand),
+    IdentifyHost(IdentifyHostCommand),
 }
 
 impl ControlCommand {
@@ -107,6 +110,10 @@ impl ControlCommand {
             Commands::Pull(pull) => {
                 pull.run(client).await?;
             }
+
+            Commands::IdentifyHost(identify) => {
+                identify.run(client).await?;
+            }
         }
         Ok(())
     }

crates/ctl/src/format.rs

@@ -4,7 +4,7 @@ use anyhow::Result;
 use fancy_duration::FancyDuration;
 use human_bytes::human_bytes;
 use krata::v1::common::{Guest, GuestMetricFormat, GuestMetricNode, GuestStatus};
-use prost_reflect::{DynamicMessage, FieldDescriptor, ReflectMessage, Value as ReflectValue};
+use prost_reflect::{DynamicMessage, ReflectMessage};
 use prost_types::Value;
 use termtree::Tree;
 
@@ -15,64 +15,59 @@ pub fn proto2dynamic(proto: impl ReflectMessage) -> Result<DynamicMessage> {
     )?)
 }
 
-pub fn proto2kv(proto: impl ReflectMessage) -> Result<HashMap<String, String>> {
-    let message = proto2dynamic(proto)?;
+pub fn value2kv(value: serde_json::Value) -> Result<HashMap<String, String>> {
     let mut map = HashMap::new();
+    fn crawl(prefix: String, map: &mut HashMap<String, String>, value: serde_json::Value) {
+        fn dot(prefix: &str, next: String) -> String {
+            if prefix.is_empty() {
+                next.to_string()
+            } else {
+                format!("{}.{}", prefix, next)
+            }
+        }
 
-    fn crawl(
-        prefix: String,
-        field: Option<&FieldDescriptor>,
-        map: &mut HashMap<String, String>,
-        value: &ReflectValue,
-    ) {
         match value {
-            ReflectValue::Message(child) => {
-                for (field, field_value) in child.fields() {
-                    let path = if prefix.is_empty() {
-                        field.json_name().to_string()
-                    } else {
-                        format!("{}.{}", prefix, field.json_name())
-                    };
-                    crawl(path, Some(&field), map, field_value);
+            serde_json::Value::Null => {
+                map.insert(prefix, "null".to_string());
+            }
+
+            serde_json::Value::String(value) => {
+                map.insert(prefix, value);
+            }
+
+            serde_json::Value::Bool(value) => {
+                map.insert(prefix, value.to_string());
+            }
+
+            serde_json::Value::Number(value) => {
+                map.insert(prefix, value.to_string());
+            }
+
+            serde_json::Value::Array(value) => {
+                for (i, item) in value.into_iter().enumerate() {
+                    let next = dot(&prefix, i.to_string());
+                    crawl(next, map, item);
                 }
             }
 
-            ReflectValue::EnumNumber(number) => {
-                if let Some(kind) = field.map(|x| x.kind()) {
-                    if let Some(e) = kind.as_enum() {
-                        if let Some(value) = e.get_value(*number) {
-                            map.insert(prefix, value.name().to_string());
-                        }
-                    }
+            serde_json::Value::Object(value) => {
+                for (key, item) in value {
+                    let next = dot(&prefix, key);
+                    crawl(next, map, item);
                 }
             }
-
-            ReflectValue::String(value) => {
-                map.insert(prefix.to_string(), value.clone());
-            }
-
-            ReflectValue::List(value) => {
-                for (x, value) in value.iter().enumerate() {
-                    crawl(format!("{}.{}", prefix, x), field, map, value);
-                }
-            }
-
-            _ => {
-                map.insert(prefix.to_string(), value.to_string());
-            }
         }
     }
-
-    crawl(
-        "".to_string(),
-        None,
-        &mut map,
-        &ReflectValue::Message(message),
-    );
-
+    crawl("".to_string(), &mut map, value);
     Ok(map)
 }
 
+pub fn proto2kv(proto: impl ReflectMessage) -> Result<HashMap<String, String>> {
+    let message = proto2dynamic(proto)?;
+    let value = serde_json::to_value(message)?;
+    value2kv(value)
+}
+
 pub fn kv2line(map: HashMap<String, String>) -> String {
     map.iter()
         .map(|(k, v)| format!("{}=\"{}\"", k, v.replace('"', "\\\"")))