feat: rebuild idm to separate transport from content

This commit is contained in:
Alex Zenla
2024-04-21 06:29:32 +00:00
parent f843abcabf
commit 4e0d843de7
24 changed files with 483 additions and 298 deletions

View File

@ -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 }

View File

@ -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: u32,
pub to: u32,
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,
})
}

View File

@ -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('"', "\\\"")))