Skip to main content

aiecho/
scanner.rs

1//! Discovery Scanner (Client) Implementation
2//!
3//! This module implements the client-side scanner that:
4//! - Sends UDP broadcast discovery requests
5//! - Collects responses from all services
6
7use std::net::SocketAddr;
8use std::time::Duration;
9use tokio::net::UdpSocket;
10use tokio::time::timeout;
11use tracing::{debug, error, info};
12
13use crate::config::ClientConfig;
14use crate::protocol::{build_discover_req, parse_message, ServiceInfo, DISCOVER_RES};
15
16/// Complete discovered service with manifest
17#[derive(Debug, Clone)]
18pub struct DiscoveredService {
19    pub service_info: ServiceInfo,
20}
21
22impl DiscoveredService {
23    /// Get service IP address
24    pub fn ip(&self) -> &str {
25        &self.service_info.ip
26    }
27
28    /// Get service port
29    pub fn port(&self) -> u16 {
30        self.service_info.port
31    }
32
33    /// Get service base URL
34    pub fn base_url(&self) -> String {
35        self.service_info.base_url()
36    }
37
38    /// Get manifest data
39    pub fn manifest(&self) -> &serde_json::Value {
40        &self.service_info.manifest
41    }
42}
43
44/// Discovery scanner that broadcasts queries and collects responses
45pub struct DiscoveryScanner {
46    config: ClientConfig,
47}
48
49impl DiscoveryScanner {
50    /// Create a new discovery scanner
51    pub fn new(config: ClientConfig) -> Self {
52        Self { config }
53    }
54
55    /// Create with default config
56    pub fn default_config() -> Self {
57        Self::new(ClientConfig::default())
58    }
59
60    /// Scan for services on the network
61    pub async fn scan(
62        &self,
63        _fetch_manifest: Option<bool>,
64    ) -> std::result::Result<Vec<DiscoveredService>, ScannerError> {
65        // Broadcast and collect responses
66        let services = self.broadcast_and_collect().await?;
67
68        if services.is_empty() {
69            return Ok(vec![]);
70        }
71
72        info!("Discovered {} service(s)", services.len());
73
74        Ok(services)
75    }
76
77    /// Send broadcast and collect all responses
78    async fn broadcast_and_collect(
79        &self,
80    ) -> std::result::Result<Vec<DiscoveredService>, ScannerError> {
81        // Create UDP socket
82        let socket = UdpSocket::bind("0.0.0.0:0")
83            .await
84            .map_err(|e| ScannerError::BindError(e.to_string()))?;
85
86        // Enable broadcast
87        let socket = socket
88            .into_std()
89            .map_err(|e| ScannerError::BindError(e.to_string()))?;
90        socket
91            .set_broadcast(true)
92            .map_err(|e| ScannerError::BindError(e.to_string()))?;
93        let socket =
94            UdpSocket::from_std(socket).map_err(|e| ScannerError::BindError(e.to_string()))?;
95
96        // Send discovery request
97        let request_msg = build_discover_req(None);
98        let broadcast_addr: SocketAddr = format!("255.255.255.255:{}", self.config.udp_port)
99            .parse()
100            .map_err(|e: std::net::AddrParseError| ScannerError::AddressError(e.to_string()))?;
101
102        socket
103            .send_to(&request_msg, broadcast_addr)
104            .await
105            .map_err(|e: std::io::Error| ScannerError::SendError(e.to_string()))?;
106
107        debug!("Sent discovery request to broadcast address");
108
109        // Collect responses with timeout
110        let mut services: Vec<DiscoveredService> = Vec::new();
111        let timeout_duration = Duration::from_secs_f64(self.config.timeout);
112        let start = std::time::Instant::now();
113
114        let mut buf = [0u8; 4096];
115
116        while start.elapsed() < timeout_duration {
117            let remaining = timeout_duration - start.elapsed();
118
119            match timeout(remaining, socket.recv_from(&mut buf)).await {
120                Ok(Ok((len, addr))) => {
121                    let data = &buf[..len];
122
123                    match parse_message(data) {
124                        Ok((cmd, payload)) => {
125                            if cmd == DISCOVER_RES {
126                                let service_info = ServiceInfo::from_payload(
127                                    &payload,
128                                    addr.ip().to_string().as_str(),
129                                );
130
131                                debug!(
132                                    "Discovered: {} @ {}:{}",
133                                    service_info.port,
134                                    addr.ip(),
135                                    service_info.port
136                                );
137
138                                services.push(DiscoveredService { service_info });
139                            }
140                        }
141                        Err(e) => {
142                            debug!("Failed to parse response from {}: {}", addr, e);
143                        }
144                    }
145                }
146                Ok(Err(e)) => {
147                    error!("Error receiving UDP packet: {}", e);
148                }
149                Err(_) => {
150                    // Timeout, check if we should continue
151                    if !services.is_empty() {
152                        break;
153                    }
154                }
155            }
156        }
157
158        Ok(services)
159    }
160}
161
162/// Scanner errors
163#[derive(Debug, thiserror::Error)]
164pub enum ScannerError {
165    #[error("Failed to bind to UDP port: {0}")]
166    BindError(String),
167
168    #[error("Failed to send UDP packet: {0}")]
169    SendError(String),
170
171    #[error("Invalid address: {0}")]
172    AddressError(String),
173
174    #[error("HTTP error: {0}")]
175    HttpError(String),
176}