feat(report): enhanced dot `report` (#797) - WIPfeat-dot-enriched-2

4 files changed, 170 insertions, 48 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 96ecaad..70f2a7a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -597,12 +597,6 @@ dependencies = [
 ]
 
 [[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"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1188,16 +1182,6 @@ 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"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1981,7 +1965,6 @@ dependencies = [
  "nix",
  "parking_lot",
  "paste",
- "petgraph",
  "pretty_assertions",
  "rand",
  "ratatui",
diff --git a/Cargo.toml b/Cargo.toml
index a9ebb18..699c572 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 5c06562..283eb6e 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 12e6eb7..b3d77a2 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<IpAddr, ()>>);
-    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<IpAddr, Node> = 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<IpAddr, ()> = 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<IpAddr, Node>,
+    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<IpAddr, Node>,
+    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<IpAddr, Node>, 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<String>,
+    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#"<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="4"><tr><td>{}</td><td>{}</td></tr><tr><td COLSPAN="2">{}</td></tr></TABLE>>"#,
+            self.addr,
+            as_label,
+            self.names.join(", ")
+        )
+    }
+}
+
+#[derive(Debug, Clone)]
+struct Edge {
+    from: usize,
+    to: usize,
+    value: HashSet<FlowId>,
+}
+
+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::<Vec<_>>()
+            .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(),
+    }
+}