1use std::net::{SocketAddr, ToSocketAddrs};
2use std::sync::Arc;
3
4use tokio::sync::{Mutex, RwLock};
5use tracing::{debug, info, warn};
6
7use aria2_protocol::bittorrent::torrent::parser::TorrentMeta;
8use aria2_protocol::bittorrent::tracker::udp_tracker_protocol::{AnnounceResponse, UdpEvent};
9
10use crate::engine::udp_tracker_client::SharedUdpClient;
11
12const DEFAULT_UDP_PORT: u16 = 6881;
13const MAX_TRACKER_ANNOUNCE_TIER: usize = 2;
14
15#[derive(Debug, Clone)]
16pub struct UdpTrackerEndpoint {
17 pub url: String,
18 pub addr: SocketAddr,
19 pub tier: u32,
20}
21
22pub struct UdpTrackerManager {
23 client: SharedUdpClient,
24 endpoints: Vec<UdpTrackerEndpoint>,
25 current_tier: usize,
26 last_announce_interval: u32,
27 last_announce_time: Option<tokio::time::Instant>,
28}
29
30impl UdpTrackerManager {
31 pub async fn new(client: SharedUdpClient) -> Self {
32 Self {
33 client,
34 endpoints: Vec::new(),
35 current_tier: 0,
36 last_announce_interval: 300,
37 last_announce_time: None,
38 }
39 }
40
41 pub fn parse_tracker_urls(&mut self, urls: &[String]) -> usize {
42 let mut added = 0usize;
43 #[allow(clippy::explicit_counter_loop)]
44 for (tier, urls_tier) in urls.iter().enumerate() {
45 for url in urls_tier.split(',') {
46 let trimmed = url.trim();
47 if let Some(endpoint) = Self::parse_udp_url(trimmed, tier as u32) {
48 debug!(
49 "Adding UDP tracker endpoint: {} -> {}",
50 trimmed, endpoint.addr
51 );
52 self.endpoints.push(endpoint);
53 added += 1;
54 }
55 }
56 }
57 if !self.endpoints.is_empty() {
58 info!(
59 "Parsed {} UDP tracker endpoints from {} URL groups",
60 added,
61 urls.len()
62 );
63 }
64 added
65 }
66
67 fn parse_udp_url(url: &str, tier: u32) -> Option<UdpTrackerEndpoint> {
68 let url_lower = url.to_lowercase();
69
70 if !url_lower.starts_with("udp://") {
71 return None;
72 }
73
74 let host_port = &url[6..];
75 let host_part = if let Some(query_idx) = host_port.find('?') {
76 &host_port[..query_idx]
77 } else if let Some(path_idx) = host_port.find('/') {
78 &host_port[..path_idx]
79 } else {
80 host_port
81 };
82
83 if let Some(addr) = Self::resolve_host(host_part) {
84 Some(UdpTrackerEndpoint {
85 url: url.to_string(),
86 addr,
87 tier,
88 })
89 } else {
90 warn!("Failed to resolve UDP tracker URL: {}", url);
91 None
92 }
93 }
94
95 fn resolve_host(host_port: &str) -> Option<SocketAddr> {
96 let (host, port_str) = if host_port.contains(':') {
97 let parts: Vec<&str> = host_port.rsplitn(2, ':').collect();
98 if parts.len() == 2 {
99 (parts[1], parts[0])
100 } else {
101 (host_port, "6881")
102 }
103 } else {
104 (host_port, "6881")
105 };
106
107 let port: u16 = port_str.parse::<u16>().unwrap_or(DEFAULT_UDP_PORT);
108 let addr_str = format!("{}:{}", host, port);
109
110 match addr_str.to_socket_addrs() {
111 Ok(mut addrs) => addrs.next(),
112 Err(e) => {
113 warn!("DNS resolution failed for {}: {}", addr_str, e);
114 None
115 }
116 }
117 }
118
119 #[allow(clippy::too_many_arguments)]
120 pub async fn announce(
121 &mut self,
122 info_hash: &[u8; 20],
123 peer_id: &[u8; 20],
124 downloaded: i64,
125 left: i64,
126 uploaded: i64,
127 event: UdpEvent,
128 num_want: i32,
129 ) -> Vec<AnnounceResponse> {
130 if self.endpoints.is_empty() {
131 debug!("No UDP tracker endpoints configured");
132 return Vec::new();
133 }
134
135 let port = self.get_bind_port().await;
136 let start_tier = self
137 .current_tier
138 .min(self.endpoints.len().saturating_sub(1));
139
140 for tier_offset in 0..MAX_TRACKER_ANNOUNCE_TIER {
141 let tier_idx = (start_tier + tier_offset) % self.endpoints.len();
142 let ep = &self.endpoints[tier_idx];
143
144 {
145 let mut c = self.client.lock().await;
146 c.add_announce(
147 &ep.addr, info_hash, peer_id, downloaded, left, uploaded, event, num_want, port,
148 )
149 .await;
150 }
151
152 let mut c = self.client.lock().await;
153 while c.process_one().await {
154 tokio::task::yield_now().await;
155 }
156 }
157
158 let results = self.process_responses().await;
159 if results.is_empty() {
160 self.advance_tier();
161 }
162
163 results
164 }
165
166 pub async fn process_responses(&mut self) -> Vec<AnnounceResponse> {
167 let mut c = self.client.lock().await;
168 let refs = c.completed_requests();
169 let results: Vec<AnnounceResponse> = refs.iter().map(|r| (*r).clone()).collect();
170
171 if !results.is_empty() {
172 self.last_announce_interval = results.first().map(|r| r.interval).unwrap_or(300);
173 self.last_announce_time = Some(tokio::time::Instant::now());
174 info!(
175 "UDP tracker announce completed: {} responses",
176 results.len()
177 );
178 }
179
180 c.handle_timeouts().await;
181 results
182 }
183
184 fn advance_tier(&mut self) {
185 if !self.endpoints.is_empty() {
186 self.current_tier = (self.current_tier + 1) % self.endpoints.len();
187 debug!("Advanced to tracker tier {}", self.current_tier);
188 }
189 }
190
191 pub fn endpoint_count(&self) -> usize {
192 self.endpoints.len()
193 }
194
195 pub fn get_announce_interval(&self) -> u32 {
196 self.last_announce_interval
197 }
198
199 pub fn should_reannounce(&self) -> bool {
200 if let Some(t) = self.last_announce_time {
201 t.elapsed().as_secs() >= self.last_announce_interval as u64
202 } else {
203 true
204 }
205 }
206
207 async fn get_bind_port(&self) -> u16 {
208 let c = self.client.lock().await;
209 c.socket()
210 .local_addr()
211 .map(|a| a.port())
212 .unwrap_or(DEFAULT_UDP_PORT)
213 }
214
215 pub fn collect_all_peers(responses: &[AnnounceResponse]) -> Vec<(String, u16)> {
216 let mut all_peers = Vec::new();
217 let mut seen = std::collections::HashSet::new();
218
219 for resp in responses {
220 for peer in &resp.peers {
221 if seen.insert(peer.clone()) {
222 all_peers.push(peer.clone());
223 }
224 }
225 }
226
227 all_peers
228 }
229}
230
231impl UdpTrackerManager {
232 pub async fn from_torrent_meta(
233 meta: &Arc<RwLock<TorrentMeta>>,
234 client: SharedUdpClient,
235 ) -> Option<Self> {
236 let mut mgr = Self::new(client).await;
237 let m = meta.read().await;
238 let urls: Vec<String> = m.announce_list.iter().flatten().cloned().collect();
239 drop(m);
240
241 if mgr.parse_tracker_urls(&urls) > 0 {
242 Some(mgr)
243 } else {
244 None
245 }
246 }
247}
248
249pub type SharedUdpTrackerManager = Arc<Mutex<UdpTrackerManager>>;
250
251impl UdpTrackerManager {
252 pub fn create_shared(client: SharedUdpClient) -> SharedUdpTrackerManager {
253 Arc::new(Mutex::new(Self::new_blocking(client)))
254 }
255
256 fn new_blocking(client: SharedUdpClient) -> Self {
257 Self {
258 client,
259 endpoints: Vec::new(),
260 current_tier: 0,
261 last_announce_interval: 300,
262 last_announce_time: None,
263 }
264 }
265}
266
267#[cfg(test)]
268mod tests {
269 use super::*;
270 use crate::engine::udp_tracker_client::UdpTrackerClient;
271
272 #[tokio::test]
273 async fn test_manager_creation() {
274 let shared = UdpTrackerClient::create_shared(0).await.unwrap();
275 let mgr = UdpTrackerManager::new(shared).await;
276 assert_eq!(mgr.endpoint_count(), 0);
277 assert!(mgr.should_reannounce());
278 }
279
280 #[test]
281 fn test_parse_udp_url_valid() {
282 let ep = UdpTrackerManager::parse_udp_url("udp://127.0.0.1:6969/announce", 0);
283 assert!(ep.is_some());
284 let ep = ep.unwrap();
285 assert_eq!(ep.tier, 0);
286 assert_eq!(ep.url, "udp://127.0.0.1:6969/announce");
287 }
288
289 #[test]
290 fn test_parse_udp_url_non_udp() {
291 let ep = UdpTrackerManager::parse_udp_url("http://tracker.example.com:6969/announce", 0);
292 assert!(ep.is_none(), "HTTP URLs should not be parsed as UDP");
293 }
294
295 #[test]
296 fn test_parse_udp_url_default_port() {
297 let ep = UdpTrackerManager::parse_udp_url("udp://127.0.0.1/announce", 0);
298 assert!(ep.is_some());
299 assert_eq!(ep.unwrap().addr.port(), 6881);
300 }
301
302 #[tokio::test]
303 async fn test_parse_tracker_urls_mixed() {
304 let shared = UdpTrackerClient::create_shared(0).await.unwrap();
305 let mut mgr = UdpTrackerManager::new(shared).await;
306
307 let urls = vec![
308 "http://tracker1.example.com:6969/announce".to_string(),
309 "udp://127.0.0.1:6969/announce".to_string(),
310 "udp://127.0.0.1:80/announce".to_string(),
311 ];
312
313 let count = mgr.parse_tracker_urls(&urls);
314 assert_eq!(count, 2, "Should parse only UDP URLs");
315 assert_eq!(mgr.endpoint_count(), 2);
316 }
317
318 #[tokio::test]
319 async fn test_multi_tier_parsing() {
320 let shared = UdpTrackerClient::create_shared(0).await.unwrap();
321 let mut mgr = UdpTrackerManager::new(shared).await;
322
323 let urls = vec![
324 "udp://127.0.0.1:6969/announce,udp://127.0.0.1:6970/announce".to_string(),
325 "udp://127.0.0.1:6980/announce".to_string(),
326 ];
327
328 let count = mgr.parse_tracker_urls(&urls);
329 assert_eq!(count, 3);
330 assert_eq!(mgr.endpoints[0].tier, 0);
331 assert_eq!(mgr.endpoints[1].tier, 0);
332 assert_eq!(mgr.endpoints[2].tier, 1);
333 }
334
335 #[tokio::test]
336 async fn test_should_reannounce_logic() {
337 let shared = UdpTrackerClient::create_shared(0).await.unwrap();
338 let mut mgr = UdpTrackerManager::new(shared).await;
339
340 assert!(mgr.should_reannounce(), "Should reannounce initially");
341 mgr.last_announce_time = Some(tokio::time::Instant::now());
342 assert!(
343 !mgr.should_reannounce(),
344 "Should not reannounce immediately after"
345 );
346
347 mgr.last_announce_interval = 0;
348 assert!(mgr.should_reannounce(), "Should reannounce with interval=0");
349 }
350
351 #[tokio::test]
352 async fn test_advance_tier_wraps() {
353 let shared = UdpTrackerClient::create_shared(0).await.unwrap();
354 let mut mgr = UdpTrackerManager::new(shared).await;
355
356 let urls = vec![
357 "udp://127.0.0.1:6969/announce".to_string(),
358 "udp://127.0.0.1:6970/announce".to_string(),
359 "udp://127.0.0.1:6980/announce".to_string(),
360 ];
361 mgr.parse_tracker_urls(&urls);
362
363 assert_eq!(mgr.current_tier, 0);
364 mgr.advance_tier();
365 assert_eq!(mgr.current_tier, 1);
366 mgr.advance_tier();
367 assert_eq!(mgr.current_tier, 2);
368 mgr.advance_tier();
369 assert_eq!(mgr.current_tier, 0, "Tier should wrap around");
370 }
371
372 #[tokio::test]
373 async fn test_collect_all_peers_dedup() {
374 let shared = UdpTrackerClient::create_shared(0).await.unwrap();
375 let _mgr = UdpTrackerManager::new(shared).await;
376
377 let resp1 = AnnounceResponse {
378 transaction_id: 1,
379 interval: 300,
380 leechers: 10,
381 seeders: 5,
382 peers: vec![
383 ("192.168.1.1".to_string(), 6881),
384 ("192.168.1.2".to_string(), 6882),
385 ],
386 };
387
388 let resp2 = AnnounceResponse {
389 transaction_id: 2,
390 interval: 200,
391 leechers: 8,
392 seeders: 3,
393 peers: vec![
394 ("192.168.1.2".to_string(), 6882),
395 ("192.168.1.3".to_string(), 6883),
396 ],
397 };
398
399 let all = UdpTrackerManager::collect_all_peers(&[resp1, resp2]);
400 assert_eq!(all.len(), 3, "Deduplicated peers should be 3");
401 }
402
403 #[tokio::test]
404 async fn test_shared_manager_creation() {
405 let shared = UdpTrackerClient::create_shared(0).await.unwrap();
406 let shared_mgr = UdpTrackerManager::create_shared(shared);
407 assert!(shared_mgr.try_lock().is_ok());
408 }
409}