netscli_core/
discover.rs

1use crate::arp::NetworkManager;
2use crate::oui::lookup_vendor;
3use crate::ping::PingScanner;
4use futures::stream::{self, StreamExt};
5use ipnet::Ipv4Net;
6use serde::Serialize;
7use std::collections::HashMap;
8use std::net::IpAddr;
9use std::sync::{
10    atomic::{AtomicUsize, Ordering},
11    Arc,
12};
13
14#[derive(Debug, Clone, Serialize)]
15pub struct Host {
16    pub ip: IpAddr,
17    pub hostname: Option<String>,
18    pub mac: Option<String>,
19    pub vendor: Option<String>,
20    pub rtt_ms: Option<u64>,
21}
22
23pub struct DiscoverEngine {
24    ping_scanner: PingScanner,
25    concurrency: usize,
26    ping_timeout_ms: u64,
27    dns_timeout_ms: u64,
28}
29
30#[derive(Debug, Clone, Copy)]
31pub enum DiscoverPhase {
32    Ping,
33    Resolve,
34}
35
36#[derive(Debug, Clone)]
37pub struct DiscoverProgress {
38    pub phase: DiscoverPhase,
39    pub completed: usize,
40    pub total: usize,
41    pub found: usize,
42    pub ip: IpAddr,
43}
44
45impl DiscoverEngine {
46    pub fn new(concurrency: usize) -> Self {
47        Self::new_with_timeouts(
48            concurrency,
49            crate::DEFAULT_PING_TIMEOUT_MS,
50            crate::DEFAULT_DNS_TIMEOUT_MS,
51        )
52    }
53
54    pub fn new_with_timeouts(
55        concurrency: usize,
56        ping_timeout_ms: u64,
57        dns_timeout_ms: u64,
58    ) -> Self {
59        let concurrency = concurrency.max(1);
60        Self {
61            ping_scanner: PingScanner::new(concurrency),
62            concurrency,
63            ping_timeout_ms: ping_timeout_ms.max(1),
64            dns_timeout_ms: dns_timeout_ms.max(1),
65        }
66    }
67
68    pub async fn scan_subnet(&self, subnet: Ipv4Net, resolve: bool) -> Vec<Host> {
69        self.scan_subnet_with_progress(subnet, resolve, None).await
70    }
71
72    pub async fn scan_subnet_with_progress(
73        &self,
74        subnet: Ipv4Net,
75        resolve: bool,
76        progress: Option<Arc<dyn Fn(DiscoverProgress) + Send + Sync>>,
77    ) -> Vec<Host> {
78        let ips: Vec<IpAddr> = subnet.hosts().map(IpAddr::V4).collect();
79
80        // 1) Ping first (fast), to avoid reverse-DNS work on dead hosts.
81        let total = ips.len();
82        let completed = Arc::new(AtomicUsize::new(0));
83        let found = Arc::new(AtomicUsize::new(0));
84        let ping_results = stream::iter(ips)
85            .map(|ip| {
86                let scanner = self.ping_scanner.clone();
87                let timeout_ms = self.ping_timeout_ms;
88                let completed = completed.clone();
89                let found = found.clone();
90                let progress = progress.clone();
91                async move {
92                    let res = scanner.ping(ip, timeout_ms).await;
93                    let done = completed.fetch_add(1, Ordering::SeqCst) + 1;
94                    // Capture `found` *after* both atomic updates settle.
95                    // Using a single load for both branches avoids the
96                    // previous inconsistency where an unreachable IP could
97                    // be reported with a `found` count that was updated by
98                    // a concurrent task between this task's two loads.
99                    if res.alive {
100                        found.fetch_add(1, Ordering::SeqCst);
101                    }
102                    let found_snapshot = found.load(Ordering::SeqCst);
103
104                    if let Some(cb) = &progress {
105                        if res.alive || done == total || done.is_multiple_of(10) {
106                            cb(DiscoverProgress {
107                                phase: DiscoverPhase::Ping,
108                                completed: done,
109                                total,
110                                found: found_snapshot,
111                                ip,
112                            });
113                        }
114                    }
115
116                    res
117                }
118            })
119            .buffer_unordered(self.concurrency)
120            .collect::<Vec<crate::ping::PingResult>>()
121            .await;
122
123        let alive: Vec<crate::ping::PingResult> =
124            ping_results.into_iter().filter(|r| r.alive).collect();
125
126        // 2) Load ARP/neighbor table once and reuse it.
127        let arp_map: HashMap<IpAddr, crate::arp::ArpEntry> = NetworkManager::get_arp_table()
128            .unwrap_or_default()
129            .into_iter()
130            .map(|e| (e.ip, e))
131            .collect();
132
133        // 3) Optionally reverse-DNS alive hosts using a single resolver.
134        let hostname_map: HashMap<IpAddr, Option<String>> = if resolve {
135            let ips = alive.iter().map(|r| r.ip).collect::<Vec<_>>();
136            let concurrency = self.concurrency.min(32);
137            let total = ips.len();
138            let completed = Arc::new(AtomicUsize::new(0));
139            let resolved = Arc::new(AtomicUsize::new(0));
140            stream::iter(ips)
141                .map(|ip| {
142                    let dns_timeout_ms = self.dns_timeout_ms;
143                    let completed = completed.clone();
144                    let resolved = resolved.clone();
145                    let progress = progress.clone();
146                    async move {
147                        let name =
148                            crate::dns::reverse_lookup_best_effort_timeout(ip, dns_timeout_ms)
149                                .await;
150                        let done = completed.fetch_add(1, Ordering::SeqCst) + 1;
151                        if name.is_some() {
152                            resolved.fetch_add(1, Ordering::SeqCst);
153                        }
154                        let resolved_count = resolved.load(Ordering::SeqCst);
155
156                        if let Some(cb) = &progress {
157                            if name.is_some() || done == total || done.is_multiple_of(5) {
158                                cb(DiscoverProgress {
159                                    phase: DiscoverPhase::Resolve,
160                                    completed: done,
161                                    total,
162                                    found: resolved_count,
163                                    ip,
164                                });
165                            }
166                        }
167
168                        (ip, name)
169                    }
170                })
171                .buffer_unordered(concurrency)
172                .collect::<Vec<(IpAddr, Option<String>)>>()
173                .await
174                .into_iter()
175                .collect()
176        } else {
177            HashMap::new()
178        };
179
180        alive
181            .into_iter()
182            .map(|r| {
183                let hostname = hostname_map.get(&r.ip).cloned().unwrap_or(None);
184                let mac_entry = arp_map.get(&r.ip);
185                let mac_str = mac_entry.map(|e| e.mac.to_string());
186                let vendor = mac_entry
187                    .and_then(|e| e.vendor.clone())
188                    .or_else(|| mac_str.as_deref().and_then(lookup_vendor));
189                Host {
190                    ip: r.ip,
191                    hostname,
192                    mac: mac_str,
193                    vendor,
194                    rtt_ms: r.rtt_ms,
195                }
196            })
197            .collect()
198    }
199}
netscli_core/discover.rs

netscli_core/
discover.rs
