From f7ef2fae8589cab67a5755e48114af43f39ab2e0 Mon Sep 17 00:00:00 2001 From: FujiApple Date: Sat, 2 Mar 2024 16:51:16 +0800 Subject: [PATCH] feat(report): enhanced dot `report` (#797) - WIP --- Cargo.lock | 17 ---- Cargo.toml | 1 - src/main.rs | 2 +- src/report/dot.rs | 198 +++++++++++++++++++++++++++++++++++++++------- 4 files changed, 170 insertions(+), 48 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 96ecaadb..70f2a7a2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -596,12 +596,6 @@ dependencies = [ "windows-sys 0.48.0", ] -[[package]] -name = "fixedbitset" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80" - [[package]] name = "fnv" version = "1.0.7" @@ -1187,16 +1181,6 @@ version = "2.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" -[[package]] -name = "petgraph" -version = "0.6.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e1d3afd2628e69da2be385eb6f2fd57c8ac7977ceeff6dc166ff1657b0e386a9" -dependencies = [ - "fixedbitset", - "indexmap 2.2.5", -] - [[package]] name = "pin-project-lite" version = "0.2.13" @@ -1981,7 +1965,6 @@ dependencies = [ "nix", "parking_lot", "paste", - "petgraph", "pretty_assertions", "rand", "ratatui", diff --git a/Cargo.toml b/Cargo.toml index a9ebb182..699c572f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -51,7 +51,6 @@ indexmap = { version = "2.2.5", default-features = false } maxminddb = "0.24.0" tracing-subscriber = { version = "0.3.18", default-features = false, features = [ "json", "env-filter" ] } tracing-chrome = "0.7.1" -petgraph = "0.6.4" csv = "1.3.0" serde_with = "3.6.1" encoding_rs_io = "0.1.7" diff --git a/src/main.rs b/src/main.rs index 5c065629..283eb6e2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -247,7 +247,7 @@ fn run_frontend( Mode::Json => report::json::report(&traces[0], args.report_cycles, &resolver)?, Mode::Pretty => report::table::report_pretty(&traces[0], args.report_cycles, &resolver)?, Mode::Markdown => report::table::report_md(&traces[0], args.report_cycles, &resolver)?, - Mode::Dot => report::dot::report(&traces[0], args.report_cycles)?, + Mode::Dot => report::dot::report(&traces[0], args.report_cycles, &resolver)?, Mode::Flows => report::flows::report(&traces[0], args.report_cycles)?, Mode::Silent => report::silent::report(&traces[0], args.report_cycles)?, } diff --git a/src/report/dot.rs b/src/report/dot.rs index 12e6eb71..b3d77a27 100644 --- a/src/report/dot.rs +++ b/src/report/dot.rs @@ -1,38 +1,178 @@ -use crate::backend::flows::FlowEntry; +use crate::backend::flows::{Flow, FlowEntry, FlowId}; use crate::TraceInfo; -use petgraph::dot::{Config, Dot}; -use petgraph::graphmap::DiGraphMap; -use std::fmt::{Debug, Formatter}; +use std::collections::{HashMap, HashSet}; +use std::fmt::Debug; use std::net::{IpAddr, Ipv4Addr}; +use trippy::dns::{AsInfo, DnsEntry, DnsResolver, Resolved, Resolver, Unresolved}; /// Run a trace and generate a dot file. -pub fn report(info: &TraceInfo, report_cycles: usize) -> anyhow::Result<()> { - struct DotWrapper<'a>(Dot<'a, &'a DiGraphMap>); - impl Debug for DotWrapper<'_> { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - self.0.fmt(f) - } - } +pub fn report( + info: &TraceInfo, + report_cycles: usize, + resolver: &DnsResolver, +) -> anyhow::Result<()> { + let mut next_id = 0; + let mut nodes: HashMap = HashMap::new(); + let mut edges: HashMap<(usize, usize), Edge> = HashMap::new(); super::wait_for_round(&info.data, report_cycles)?; let trace = info.data.read().clone(); - let mut graph: DiGraphMap = DiGraphMap::new(); - for (flow, _id) in trace.flows() { - for (fst, snd) in flow.entries.windows(2).map(|pair| (pair[0], pair[1])) { - match (fst, snd) { - (FlowEntry::Known(addr1), FlowEntry::Known(addr2)) => { - graph.add_edge(addr1, addr2, ()); - } - (FlowEntry::Known(addr1), FlowEntry::Unknown) => { - graph.add_edge(addr1, IpAddr::V4(Ipv4Addr::UNSPECIFIED), ()); - } - (FlowEntry::Unknown, FlowEntry::Known(addr2)) => { - graph.add_edge(IpAddr::V4(Ipv4Addr::UNSPECIFIED), addr2, ()); - } - _ => {} - } - } + for (flow, flow_id) in trace.flows() { + process_flow_entries( + &mut nodes, + &mut edges, + flow, + *flow_id, + &mut next_id, + resolver, + ); } - let dot = DotWrapper(Dot::with_config(&graph, &[Config::EdgeNoLabel])); - print!("{dot:?}"); + generate_dot_graph(&nodes, &edges); Ok(()) } + +fn create_or_get_node_id( + nodes: &mut HashMap, + entry: FlowEntry, + next_id: &mut usize, + resolver: &DnsResolver, +) -> usize { + match entry { + FlowEntry::Known(addr) => *nodes + .entry(addr) + .or_insert_with(|| create_node(next_id, addr, resolver)) + .id(), + FlowEntry::Unknown => *nodes + .entry(UNSPECIFIED_IP) + .or_insert_with(|| create_unknown_node(next_id)) + .id(), + } +} + +fn process_flow_entries( + nodes: &mut HashMap, + edges: &mut HashMap<(usize, usize), Edge>, + flow: &Flow, + flow_id: FlowId, + next_id: &mut usize, + resolver: &DnsResolver, +) { + for window in flow.entries.windows(2) { + if let [fst, snd] = *window { + let fst_id = create_or_get_node_id(nodes, fst, next_id, resolver); + let snd_id = create_or_get_node_id(nodes, snd, next_id, resolver); + edges + .entry((fst_id, snd_id)) + .or_insert_with(|| Edge::new(fst_id, snd_id)) + .value + .insert(flow_id); + } + } +} + +fn generate_dot_graph(nodes: &HashMap, edges: &HashMap<(usize, usize), Edge>) { + println!("digraph {{"); + println!(" node [shape=plaintext]"); + for node in nodes.values() { + println!(" {} [ label = {} ]", node.id, node.to_label_string()); + } + for edge in edges.values() { + println!( + " {} -> {} [ label = \"[{}]\" ]", + edge.from, + edge.to, + edge.to_label_string() + ); + } + println!("}}"); +} + +const UNSPECIFIED_IP: IpAddr = IpAddr::V4(Ipv4Addr::UNSPECIFIED); + +#[derive(Debug, Clone)] +struct Node { + id: usize, + addr: IpAddr, + names: Vec, + as_info: AsInfo, +} + +impl Node { + fn id(&self) -> &usize { + &self.id + } + + fn to_label_string(&self) -> String { + let as_label = if self.as_info.asn.is_empty() { + "n/a".to_string() + } else { + format!("AS{}", self.as_info.asn) + }; + + format!( + r#"<
{}{}
{}
>"#, + self.addr, + as_label, + self.names.join(", ") + ) + } +} + +#[derive(Debug, Clone)] +struct Edge { + from: usize, + to: usize, + value: HashSet, +} + +impl Edge { + fn new(from: usize, to: usize) -> Self { + Self { + from, + to, + value: HashSet::new(), + } + } + + fn to_label_string(&self) -> String { + self.value + .iter() + .map(|flow_id| flow_id.0.to_string()) + .collect::>() + .join(", ") + } +} + +// Utility functions to create nodes +fn create_node(next_id: &mut usize, addr: IpAddr, resolver: &DnsResolver) -> Node { + let id = *next_id; + *next_id += 1; + + let entry = resolver.reverse_lookup_with_asinfo(addr); + let (addr, names, as_info) = match entry { + DnsEntry::Resolved(Resolved::WithAsInfo(addr, names, as_info)) => (addr, names, as_info), + DnsEntry::Resolved(Resolved::Normal(addr, names)) => (addr, names, AsInfo::default()), + DnsEntry::NotFound(Unresolved::WithAsInfo(addr, as_info)) => { + (addr, vec![String::from("unknown")], as_info) + } + _ => (addr, vec![String::from("unknown")], AsInfo::default()), + }; + + Node { + id, + addr, + names, + as_info, + } +} + +fn create_unknown_node(next_id: &mut usize) -> Node { + let id = *next_id; + *next_id += 1; + + Node { + id, + addr: UNSPECIFIED_IP, + names: vec![String::from("unknown")], + as_info: AsInfo::default(), + } +}