feat: improve dot reporting mode (#797)

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,20 +1,19 @@
-use crate::backend::flows::FlowEntry;
+use crate::backend::flows::{FlowEntry, FlowId, Ttl};
 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::collections::{HashMap, HashSet};
+use std::fmt::{Display, Formatter};
 use std::net::{IpAddr, Ipv4Addr};
 use std::sync::Arc;
 use std::thread::sleep;
 use std::time::Duration;
-use trippy::dns::{DnsResolver, Resolver};
+use trippy::dns::{AsInfo, DnsEntry, DnsResolver, Resolved, Resolver, Unresolved};
 
 /// Generate a CSV report of trace data.
 pub fn run_report_csv(
@@ -278,34 +277,191 @@ pub fn run_report_silent(info: &TraceInfo, report_cycles: usize) -> anyhow::Resu
 }
 
 /// 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<'_> {
+pub fn run_report_dot(
+    info: &TraceInfo,
+    report_cycles: usize,
+    resolver: &DnsResolver,
+) -> anyhow::Result<()> {
+    #[derive(Debug)]
+    struct Weights(Vec<FlowId>);
+    impl Display for Weights {
         fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-            self.0.fmt(f)
+            write!(f, "{}", self.0.iter().format(", "))
         }
     }
+
+    #[derive(Debug, Clone)]
+    struct Node {
+        id: usize,
+        addr: IpAddr,
+        ttl: Ttl,
+        names: Vec<String>,
+        as_info: AsInfo,
+    }
+
+    #[derive(Debug, Clone)]
+    struct Edge {
+        from: usize,
+        to: usize,
+        value: HashSet<FlowId>,
+    }
+
+    // impl Display for Node<'_> {
+    //     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+    //         self.entry.fmt(f)
+    //     }
+    // }
+
+    let mut next_id = 0;
+    let mut nodes: HashMap<IpAddr, Node> = HashMap::new();
+    let mut edges: HashMap<(usize, usize), Edge> = HashMap::new();
+
+    // Nodes should be unique for (addr, ttl) so that Unknown nodes at the different ttl's are not conflated (ok to conflate at same ttl?)
+
     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 (flow, 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, ());
+            let fst_id = match fst {
+                FlowEntry::Known(ttl, addr) => {
+                    nodes
+                        .entry(addr)
+                        .or_insert_with(|| {
+                            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,
+                                names,
+                                as_info,
+                            }
+                        })
+                        .id
                 }
-                (FlowEntry::Known(addr1), FlowEntry::Unknown) => {
-                    graph.add_edge(addr1, IpAddr::V4(Ipv4Addr::UNSPECIFIED), ());
+                FlowEntry::Unknown(ttl) => {
+                    nodes
+                        .entry(IpAddr::V4(Ipv4Addr::UNSPECIFIED))
+                        .or_insert_with(|| {
+                            let id = next_id;
+                            next_id += 1;
+                            Node {
+                                id,
+                                addr: IpAddr::V4(Ipv4Addr::UNSPECIFIED),
+                                ttl,
+                                names: vec![String::from("unknown")],
+                                as_info: AsInfo::default(),
+                            }
+                        })
+                        .id
                 }
-                (FlowEntry::Unknown, FlowEntry::Known(addr2)) => {
-                    graph.add_edge(IpAddr::V4(Ipv4Addr::UNSPECIFIED), addr2, ());
+            };
+            let snd_id = match snd {
+                FlowEntry::Known(ttl, addr) => {
+                    nodes
+                        .entry(addr)
+                        .or_insert_with(|| {
+                            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,
+                                names,
+                                as_info,
+                            }
+                        })
+                        .id
                 }
-                _ => {}
-            }
+                FlowEntry::Unknown(ttl) => {
+                    nodes
+                        .entry(IpAddr::V4(Ipv4Addr::UNSPECIFIED))
+                        .or_insert_with(|| {
+                            let id = next_id;
+                            next_id += 1;
+                            Node {
+                                id,
+                                addr: IpAddr::V4(Ipv4Addr::UNSPECIFIED),
+                                ttl,
+                                names: vec![String::from("unknown")],
+                                as_info: AsInfo::default(),
+                            }
+                        })
+                        .id
+                }
+            };
+
+            // Edges
+
+            edges
+                .entry((fst_id, snd_id))
+                .or_insert_with(|| Edge {
+                    from: fst_id,
+                    to: snd_id,
+                    value: HashSet::new(),
+                })
+                .value
+                .insert(*flow_id);
         }
     }
-    let dot = DotWrapper(Dot::with_config(&graph, &[Config::EdgeNoLabel]));
-    print!("{dot:?}");
+
+    println!("digraph {{");
+    println!("    node [shape=plaintext]");
+    for node in nodes.values().sorted_by_key(|node| node.id) {
+        let as_label = if node.as_info.asn.is_empty() {
+            format!("n/a")
+        } else {
+            format!("AS{}", node.as_info.asn)
+        };
+        let label = format!(
+            r#"<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="4"><tr><td>{}</td><td>{}</td></tr><tr><td COLSPAN="2">{}</td></tr></TABLE>>"#,
+            node.addr,
+            as_label,
+            node.names.join(", ")
+        );
+
+        println!("    {} [ label = {} ]", node.id, label);
+    }
+    for (_, edge) in edges {
+        println!(
+            "    {} -> {} [ label = \"[{}]\" ]",
+            edge.from,
+            edge.to,
+            edge.value
+                .iter()
+                .map(|flow_id| flow_id.0)
+                .sorted()
+                .join(", ")
+        );
+    }
+    println!("}}");
+
     Ok(())
 }