Feat dot enriched

Cargo.toml

@@ -50,7 +50,6 @@ indexmap = { version = "2.1.0", default-features = false }
 maxminddb = "0.23.0"
 tracing-subscriber = { version = "0.3.18", default-features = false, features = [ "json", "env-filter" ] }
 tracing-chrome = "0.7.1"
-petgraph = "0.6.4"
 
 # Library dependencies (Linux)
 [target.'cfg(target_os = "linux")'.dependencies]

src/backend/flows.rs

@@ -21,6 +21,10 @@ use std::net::IpAddr;
 )]
 pub struct FlowId(pub u64);
 
+/// A time-to-live value.
+#[derive(Debug, Clone, Copy, Default, Ord, PartialOrd, Eq, PartialEq, Hash)]
+pub struct Ttl(pub u8);
+
 impl Display for FlowId {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         write!(f, "{}", self.0)
@@ -98,11 +102,12 @@ impl Flow {
     pub fn from_hops(hops: impl IntoIterator<Item = Option<IpAddr>>) -> Self {
         let entries = hops
             .into_iter()
-            .map(|addr| {
+            .enumerate()
+            .map(|(i, addr)| {
                 if let Some(addr) = addr {
-                    FlowEntry::Known(addr)
+                    FlowEntry::Known(Ttl(i as u8), addr)
                 } else {
-                    FlowEntry::Unknown
+                    FlowEntry::Unknown(Ttl(i as u8))
                 }
             })
             .collect();
@@ -127,10 +132,10 @@ impl Flow {
         let mut additions = 0;
         for (old, new) in self.entries.iter().zip(&flow.entries) {
             match (old, new) {
-                (FlowEntry::Known(fst), FlowEntry::Known(snd)) if fst != snd => {
+                (FlowEntry::Known(_, fst), FlowEntry::Known(_, snd)) if fst != snd => {
                     return CheckStatus::NoMatch;
                 }
-                (FlowEntry::Unknown, FlowEntry::Known(_)) => additions += 1,
+                (FlowEntry::Unknown(_), FlowEntry::Known(_, _)) => additions += 1,
                 _ => {}
             }
         }
@@ -149,7 +154,7 @@ impl Flow {
             .zip_longest(flow.entries.iter())
             .map(|eob| match eob {
                 EitherOrBoth::Both(left, right) => match (left, right) {
-                    (FlowEntry::Unknown, FlowEntry::Known(_)) => *right,
+                    (FlowEntry::Unknown(_), FlowEntry::Known(_, _)) => *right,
                     _ => *left,
                 },
                 EitherOrBoth::Left(left) => *left,
@@ -177,20 +182,20 @@ pub enum CheckStatus {
 }
 
 /// An entry in a `Flow`.
-#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
+#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash, Ord, PartialOrd)]
 pub enum FlowEntry {
     /// An unknown flow entry.
-    Unknown,
+    Unknown(Ttl),
     /// A known flow entry with an `IpAddr`.
-    Known(IpAddr),
+    Known(Ttl, IpAddr),
 }
 
 impl Display for FlowEntry {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
-            Self::Unknown => f.write_str("*"),
-            Self::Known(addr) => {
-                write!(f, "{addr}")
+            Self::Unknown(ttl) => write!(f, "{} *", ttl.0),
+            Self::Known(ttl, addr) => {
+                write!(f, "{} {addr}", ttl.0)
             }
         }
     }

src/main.rs

@@ -204,7 +204,7 @@ fn run_frontend(
         Mode::Json => report::run_report_json(&traces[0], args.report_cycles, &resolver)?,
         Mode::Pretty => report::run_report_table_pretty(&traces[0], args.report_cycles, &resolver)?,
         Mode::Markdown => report::run_report_table_md(&traces[0], args.report_cycles, &resolver)?,
-        Mode::Dot => report::run_report_dot(&traces[0], args.report_cycles)?,
+        Mode::Dot => report::run_report_dot(&traces[0], args.report_cycles, &resolver)?,
         Mode::Flows => report::run_report_flows(&traces[0], args.report_cycles)?,
         Mode::Silent => report::run_report_silent(&traces[0], args.report_cycles)?,
     }

src/report.rs

@@ -1,16 +1,11 @@
-use crate::backend::flows::FlowEntry;
 use crate::backend::trace::Trace;
 use crate::TraceInfo;
 use anyhow::anyhow;
 use comfy_table::presets::{ASCII_MARKDOWN, UTF8_FULL};
 use comfy_table::{ContentArrangement, Table};
 use itertools::Itertools;
 use parking_lot::RwLock;
-use petgraph::dot::{Config, Dot};
-use petgraph::graphmap::DiGraphMap;
 use serde::{Serialize, Serializer};
-use std::fmt::{Debug, Formatter};
-use std::net::{IpAddr, Ipv4Addr};
 use std::sync::Arc;
 use std::thread::sleep;
 use std::time::Duration;
@@ -277,36 +272,192 @@ pub fn run_report_silent(info: &TraceInfo, report_cycles: usize) -> anyhow::Resu
     Ok(())
 }
 
-/// Run a trace and generate a dot file.
-pub fn run_report_dot(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 use dot::run_report_dot;
+
+pub mod dot {
+    use crate::backend::flows::{Flow, FlowEntry, FlowId, Ttl};
+    use crate::report::wait_for_round;
+    use crate::TraceInfo;
+    use std::collections::{HashMap, HashSet};
+    use std::net::{IpAddr, Ipv4Addr};
+    use trippy::dns::{AsInfo, DnsEntry, DnsResolver, Resolved, Resolver, Unresolved};
+
+    /// Run a trace and generate a dot file.
+    pub fn run_report_dot(
+        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();
+        wait_for_round(&info.data, report_cycles)?;
+        let trace = info.data.read().clone();
+        for (flow, flow_id) in trace.flows() {
+            process_flow_entries(
+                &mut nodes,
+                &mut edges,
+                flow,
+                *flow_id,
+                &mut next_id,
+                resolver,
+            );
         }
+        generate_dot_graph(&nodes, &edges);
+        Ok(())
     }
-    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, ());
-                }
-                _ => {}
+
+    fn create_or_get_node_id(
+        nodes: &mut HashMap<IpAddr, Node>,
+        entry: FlowEntry,
+        next_id: &mut usize,
+        resolver: &DnsResolver,
+    ) -> usize {
+        match entry {
+            FlowEntry::Known(ttl, addr) => *nodes
+                .entry(addr)
+                .or_insert_with(|| create_node(next_id, ttl, addr, resolver))
+                .id(),
+            FlowEntry::Unknown(ttl) => *nodes
+                .entry(UNSPECIFIED_IP)
+                .or_insert_with(|| create_unknown_node(next_id, ttl))
+                .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);
             }
         }
     }
-    let dot = DotWrapper(Dot::with_config(&graph, &[Config::EdgeNoLabel]));
-    print!("{dot:?}");
-    Ok(())
+
+    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,
+        _ttl: Ttl,
+        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, ttl: Ttl, 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,
+            _ttl: ttl,
+            names,
+            as_info,
+        }
+    }
+
+    fn create_unknown_node(next_id: &mut usize, ttl: Ttl) -> Node {
+        let id = *next_id;
+        *next_id += 1;
+
+        Node {
+            id,
+            addr: UNSPECIFIED_IP,
+            _ttl: ttl,
+            names: vec![String::from("unknown")],
+            as_info: AsInfo::default(),
+        }
+    }
 }
 
 /// Run a trace and report all flows observed.