async_traceroute/traceroute/
terminal.rs

1use std::collections::HashMap;
2use std::net::Ipv4Addr;
3
4use futures::pin_mut;
5use futures_util::StreamExt;
6
7use crate::traceroute::probe::ProbeResult;
8use crate::traceroute::Traceroute;
9
10type Ttl = u8;
11
12pub struct TracerouteTerminal {
13    traceroute: Traceroute,
14    current_ttl: u8,
15    nqueries: u16,
16    max_ttl: u8,
17    hop_by_ttl: HashMap<Ttl, PrintableHop>,
18}
19
20impl TracerouteTerminal {
21    pub fn new(traceroute: Traceroute) -> Self {
22        let nqueries = traceroute.get_nqueries();
23        let max_ttl = traceroute.get_max_ttl();
24        Self {
25            traceroute,
26            current_ttl: 1,
27            nqueries,
28            max_ttl,
29            hop_by_ttl: HashMap::with_capacity(max_ttl as usize)
30        }
31    }
32    
33    pub async fn print_trace(mut self) {
34        let traceroute_stream= self.traceroute.trace();
35        pin_mut!(traceroute_stream);
36
37        let mut skip = true;
38        self.hop_by_ttl.insert(1, PrintableHop::new(1));
39        Self::print_current_hop_index(self.current_ttl);
40
41        'outer: while let Some(probe_result) = traceroute_stream.next().await {
42            let ttl = match probe_result {
43                Ok(probe_result) => {
44                    let probe_result_ttl = probe_result.ttl();
45
46                    let hop = self.hop_by_ttl
47                        .entry(probe_result_ttl)
48                        .or_insert(PrintableHop::new(probe_result_ttl));
49
50                    hop.add_printable_probe_result(
51                        PrintableProbeResult::ProbeResult {
52                            probe_result,
53                            printed: false
54                        }
55                    );
56
57                    probe_result_ttl
58                },
59                Err(probe_error) => {
60                    let probe_error_ttl = probe_error.get_ttl();
61
62                    let hop = self.hop_by_ttl
63                        .entry(probe_error_ttl)
64                        .or_insert(PrintableHop::new(probe_error_ttl));
65
66                    hop.add_printable_probe_result(
67                        PrintableProbeResult::Timeout {
68                            printed: false
69                        }
70                    );
71
72                    probe_error_ttl
73                },
74            };
75
76            if ttl > self.current_ttl && skip {
77                skip = false;
78                
79                for current_ttl in self.current_ttl+1..=ttl {
80                    for _ in 0..self.nqueries {
81                        print!("* ");
82                    }
83
84                    println!();
85                    Self::print_current_hop_index(current_ttl);
86                }
87                
88                self.current_ttl = ttl;
89            }
90
91            while let Some(hop) = self.hop_by_ttl.get_mut(&self.current_ttl) {
92                hop.print();
93
94                if hop.tot_completed_queries() == self.nqueries as usize {
95                    self.current_ttl += 1;
96
97                    if self.current_ttl > self.max_ttl {
98                        break 'outer;
99                    } else {
100                        println!();
101                        Self::print_current_hop_index(self.current_ttl);
102                    }
103                } else {
104                    break;
105                }
106            }
107        }
108
109        println!()
110    }
111    fn print_current_hop_index(ttl: u8) {
112        if ttl < 10 {
113            print!(" {}  ", ttl);
114        } else {
115            print!("{}  ", ttl);
116        }
117    }
118}
119
120enum PrintableProbeResult {
121    Timeout { printed: bool },
122    ProbeResult { probe_result: ProbeResult, printed: bool },
123}
124
125
126impl PrintableProbeResult {
127    fn mark_as_printed(&mut self) {
128        let printed = match self {
129            PrintableProbeResult::Timeout { ref mut printed } => printed,
130            PrintableProbeResult::ProbeResult {ref mut printed, .. } => printed,
131        };
132        
133        *printed = true;
134    }
135
136    fn is_printed(&self) -> bool {
137        match self {
138            PrintableProbeResult::Timeout { printed } => *printed,
139            PrintableProbeResult::ProbeResult { printed, .. } => *printed
140        }
141    }
142}
143
144struct PrintableHop {
145    ttl: u8,
146    probe_results: Vec<PrintableProbeResult>,
147    printed_hostnames: Vec<String>,
148    printed_ip_addr: Vec<Ipv4Addr>,
149}
150
151impl PrintableHop {
152    fn new(ttl: u8) -> PrintableHop {
153        Self {
154            ttl,
155            probe_results: Vec::with_capacity(10),
156            printed_hostnames: Vec::with_capacity(10),
157            printed_ip_addr: Vec::with_capacity(10),
158        }
159    }
160
161    fn add_printable_probe_result(&mut self, printable_probe_result: PrintableProbeResult) {
162        self.probe_results.push(printable_probe_result)
163    }
164
165    fn tot_completed_queries(&self) -> usize {
166        self.probe_results.len()
167    }
168
169    fn print(&mut self) {
170        let printed_hostnames = &mut self.printed_hostnames;
171        let printed_ip_addr = &mut self.printed_ip_addr;
172
173        for printable_probe_result in self.probe_results.iter_mut() {
174            if printable_probe_result.is_printed() {
175                continue;
176            }
177
178            printable_probe_result.mark_as_printed();
179
180            match printable_probe_result {
181                PrintableProbeResult::Timeout { .. } => {
182                    print!("* ");
183                },
184                PrintableProbeResult::ProbeResult { probe_result, .. } => {
185                    let rtt = probe_result.rtt();
186                    let hostname = probe_result.get_hostname();
187                    if hostname.is_none() || printed_hostnames.contains(&hostname.as_ref().unwrap()) {
188                        if printed_ip_addr.contains(&probe_result.from_address()) {
189                            let mut rtt_micros = rtt.as_micros().to_string();
190                            rtt_micros.insert(2, '.');
191                            print!("{:2} ms  ", rtt_micros);
192                        } else {
193                            print!("{probe_result}  ");
194                            printed_ip_addr.push(probe_result.from_address());
195                        }
196                    } else {
197                        print!("{probe_result}  ");
198                        printed_hostnames.push(hostname.unwrap());
199                        printed_ip_addr.push(probe_result.from_address());
200                    }
201                }
202            }
203        }
204    }
205}
async_traceroute/traceroute/terminal.rs

async_traceroute/traceroute/
terminal.rs
