aria2_core/engine/
udp_tracker_client.rs

1use std::collections::{HashMap, VecDeque};
2use std::net::SocketAddr;
3use std::sync::Arc;
4use std::time::{Duration, Instant};
5
6use tokio::sync::Mutex;
7use tracing::{debug, info, warn};
8
9use aria2_protocol::bittorrent::tracker::udp_tracker_protocol::{
10    AnnounceResponse, CONNECTION_TIMEOUT_SECS, ConnectResponse, ScrapeResult, UdpAction, UdpError,
11    UdpEvent, UdpState, build_announce_request, build_connect_request, build_scrape_request,
12    parse_announce_response, parse_connect_response, parse_scrape_response,
13};
14
15const REQUEST_TIMEOUT_SECS: u64 = 15;
16const MAX_RETRIES: u32 = 3;
17
18struct ConnectionState {
19    id: u64,
20    updated_at: Instant,
21}
22
23#[derive(Debug)]
24pub struct UdpTrackerRequest {
25    pub remote_addr: SocketAddr,
26    pub info_hash: [u8; 20],
27    pub peer_id: [u8; 20],
28    pub downloaded: i64,
29    pub left: i64,
30    pub uploaded: i64,
31    pub event: UdpEvent,
32    pub num_want: i32,
33    pub port: u16,
34    pub state: UdpState,
35    pub error: Option<UdpError>,
36    pub dispatched_at: Option<Instant>,
37    pub fail_count: u32,
38    pub reply: Option<AnnounceResponse>,
39    /// Scrape results populated when this is a scrape request
40    pub scrape_results: Option<Vec<ScrapeResult>>,
41    /// Info hashes for scrape requests (can be multiple)
42    pub scrape_info_hashes: Vec<[u8; 20]>,
43    txn_id: u32,
44}
45
46impl UdpTrackerRequest {
47    #[allow(clippy::too_many_arguments)]
48    fn new(
49        addr: SocketAddr,
50        info_hash: [u8; 20],
51        peer_id: [u8; 20],
52        downloaded: i64,
53        left: i64,
54        uploaded: i64,
55        event: UdpEvent,
56        num_want: i32,
57        port: u16,
58    ) -> Self {
59        Self {
60            remote_addr: addr,
61            info_hash,
62            peer_id,
63            downloaded,
64            left,
65            uploaded,
66            event,
67            num_want,
68            port,
69            state: UdpState::Pending,
70            error: None,
71            dispatched_at: None,
72            fail_count: 0,
73            reply: None,
74            scrape_results: None,
75            scrape_info_hashes: Vec::new(),
76            txn_id: 0,
77        }
78    }
79}
80
81pub struct UdpTrackerClient {
82    socket: Arc<tokio::net::UdpSocket>,
83    conn_cache: HashMap<SocketAddr, ConnectionState>,
84    pending: VecDeque<UdpTrackerRequest>,
85    inflight: VecDeque<UdpTrackerRequest>,
86    waiting_for_conn: VecDeque<UdpTrackerRequest>,
87    txn_map: HashMap<u32, usize>,
88    next_txn_id: u32,
89}
90
91impl UdpTrackerClient {
92    pub async fn new(bind_port: u16) -> Result<Self, String> {
93        let addr = format!("0.0.0.0:{}", bind_port);
94        let socket = tokio::net::UdpSocket::bind(&addr)
95            .await
96            .map_err(|e| format!("UDP bind failed on {}: {}", addr, e))?;
97
98        info!("UdpTrackerClient bound to {}", addr);
99
100        Ok(Self {
101            socket: Arc::new(socket),
102            conn_cache: HashMap::new(),
103            pending: VecDeque::new(),
104            inflight: VecDeque::new(),
105            waiting_for_conn: VecDeque::new(),
106            txn_map: HashMap::new(),
107            next_txn_id: Self::initial_txn_id(),
108        })
109    }
110
111    #[allow(clippy::too_many_arguments)]
112    pub async fn add_announce(
113        &mut self,
114        addr: &SocketAddr,
115        info_hash: &[u8; 20],
116        peer_id: &[u8; 20],
117        downloaded: i64,
118        left: i64,
119        uploaded: i64,
120        event: UdpEvent,
121        num_want: i32,
122        port: u16,
123    ) {
124        let req = UdpTrackerRequest::new(
125            *addr, *info_hash, *peer_id, downloaded, left, uploaded, event, num_want, port,
126        );
127        self.pending.push_back(req);
128        debug!("Added announce request for {}", addr);
129    }
130
131    /// Add a SCRAPE request to query statistics for one or more info hashes
132    ///
133    /// # Arguments
134    /// * `addr` - UDP tracker socket address
135    /// * `info_hashes` - Slice of 20-byte info hashes to query (max ~74 per request)
136    pub async fn add_scrape(&mut self, addr: &SocketAddr, info_hashes: &[[u8; 20]]) {
137        // Use first info hash for the request struct (scrape can have multiple)
138        let first_ih = if info_hashes.is_empty() {
139            [0u8; 20]
140        } else {
141            info_hashes[0]
142        };
143
144        let mut req =
145            UdpTrackerRequest::new(*addr, first_ih, [0u8; 20], 0, 0, 0, UdpEvent::None, 0, 0);
146        req.scrape_info_hashes = info_hashes.to_vec();
147        self.pending.push_back(req);
148        debug!(
149            "Added scrape request for {} ({} hashes)",
150            addr,
151            info_hashes.len()
152        );
153    }
154
155    pub async fn process_one(&mut self) -> bool {
156        loop {
157            if self.pending.is_empty() && self.waiting_for_conn.is_empty() {
158                return false;
159            }
160
161            if let Some(mut req) = self.pending.pop_front() {
162                let host_key = req.remote_addr;
163
164                if let Some(conn) = self.conn_cache.get(&host_key) {
165                    if conn.updated_at.elapsed().as_secs() < CONNECTION_TIMEOUT_SECS {
166                        // Route to appropriate send method based on request type
167                        if !req.scrape_info_hashes.is_empty() {
168                            return self.send_scrape(&mut req, conn.id).await;
169                        }
170                        return self.send_announce(&mut req, conn.id).await;
171                    } else {
172                        self.conn_cache.remove(&host_key);
173                        debug!("Connection cache expired for {}", host_key);
174                    }
175                }
176
177                if !self.is_connecting_to(&host_key) {
178                    self.waiting_for_conn.push_back(req);
179                    return self.send_connect(host_key).await;
180                }
181
182                self.waiting_for_conn.push_back(req);
183                debug!("Waiting for connection to {}", host_key);
184            } else if let Some(req) = self.waiting_for_conn.pop_front() {
185                self.pending.push_front(req);
186            } else {
187                return false;
188            }
189        }
190    }
191
192    async fn send_announce(&mut self, req: &mut UdpTrackerRequest, conn_id: u64) -> bool {
193        let txn_id = self.next_txn();
194        req.txn_id = txn_id;
195        req.dispatched_at = Some(Instant::now());
196        req.state = UdpState::Pending;
197
198        let payload = build_announce_request(
199            conn_id,
200            txn_id,
201            &req.info_hash,
202            &req.peer_id,
203            req.downloaded,
204            req.left,
205            req.uploaded,
206            req.event,
207            0,
208            0,
209            req.num_want,
210            req.port,
211        );
212
213        match self.socket.send_to(&payload, req.remote_addr).await {
214            Ok(len) => {
215                self.txn_map.insert(txn_id, self.inflight.len());
216                self.inflight.push_back(std::mem::replace(
217                    req,
218                    UdpTrackerRequest::new(
219                        req.remote_addr,
220                        req.info_hash,
221                        req.peer_id,
222                        req.downloaded,
223                        req.left,
224                        req.uploaded,
225                        req.event,
226                        req.num_want,
227                        req.port,
228                    ),
229                ));
230                debug!(
231                    "Sent ANNOUNCE {} bytes to {} (txn={})",
232                    len, req.remote_addr, txn_id
233                );
234                true
235            }
236            Err(e) => {
237                warn!("Send ANNOUNCE to {} failed: {}", req.remote_addr, e);
238                req.fail_count += 1;
239                req.error = Some(UdpError::Network);
240                if req.fail_count < MAX_RETRIES {
241                    self.pending.push_front(std::mem::replace(
242                        req,
243                        UdpTrackerRequest::new(
244                            req.remote_addr,
245                            req.info_hash,
246                            req.peer_id,
247                            req.downloaded,
248                            req.left,
249                            req.uploaded,
250                            req.event,
251                            req.num_want,
252                            req.port,
253                        ),
254                    ));
255                }
256                true
257            }
258        }
259    }
260
261    async fn send_scrape(&mut self, req: &mut UdpTrackerRequest, conn_id: u64) -> bool {
262        let txn_id = self.next_txn();
263        req.txn_id = txn_id;
264        req.dispatched_at = Some(Instant::now());
265        req.state = UdpState::Pending;
266
267        // Build scrape payload with all info hashes from this request
268        let hashes: Vec<[u8; 20]> = req.scrape_info_hashes.clone();
269        if hashes.is_empty() {
270            warn!("Scrape request with no info hashes for {}", req.remote_addr);
271            return true;
272        }
273
274        let payload = build_scrape_request(conn_id, txn_id, &hashes);
275
276        match self.socket.send_to(&payload, req.remote_addr).await {
277            Ok(len) => {
278                self.txn_map.insert(txn_id, self.inflight.len());
279                // Preserve scrape_info_hashes when replacing the request
280                let mut replacement = UdpTrackerRequest::new(
281                    req.remote_addr,
282                    req.info_hash,
283                    req.peer_id,
284                    req.downloaded,
285                    req.left,
286                    req.uploaded,
287                    req.event,
288                    req.num_want,
289                    req.port,
290                );
291                replacement.scrape_info_hashes = std::mem::take(&mut req.scrape_info_hashes);
292                self.inflight.push_back(std::mem::replace(req, replacement));
293                debug!(
294                    "Sent SCRAPE {} bytes to {} (txn={}, {} hashes)",
295                    len,
296                    req.remote_addr,
297                    txn_id,
298                    hashes.len()
299                );
300                true
301            }
302            Err(e) => {
303                warn!("Send SCRAPE to {} failed: {}", req.remote_addr, e);
304                req.fail_count += 1;
305                req.error = Some(UdpError::Network);
306                if req.fail_count < MAX_RETRIES {
307                    self.pending.push_front(std::mem::replace(
308                        req,
309                        UdpTrackerRequest::new(
310                            req.remote_addr,
311                            req.info_hash,
312                            req.peer_id,
313                            req.downloaded,
314                            req.left,
315                            req.uploaded,
316                            req.event,
317                            req.num_want,
318                            req.port,
319                        ),
320                    ));
321                }
322                true
323            }
324        }
325    }
326
327    async fn send_connect(&mut self, addr: SocketAddr) -> bool {
328        let txn_id = self.next_txn();
329
330        let payload = build_connect_request(txn_id);
331
332        match self.socket.send_to(&payload, addr).await {
333            Ok(len) => {
334                let mut dummy_req = UdpTrackerRequest::new(
335                    addr,
336                    [0u8; 20],
337                    [0u8; 20],
338                    0,
339                    0,
340                    0,
341                    UdpEvent::None,
342                    0,
343                    0,
344                );
345                dummy_req.txn_id = txn_id;
346                dummy_req.dispatched_at = Some(Instant::now());
347                dummy_req.state = UdpState::Pending;
348                self.txn_map.insert(txn_id, self.inflight.len());
349                self.inflight.push_back(dummy_req);
350                debug!("Sent CONNECT {} bytes to {} (txn={})", len, addr, txn_id);
351                true
352            }
353            Err(e) => {
354                warn!("Send CONNECT to {} failed: {}", addr, e);
355                true
356            }
357        }
358    }
359
360    pub async fn handle_response(&mut self, data: &[u8], from: &SocketAddr) {
361        if data.len() < 4 {
362            warn!("Short response from {}: {} bytes", from, data.len());
363            return;
364        }
365
366        let action_val = i32::from_be_bytes([data[0], data[1], data[2], data[3]]);
367        let txn_id = if data.len() >= 8 {
368            u32::from_be_bytes([data[4], data[5], data[6], data[7]])
369        } else {
370            warn!("Response too short for txn_id from {}", from);
371            return;
372        };
373
374        let idx = match self.txn_map.remove(&txn_id) {
375            Some(i) => i,
376            None => {
377                debug!("Unknown txn_id {} from {}", txn_id, from);
378                return;
379            }
380        };
381
382        if idx >= self.inflight.len() {
383            warn!(
384                "Invalid index {} for txn_id {} (inflight={})",
385                idx,
386                txn_id,
387                self.inflight.len()
388            );
389            return;
390        }
391
392        let mut req = self.inflight.remove(idx).unwrap_or_else(|| {
393            UdpTrackerRequest::new(*from, [0u8; 20], [0u8; 20], 0, 0, 0, UdpEvent::None, 0, 0)
394        });
395
396        match UdpAction::from_i32(action_val) {
397            Some(UdpAction::Connect) => match parse_connect_response(data) {
398                Ok(ConnectResponse { connection_id, .. }) => {
399                    info!(
400                        "CONNECT response from {}, conn_id=0x{:016X}",
401                        from, connection_id
402                    );
403                    self.conn_cache.insert(
404                        *from,
405                        ConnectionState {
406                            id: connection_id,
407                            updated_at: Instant::now(),
408                        },
409                    );
410
411                    while let Some(waiting) = self.waiting_for_conn.pop_front() {
412                        self.pending.push_front(waiting);
413                    }
414                }
415                Err(e) => {
416                    warn!("Parse CONNECT response from {} failed: {}", from, e);
417                    req.error = Some(UdpError::TrackerError);
418                }
419            },
420            Some(UdpAction::Announce) => {
421                match parse_announce_response(data) {
422                    Ok(resp) => {
423                        info!(
424                            "ANNOUNCE response from {}: {} peers, interval={}s",
425                            from,
426                            resp.peers.len(),
427                            resp.interval
428                        );
429                        req.reply = Some(resp);
430                        req.state = UdpState::Complete;
431                        req.error = Some(UdpError::Success);
432                    }
433                    Err(e) => {
434                        warn!("Parse ANNOUNCE response from {} failed: {}", from, e);
435                        req.error = Some(UdpError::TrackerError);
436                    }
437                }
438                self.pending.push_back(req);
439            }
440            Some(UdpAction::Error) => {
441                let msg_len = (data.len() - 8).min(256);
442                let msg = String::from_utf8_lossy(&data[8..8 + msg_len]);
443                warn!("Tracker error from {}: {}", from, msg);
444                req.error = Some(UdpError::TrackerError);
445                self.pending.push_back(req);
446            }
447            Some(UdpAction::Scrape) => match parse_scrape_response(data) {
448                Ok(results) => {
449                    info!(
450                        "SCRAPE response from {}: {} info hashes scraped",
451                        from,
452                        results.len()
453                    );
454                    // Store scrape results in a dedicated collection
455                    for result in &results {
456                        debug!(
457                            "  seeders={} leechers={} completed={}",
458                            result.seeders, result.leechers, result.completed
459                        );
460                    }
461                    req.scrape_results = Some(results);
462                    req.state = UdpState::Complete;
463                    req.error = Some(UdpError::Success);
464                    self.pending.push_back(req);
465                }
466                Err(e) => {
467                    warn!("Parse SCRAPE response from {} failed: {}", from, e);
468                    req.error = Some(UdpError::TrackerError);
469                    self.pending.push_back(req);
470                }
471            },
472            _ => {
473                warn!("Unknown action {} from {}", action_val, from);
474                req.error = Some(UdpError::TrackerError);
475                self.pending.push_back(req);
476            }
477        }
478    }
479
480    pub async fn handle_timeouts(&mut self) {
481        let now = Instant::now();
482        let expired: Vec<usize> = self
483            .inflight
484            .iter()
485            .enumerate()
486            .filter(|(_, r)| {
487                r.dispatched_at.is_some_and(|t| {
488                    t.duration_since(now) > Duration::from_secs(REQUEST_TIMEOUT_SECS)
489                })
490            })
491            .map(|(i, _)| i)
492            .collect();
493
494        for idx in expired.into_iter().rev() {
495            if idx < self.inflight.len() {
496                let mut req = self.inflight.remove(idx).unwrap_or_else(|| {
497                    UdpTrackerRequest::new(
498                        std::net::SocketAddr::V4(std::net::SocketAddrV4::new(
499                            std::net::Ipv4Addr::UNSPECIFIED,
500                            0,
501                        )),
502                        [0u8; 20],
503                        [0u8; 20],
504                        0,
505                        0,
506                        0,
507                        UdpEvent::None,
508                        0,
509                        0,
510                    )
511                });
512                if req.txn_id != 0 {
513                    self.txn_map.remove(&req.txn_id);
514                }
515                req.fail_count += 1;
516                if req.fail_count < MAX_RETRIES {
517                    debug!(
518                        "Timeout retry {}/{} for txn_id={}",
519                        req.fail_count, MAX_RETRIES, req.txn_id
520                    );
521                    req.dispatched_at = None;
522                    self.pending.push_back(req);
523                } else {
524                    warn!("Max retries exceeded for txn_id={}", req.txn_id);
525                    req.error = Some(UdpError::Timeout);
526                }
527            }
528        }
529
530        let stale_addrs: Vec<SocketAddr> = self
531            .conn_cache
532            .iter()
533            .filter(|(_, s)| s.updated_at.elapsed().as_secs() > CONNECTION_TIMEOUT_SECS)
534            .map(|(&a, _)| a)
535            .collect();
536
537        for addr in stale_addrs {
538            self.conn_cache.remove(&addr);
539            debug!("Removed stale connection cache for {}", addr);
540        }
541    }
542
543    pub fn no_pending(&self) -> bool {
544        self.pending.is_empty() && self.inflight.is_empty() && self.waiting_for_conn.is_empty()
545    }
546
547    pub fn completed_requests(&self) -> Vec<&AnnounceResponse> {
548        self.pending
549            .iter()
550            .filter_map(|r| r.reply.as_ref())
551            .collect()
552    }
553
554    /// Get all completed scrape results from pending requests
555    pub fn completed_scrape_results(&self) -> Vec<&Vec<ScrapeResult>> {
556        self.pending
557            .iter()
558            .filter_map(|r| r.scrape_results.as_ref())
559            .collect()
560    }
561
562    pub fn socket(&self) -> Arc<tokio::net::UdpSocket> {
563        Arc::clone(&self.socket)
564    }
565
566    fn is_connecting_to(&self, addr: &SocketAddr) -> bool {
567        self.inflight
568            .iter()
569            .any(|r| &r.remote_addr == addr && r.reply.is_none())
570    }
571
572    pub(crate) fn next_txn(&mut self) -> u32 {
573        let id = self.next_txn_id;
574        self.next_txn_id = id.wrapping_add(1);
575        if self.next_txn_id == 0 {
576            self.next_txn_id = 1;
577        }
578        id
579    }
580
581    fn initial_txn_id() -> u32 {
582        use std::time::{SystemTime, UNIX_EPOCH};
583        let dur = SystemTime::now()
584            .duration_since(UNIX_EPOCH)
585            .unwrap_or_default();
586        ((dur.as_nanos() & 0xFFFFFFFF) as u32).max(1)
587    }
588}
589
590pub type SharedUdpClient = Arc<Mutex<UdpTrackerClient>>;
591
592impl UdpTrackerClient {
593    pub async fn create_shared(bind_port: u16) -> Result<SharedUdpClient, String> {
594        let client = Self::new(bind_port).await?;
595        Ok(Arc::new(Mutex::new(client)))
596    }
597}
598
599#[cfg(test)]
600mod tests {
601    use super::*;
602
603    #[tokio::test]
604    async fn test_client_creation() {
605        let client = UdpTrackerClient::new(0).await;
606        assert!(
607            client.is_ok(),
608            "UDP client creation should succeed with port 0"
609        );
610        let c = client.unwrap();
611        assert!(c.no_pending());
612        assert!(c.completed_requests().is_empty());
613    }
614
615    #[tokio::test]
616    async fn test_add_announce_request() {
617        let mut client = UdpTrackerClient::new(0).await.unwrap();
618        let addr: SocketAddr = "127.0.0.1:6969".parse().unwrap();
619        let ih = [0xABu8; 20];
620        let pid = [0xCDu8; 20];
621
622        client
623            .add_announce(&addr, &ih, &pid, 0, 1000, 0, UdpEvent::Started, 50, 6881)
624            .await;
625        assert_eq!(client.pending.len(), 1);
626
627        client
628            .add_announce(&addr, &ih, &pid, 500, 500, 0, UdpEvent::None, -1, 6881)
629            .await;
630        assert_eq!(client.pending.len(), 2);
631    }
632
633    #[tokio::test]
634    async fn test_process_one_needs_connection() {
635        let mut client = UdpTrackerClient::new(0).await.unwrap();
636        let addr: SocketAddr = "127.0.0.1:6969".parse().unwrap();
637        let ih = [0x12u8; 20];
638        let pid = [0x34u8; 20];
639
640        client
641            .add_announce(&addr, &ih, &pid, 0, 1000, 0, UdpEvent::Started, 50, 6881)
642            .await;
643        let processed = client.process_one().await;
644        assert!(processed, "Should have processed the connect step");
645        assert!(
646            !client.inflight.is_empty(),
647            "Should have an in-flight CONNECT"
648        );
649    }
650
651    #[tokio::test]
652    async fn test_no_pending_returns_false_when_empty() {
653        let mut client = UdpTrackerClient::new(0).await.unwrap();
654        assert!(client.no_pending());
655        let processed = client.process_one().await;
656        assert!(!processed, "process_one should return false when empty");
657    }
658
659    #[tokio::test]
660    async fn test_handle_connect_response() {
661        let mut client = UdpTrackerClient::new(0).await.unwrap();
662        let addr: SocketAddr = "127.0.0.1:6969".parse().unwrap();
663        let ih = [0x11u8; 20];
664        let pid = [0x22u8; 20];
665
666        client
667            .add_announce(&addr, &ih, &pid, 0, 1000, 0, UdpEvent::Started, 50, 6881)
668            .await;
669        client.process_one().await;
670
671        let mut resp_data = vec![0u8; 16];
672        resp_data[0..4].copy_from_slice(&0i32.to_be_bytes());
673        let txn_id = client.inflight.front().map(|r| r.txn_id).unwrap_or(0);
674        resp_data[4..8].copy_from_slice(&txn_id.to_be_bytes());
675        resp_data[8..16].copy_from_slice(&0x123456789ABCDEF0u64.to_be_bytes());
676
677        client.handle_response(&resp_data, &addr).await;
678        assert!(
679            client.conn_cache.contains_key(&addr),
680            "Should cache connection after CONNECT response"
681        );
682    }
683
684    #[tokio::test]
685    async fn test_handle_announce_response_with_peers() {
686        let mut client = UdpTrackerClient::new(0).await.unwrap();
687        let addr: SocketAddr = "127.0.0.1:6969".parse().unwrap();
688
689        let ih = [0x33u8; 20];
690        let pid = [0x44u8; 20];
691        let txn_id = client.next_txn();
692
693        client.txn_map.insert(txn_id, 0);
694        let mut dummy_req =
695            UdpTrackerRequest::new(addr, ih, pid, 0, 1000, 0, UdpEvent::Started, 50, 6881);
696        dummy_req.txn_id = txn_id;
697        dummy_req.dispatched_at = Some(Instant::now());
698        client.inflight.push_back(dummy_req);
699
700        let mut resp_data = vec![0u8; 26];
701        resp_data[0..4].copy_from_slice(&1i32.to_be_bytes());
702        resp_data[4..8].copy_from_slice(&txn_id.to_be_bytes());
703        resp_data[8..12].copy_from_slice(&900u32.to_be_bytes());
704        resp_data[12..16].copy_from_slice(&5u32.to_be_bytes());
705        resp_data[16..20].copy_from_slice(&3u32.to_be_bytes());
706        resp_data.extend_from_slice(&[10, 0, 0, 1, 0x1A, 0x04, 192, 168, 1, 100, 0x1F, 0x90]);
707
708        client.handle_response(&resp_data, &addr).await;
709
710        let completed = client.completed_requests();
711        assert!(
712            !completed.is_empty(),
713            "Should have at least one completed announce"
714        );
715        assert!(
716            completed[0].peers.len() >= 2,
717            "Should have at least 2 peers"
718        );
719        assert_eq!(completed[0].interval, 900);
720        assert_eq!(completed[0].leechers, 5);
721        assert_eq!(completed[0].seeders, 3);
722    }
723
724    #[tokio::test]
725    async fn test_handle_error_response() {
726        let mut client = UdpTrackerClient::new(0).await.unwrap();
727        let addr: SocketAddr = "127.0.0.1:6969".parse().unwrap();
728        let ih = [0x55u8; 20];
729        let pid = [0x66u8; 20];
730
731        client
732            .add_announce(&addr, &ih, &pid, 0, 1000, 0, UdpEvent::Started, 50, 6881)
733            .await;
734        client.process_one().await;
735
736        let txn_id = client.inflight.front().map(|r| r.txn_id).unwrap_or(0);
737        let mut err_data = vec![0u8; 23];
738        err_data[0..4].copy_from_slice(&3i32.to_be_bytes());
739        err_data[4..8].copy_from_slice(&txn_id.to_be_bytes());
740        err_data[8..23].copy_from_slice(b"tracker offline");
741
742        client.handle_response(&err_data, &addr).await;
743        assert!(
744            !client.conn_cache.contains_key(&addr),
745            "Error should not create cache entry"
746        );
747    }
748
749    #[tokio::test]
750    async fn test_timeout_cleaning() {
751        let mut client = UdpTrackerClient::new(0).await.unwrap();
752        let addr: SocketAddr = "127.0.0.1:6969".parse().unwrap();
753        let ih = [0x77u8; 20];
754        let pid = [0x88u8; 20];
755
756        client
757            .add_announce(&addr, &ih, &pid, 0, 1000, 0, UdpEvent::Started, 50, 6881)
758            .await;
759        client.process_one().await;
760        assert_eq!(client.inflight.len(), 1);
761
762        tokio::time::sleep(Duration::from_millis(100)).await;
763        client.handle_timeouts().await;
764        assert_eq!(client.inflight.len(), 1, "Not yet timed out");
765
766        for req in &mut client.inflight {
767            req.dispatched_at =
768                Some(Instant::now() - Duration::from_secs(REQUEST_TIMEOUT_SECS + 1));
769        }
770        client.handle_timeouts().await;
771        assert!(
772            client.inflight.is_empty()
773                || !client.pending.is_empty()
774                || !client.waiting_for_conn.is_empty(),
775            "Timed-out request should be moved"
776        );
777    }
778
779    #[tokio::test]
780    async fn test_txn_id_generation() {
781        let mut client = UdpTrackerClient::new(0).await.unwrap();
782        let mut txn_ids = Vec::new();
783        for _ in 0..5 {
784            txn_ids.push(client.next_txn());
785            client.next_txn();
786        }
787        let unique: std::collections::HashSet<_> = txn_ids.iter().cloned().collect();
788        assert_eq!(
789            unique.len(),
790            txn_ids.len(),
791            "Transaction IDs should all be unique"
792        );
793    }
794
795    #[tokio::test]
796    async fn test_shared_client_creation() {
797        let shared = UdpTrackerClient::create_shared(0).await;
798        assert!(shared.is_ok(), "Shared client creation should succeed");
799    }
800
801    // --- Scrape tests ---
802
803    #[tokio::test]
804    async fn test_add_scrape_request() {
805        let mut client = UdpTrackerClient::new(0).await.unwrap();
806        let addr: SocketAddr = "127.0.0.1:6969".parse().unwrap();
807        let hashes = [[0xAAu8; 20], [0xBBu8; 20], [0xCCu8; 20]];
808
809        client.add_scrape(&addr, &hashes).await;
810        assert_eq!(client.pending.len(), 1);
811
812        let req = &client.pending[0];
813        assert!(!req.scrape_info_hashes.is_empty());
814        assert_eq!(req.scrape_info_hashes.len(), 3);
815        assert_eq!(req.scrape_info_hashes[0], [0xAAu8; 20]);
816    }
817
818    #[tokio::test]
819    async fn test_handle_scrape_response_single_hash() {
820        let mut client = UdpTrackerClient::new(0).await.unwrap();
821        let addr: SocketAddr = "127.0.0.1:6969".parse().unwrap();
822        let hashes = [[0x11u8; 20]];
823
824        // Manually set up an in-flight scrape request
825        let txn_id = client.next_txn();
826        let mut req =
827            UdpTrackerRequest::new(addr, hashes[0], [0u8; 20], 0, 0, 0, UdpEvent::None, 0, 0);
828        req.txn_id = txn_id;
829        req.dispatched_at = Some(Instant::now());
830        req.scrape_info_hashes = hashes.to_vec();
831        client.inflight.push_back(req);
832        client.txn_map.insert(txn_id, 0);
833
834        // Build scrape response: action=2, txn_id, seeders=42, leechers=10, completed=999
835        let mut resp_data = vec![0u8; 20];
836        resp_data[0..4].copy_from_slice(&(UdpAction::Scrape as i32).to_be_bytes());
837        resp_data[4..8].copy_from_slice(&txn_id.to_be_bytes());
838        resp_data[8..12].copy_from_slice(&42u32.to_be_bytes());
839        resp_data[12..16].copy_from_slice(&10u32.to_be_bytes());
840        resp_data[16..20].copy_from_slice(&999u32.to_be_bytes());
841
842        client.handle_response(&resp_data, &addr).await;
843
844        let scrape_results = client.completed_scrape_results();
845        assert_eq!(scrape_results.len(), 1);
846        assert_eq!(scrape_results[0].len(), 1);
847        assert_eq!(scrape_results[0][0].seeders, 42);
848        assert_eq!(scrape_results[0][0].leechers, 10);
849        assert_eq!(scrape_results[0][0].completed, 999);
850    }
851
852    #[tokio::test]
853    async fn test_handle_scrape_response_multi_hash() {
854        let mut client = UdpTrackerClient::new(0).await.unwrap();
855        let addr: SocketAddr = "127.0.0.1:6969".parse().unwrap();
856        let hashes = [[0x22u8; 20], [0x33u8; 20]];
857
858        let txn_id = client.next_txn();
859        let mut req =
860            UdpTrackerRequest::new(addr, hashes[0], [0u8; 20], 0, 0, 0, UdpEvent::None, 0, 0);
861        req.txn_id = txn_id;
862        req.dispatched_at = Some(Instant::now());
863        req.scrape_info_hashes = hashes.to_vec();
864        client.inflight.push_back(req);
865        client.txn_map.insert(txn_id, 0);
866
867        // Response for 2 info hashes
868        let mut resp_data = vec![0u8; 32]; // 8 header + 12*2
869        resp_data[0..4].copy_from_slice(&(UdpAction::Scrape as i32).to_be_bytes());
870        resp_data[4..8].copy_from_slice(&txn_id.to_be_bytes());
871        // Hash 1
872        resp_data[8..12].copy_from_slice(&100u32.to_be_bytes());
873        resp_data[12..16].copy_from_slice(&50u32.to_be_bytes());
874        resp_data[16..20].copy_from_slice(&200u32.to_be_bytes());
875        // Hash 2
876        resp_data[20..24].copy_from_slice(&5u32.to_be_bytes());
877        resp_data[24..28].copy_from_slice(&3u32.to_be_bytes());
878        resp_data[28..32].copy_from_slice(&7u32.to_be_bytes());
879
880        client.handle_response(&resp_data, &addr).await;
881
882        let scrape_results = client.completed_scrape_results();
883        assert_eq!(scrape_results.len(), 1);
884        assert_eq!(scrape_results[0].len(), 2);
885        assert_eq!(scrape_results[0][0].seeders, 100);
886        assert_eq!(scrape_results[0][1].seeders, 5);
887    }
888
889    #[tokio::test]
890    async fn test_scrape_error_action_returns_error() {
891        let mut client = UdpTrackerClient::new(0).await.unwrap();
892        let addr: SocketAddr = "127.0.0.1:6969".parse().unwrap();
893
894        let txn_id = client.next_txn();
895        let mut req =
896            UdpTrackerRequest::new(addr, [0x99u8; 20], [0u8; 20], 0, 0, 0, UdpEvent::None, 0, 0);
897        req.txn_id = txn_id;
898        req.dispatched_at = Some(Instant::now());
899        req.scrape_info_hashes = vec![[0x99u8; 20]];
900        client.inflight.push_back(req);
901        client.txn_map.insert(txn_id, 0);
902
903        // Send error action instead of scrape action
904        let mut err_data = vec![0u8; 23];
905        err_data[0..4].copy_from_slice(&3i32.to_be_bytes()); // Error action
906        err_data[4..8].copy_from_slice(&txn_id.to_be_bytes());
907        err_data[8..23].copy_from_slice(b"scrape failed!!");
908
909        client.handle_response(&err_data, &addr).await;
910
911        // Should not have successful scrape results
912        let scrape_results = client.completed_scrape_results();
913        assert!(
914            scrape_results.is_empty(),
915            "Error response should not produce scrape results"
916        );
917    }
918}
aria2_core/engine/udp_tracker_client.rs

aria2_core/engine/
udp_tracker_client.rs
