1use 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#[derive(Debug, Clone)]
18pub struct DiscoveredService {
19 pub service_info: ServiceInfo,
20}
21
22impl DiscoveredService {
23 pub fn ip(&self) -> &str {
25 &self.service_info.ip
26 }
27
28 pub fn port(&self) -> u16 {
30 self.service_info.port
31 }
32
33 pub fn base_url(&self) -> String {
35 self.service_info.base_url()
36 }
37
38 pub fn manifest(&self) -> &serde_json::Value {
40 &self.service_info.manifest
41 }
42}
43
44pub struct DiscoveryScanner {
46 config: ClientConfig,
47}
48
49impl DiscoveryScanner {
50 pub fn new(config: ClientConfig) -> Self {
52 Self { config }
53 }
54
55 pub fn default_config() -> Self {
57 Self::new(ClientConfig::default())
58 }
59
60 pub async fn scan(
62 &self,
63 _fetch_manifest: Option<bool>,
64 ) -> std::result::Result<Vec<DiscoveredService>, ScannerError> {
65 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 async fn broadcast_and_collect(
79 &self,
80 ) -> std::result::Result<Vec<DiscoveredService>, ScannerError> {
81 let socket = UdpSocket::bind("0.0.0.0:0")
83 .await
84 .map_err(|e| ScannerError::BindError(e.to_string()))?;
85
86 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 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 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 if !services.is_empty() {
152 break;
153 }
154 }
155 }
156 }
157
158 Ok(services)
159 }
160}
161
162#[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}