1use std::fmt;
2use std::net::SocketAddr;
3use std::time::Duration;
4use tokio::time::timeout;
5use tracing::{debug, info, warn};
6
7use aria2_protocol::bittorrent::bencode::codec::BencodeValue;
8use aria2_protocol::bittorrent::extension::ut_metadata::{
9 ExtensionHandshake, MetadataCollector, UtMetadataMsg,
10};
11use aria2_protocol::bittorrent::peer::connection::{PeerAddr, PeerConnection};
12
13const METADATA_MAX_SIZE: u64 = 100 * 1024 * 1024;
14const PIECE_SIZE_MIN: u32 = 1024;
15const PIECE_SIZE_MAX: u32 = 65536;
16const DEFAULT_MAX_ATTEMPTS: usize = 3;
17
18#[derive(Debug, Clone)]
19pub enum MetadataExchangeError {
20 NoPeersAvailable,
21 AllPeersFailed { attempts: usize, last_error: String },
22 PeerConnectFailed { addr: String, reason: String },
23 PeerTimeout { addr: String },
24 UnsupportedPeer { addr: String, reason: String },
25 InvalidMetadataSize { size: u64 },
26 MetadataTooLarge { size: u64, max: u64 },
27 BencodeDecodeFailed { detail: String },
28 PieceRejected { piece: u32 },
29 PieceTimeout { piece: u32 },
30 IncompleteMetadata { expected: u64, received: u64 },
31 IoError(String),
32}
33
34impl MetadataExchangeError {
35 fn is_fatal(&self) -> bool {
36 matches!(
37 self,
38 MetadataExchangeError::MetadataTooLarge { .. }
39 | MetadataExchangeError::BencodeDecodeFailed { .. }
40 | MetadataExchangeError::NoPeersAvailable
41 )
42 }
43
44 pub fn addr(&self) -> Option<&str> {
45 match self {
46 MetadataExchangeError::PeerConnectFailed { addr, .. } => Some(addr),
47 MetadataExchangeError::PeerTimeout { addr } => Some(addr),
48 MetadataExchangeError::UnsupportedPeer { addr, .. } => Some(addr),
49 _ => None,
50 }
51 }
52}
53
54impl fmt::Display for MetadataExchangeError {
55 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
56 match self {
57 MetadataExchangeError::NoPeersAvailable => {
58 write!(f, "No peers available for metadata fetch")
59 }
60 MetadataExchangeError::AllPeersFailed {
61 attempts,
62 last_error,
63 } => {
64 write!(
65 f,
66 "All {} peers failed, last error: {}",
67 attempts, last_error
68 )
69 }
70 MetadataExchangeError::PeerConnectFailed { addr, reason } => {
71 write!(f, "Connect to {} failed: {}", addr, reason)
72 }
73 MetadataExchangeError::PeerTimeout { addr } => {
74 write!(f, "Connect to {} timed out", addr)
75 }
76 MetadataExchangeError::UnsupportedPeer { addr, reason } => {
77 write!(f, "Peer {} unsupported: {}", addr, reason)
78 }
79 MetadataExchangeError::InvalidMetadataSize { size } => {
80 write!(f, "Invalid metadata_size: {}", size)
81 }
82 MetadataExchangeError::MetadataTooLarge { size, max } => {
83 write!(f, "metadata_size too large: {} (max {})", size, max)
84 }
85 MetadataExchangeError::BencodeDecodeFailed { detail } => {
86 write!(f, "Bencode decode failed: {}", detail)
87 }
88 MetadataExchangeError::PieceRejected { piece } => {
89 write!(f, "Piece {} rejected by peer", piece)
90 }
91 MetadataExchangeError::PieceTimeout { piece } => {
92 write!(f, "ut_metadata timeout for piece {}", piece)
93 }
94 MetadataExchangeError::IncompleteMetadata { expected, received } => {
95 write!(
96 f,
97 "Incomplete metadata collection: expected {} bytes, received {}",
98 expected, received
99 )
100 }
101 MetadataExchangeError::IoError(msg) => write!(f, "IO error: {}", msg),
102 }
103 }
104}
105
106impl std::error::Error for MetadataExchangeError {}
107
108pub struct MetadataExchangeConfig {
109 pub max_peers_to_try: usize,
110 pub connect_timeout: Duration,
111 pub request_timeout: Duration,
112 pub piece_size: u32,
113 pub max_attempts: usize,
114}
115
116impl Default for MetadataExchangeConfig {
117 fn default() -> Self {
118 Self {
119 max_peers_to_try: 5,
120 connect_timeout: Duration::from_secs(15),
121 request_timeout: Duration::from_secs(10),
122 piece_size: 16 * 1024,
123 max_attempts: DEFAULT_MAX_ATTEMPTS,
124 }
125 }
126}
127
128impl MetadataExchangeConfig {
129 pub fn with_piece_size(mut self, size: u32) -> Self {
130 if !(PIECE_SIZE_MIN..=PIECE_SIZE_MAX).contains(&size) {
131 warn!(
132 "piece_size={} is out of valid range [{}-{}], clamping",
133 size, PIECE_SIZE_MIN, PIECE_SIZE_MAX
134 );
135 self.piece_size = size.clamp(PIECE_SIZE_MIN, PIECE_SIZE_MAX);
136 } else {
137 self.piece_size = size;
138 }
139 self
140 }
141
142 pub fn with_max_attempts(mut self, attempts: usize) -> Self {
143 self.max_attempts = attempts.max(1);
144 self
145 }
146}
147
148pub struct MetadataExchangeSession {
149 config: MetadataExchangeConfig,
150}
151
152impl MetadataExchangeSession {
153 pub fn new(config: MetadataExchangeConfig) -> Self {
154 Self { config }
155 }
156
157 pub async fn fetch_metadata(
158 &self,
159 info_hash: &[u8; 20],
160 peers: &[SocketAddr],
161 ) -> Result<Vec<u8>, MetadataExchangeError> {
162 if peers.is_empty() {
163 return Err(MetadataExchangeError::NoPeersAvailable);
164 }
165
166 let mut attempt_count = 0usize;
167 let mut last_error = String::new();
168
169 for peer_addr in peers.iter().take(self.config.max_peers_to_try) {
170 if attempt_count >= self.config.max_attempts {
171 break;
172 }
173
174 match self.exchange_with_peer(info_hash, peer_addr).await {
175 Ok(torrent_bytes) => {
176 info!(
177 "Metadata fetched successfully from {} ({} bytes)",
178 peer_addr,
179 torrent_bytes.len()
180 );
181 return Ok(torrent_bytes);
182 }
183 Err(e) => {
184 attempt_count += 1;
185 last_error = e.to_string();
186
187 if e.is_fatal() {
188 warn!("Fatal metadata exchange error with {}: {}", peer_addr, e);
189 return Err(e);
190 }
191
192 warn!(
193 "Recoverable metadata exchange error with {} (attempt {}/{}): {}",
194 peer_addr, attempt_count, self.config.max_attempts, e
195 );
196 }
197 }
198 }
199
200 Err(MetadataExchangeError::AllPeersFailed {
201 attempts: attempt_count,
202 last_error,
203 })
204 }
205
206 async fn exchange_with_peer(
207 &self,
208 info_hash: &[u8; 20],
209 peer_addr: &SocketAddr,
210 ) -> Result<Vec<u8>, MetadataExchangeError> {
211 let addr_str = peer_addr.to_string();
212 let addr = PeerAddr::new(&peer_addr.ip().to_string(), peer_addr.port());
213
214 let conn_result = timeout(
215 self.config.connect_timeout,
216 PeerConnection::connect(&addr, info_hash),
217 )
218 .await;
219 let mut conn = match conn_result {
220 Ok(Ok(c)) => c,
221 Ok(Err(e)) => {
222 return Err(MetadataExchangeError::PeerConnectFailed {
223 addr: addr_str.clone(),
224 reason: e.to_string(),
225 });
226 }
227 Err(_) => {
228 return Err(MetadataExchangeError::PeerTimeout {
229 addr: addr_str.clone(),
230 });
231 }
232 };
233
234 debug!("Connected to {}, sending extension handshake", peer_addr);
235
236 let hs_payload = ExtensionHandshake::new(0).to_bencode();
237 let hs_encoded = hs_payload.encode();
238
239 conn.stream_write(&[0])
240 .await
241 .map_err(|e| MetadataExchangeError::IoError(format!("stream_write failed: {}", e)))?;
242 conn.stream_write(&hs_encoded)
243 .await
244 .map_err(|e| MetadataExchangeError::IoError(format!("stream_write failed: {}", e)))?;
245 conn.stream_flush()
246 .await
247 .map_err(|e| MetadataExchangeError::IoError(format!("stream_flush failed: {}", e)))?;
248
249 debug!("Extension handshake sent to {}", peer_addr);
250
251 let remote_hs_data = self.read_extension_message(&mut conn).await?;
252 let remote_hs = ExtensionHandshake::parse(&remote_hs_data).ok_or_else(|| {
253 MetadataExchangeError::BencodeDecodeFailed {
254 detail: "Failed to parse remote extension handshake".to_string(),
255 }
256 })?;
257
258 let metadata_size = match remote_hs.metadata_size {
259 Some(size) => size,
260 None => {
261 warn!(
262 "Peer {} reported metadata_size=None, skipping...",
263 peer_addr
264 );
265 return Err(MetadataExchangeError::UnsupportedPeer {
266 addr: addr_str,
267 reason: "Remote did not provide metadata_size".to_string(),
268 });
269 }
270 };
271
272 if metadata_size == 0 {
273 warn!("Peer {} reported metadata_size=0, skipping...", peer_addr);
274 return Err(MetadataExchangeError::InvalidMetadataSize { size: 0 });
275 }
276
277 if metadata_size > METADATA_MAX_SIZE {
278 return Err(MetadataExchangeError::MetadataTooLarge {
279 size: metadata_size,
280 max: METADATA_MAX_SIZE,
281 });
282 }
283
284 debug!("Remote reports metadata_size={} bytes", metadata_size);
285
286 let num_pieces = metadata_size.div_ceil(self.config.piece_size as u64) as u32;
287 let mut collector = MetadataCollector::new(metadata_size, self.config.piece_size);
288
289 for piece_idx in 0..num_pieces {
290 if collector.is_complete() {
291 break;
292 }
293
294 let req_msg = UtMetadataMsg::Request(piece_idx);
295 let encoded = req_msg.encode(20);
296
297 conn.stream_write(&encoded).await.map_err(|e| {
298 MetadataExchangeError::IoError(format!("stream_write failed: {}", e))
299 })?;
300 conn.stream_flush().await.map_err(|e| {
301 MetadataExchangeError::IoError(format!("stream_flush failed: {}", e))
302 })?;
303
304 match timeout(
305 self.config.request_timeout,
306 self.read_ut_metadata_response(&mut conn),
307 )
308 .await
309 {
310 Ok(Ok(UtMetadataMsg::Data(recv_piece, data))) => {
311 collector.add_piece(recv_piece, &data);
312 debug!(
313 "Received piece {}/{} ({} bytes)",
314 recv_piece + 1,
315 num_pieces,
316 data.len()
317 );
318 }
319 Ok(Ok(UtMetadataMsg::Reject(_))) => {
320 debug!("Piece {} rejected by {}", piece_idx, peer_addr);
321 return Err(MetadataExchangeError::PieceRejected { piece: piece_idx });
322 }
323 Ok(Ok(UtMetadataMsg::Request(_))) => {
324 return Err(MetadataExchangeError::BencodeDecodeFailed {
325 detail: "Unexpected Request message type from peer".to_string(),
326 });
327 }
328 Ok(Err(inner_err)) => {
329 return Err(MetadataExchangeError::BencodeDecodeFailed {
330 detail: format!("ut_metadata error for piece {}: {}", piece_idx, inner_err),
331 });
332 }
333 Err(_) => {
334 return Err(MetadataExchangeError::PieceTimeout { piece: piece_idx });
335 }
336 }
337 }
338
339 collector.assemble().ok_or_else(|| {
340 let received = (collector.progress() * metadata_size as f64) as u64;
341 MetadataExchangeError::IncompleteMetadata {
342 expected: metadata_size,
343 received,
344 }
345 })
346 }
347
348 async fn read_extension_message(
349 &self,
350 conn: &mut PeerConnection,
351 ) -> Result<BencodeValue, MetadataExchangeError> {
352 let mut len_buf = [0u8; 4];
353 conn.stream_read_exact(&mut len_buf).await.map_err(|e| {
354 MetadataExchangeError::IoError(format!("Read message length failed: {}", e))
355 })?;
356
357 let msg_len = u32::from_be_bytes(len_buf) as usize;
358 if msg_len == 0 || msg_len > 10 * 1024 * 1024 {
359 return Err(MetadataExchangeError::BencodeDecodeFailed {
360 detail: "Invalid extension message length".to_string(),
361 });
362 }
363
364 let mut payload = vec![0u8; msg_len];
365 conn.stream_read_exact(&mut payload).await.map_err(|e| {
366 MetadataExchangeError::IoError(format!("Read message body failed: {}", e))
367 })?;
368
369 if payload.first().copied() != Some(20u8) {
370 return Err(MetadataExchangeError::BencodeDecodeFailed {
371 detail: "Expected extended message ID 20".to_string(),
372 });
373 }
374
375 BencodeValue::decode(&payload[1..])
376 .map(|(v, _)| v)
377 .map_err(|e| MetadataExchangeError::BencodeDecodeFailed {
378 detail: format!("Decode BEncode failed: {}", e),
379 })
380 }
381
382 async fn read_ut_metadata_response(
383 &self,
384 conn: &mut PeerConnection,
385 ) -> Result<UtMetadataMsg, MetadataExchangeError> {
386 let mut len_buf = [0u8; 4];
387 conn.stream_read_exact(&mut len_buf).await.map_err(|e| {
388 MetadataExchangeError::IoError(format!("Read ut_metadata length failed: {}", e))
389 })?;
390
391 let msg_len = u32::from_be_bytes(len_buf) as usize;
392 if msg_len == 0 || msg_len > 10 * 1024 * 1024 {
393 return Err(MetadataExchangeError::BencodeDecodeFailed {
394 detail: "Invalid ut_metadata message length".to_string(),
395 });
396 }
397
398 let mut payload = vec![0u8; msg_len];
399 conn.stream_read_exact(&mut payload).await.map_err(|e| {
400 MetadataExchangeError::IoError(format!("Read ut_metadata body failed: {}", e))
401 })?;
402
403 UtMetadataMsg::decode(&payload).map_err(|e| MetadataExchangeError::BencodeDecodeFailed {
404 detail: format!("ut_metadata decode failed: {}", e),
405 })
406 }
407}
408
409#[cfg(test)]
410mod tests {
411 use super::*;
412
413 #[test]
414 fn test_config_default() {
415 let cfg = MetadataExchangeConfig::default();
416 assert_eq!(cfg.max_peers_to_try, 5);
417 assert_eq!(cfg.piece_size, 16 * 1024);
418 assert_eq!(cfg.max_attempts, DEFAULT_MAX_ATTEMPTS);
419 }
420
421 #[test]
422 fn test_fetch_metadata_no_peers() {
423 let session = MetadataExchangeSession::new(MetadataExchangeConfig::default());
424 let target_hash = [0u8; 20];
425 let peers: Vec<SocketAddr> = vec![];
426
427 let rt = tokio::runtime::Runtime::new().unwrap();
428 let result = rt.block_on(session.fetch_metadata(&target_hash, &peers));
429
430 assert!(result.is_err());
431 match result.unwrap_err() {
432 MetadataExchangeError::NoPeersAvailable => {}
433 other => panic!("Expected NoPeersAvailable, got {:?}", other),
434 }
435 }
436
437 #[tokio::test]
438 async fn test_read_extension_message_invalid_length_zero() {
439 let session = MetadataExchangeSession::new(MetadataExchangeConfig::default());
440 let _ = session;
441 }
442
443 #[test]
444 fn test_error_enum_variant_count() {
445 let _ = MetadataExchangeError::NoPeersAvailable;
446 let _ = MetadataExchangeError::AllPeersFailed {
447 attempts: 0,
448 last_error: String::new(),
449 };
450 let _ = MetadataExchangeError::PeerConnectFailed {
451 addr: String::new(),
452 reason: String::new(),
453 };
454 let _ = MetadataExchangeError::PeerTimeout {
455 addr: String::new(),
456 };
457 let _ = MetadataExchangeError::UnsupportedPeer {
458 addr: String::new(),
459 reason: String::new(),
460 };
461 let _ = MetadataExchangeError::InvalidMetadataSize { size: 0 };
462 let _ = MetadataExchangeError::MetadataTooLarge { size: 0, max: 0 };
463 let _ = MetadataExchangeError::BencodeDecodeFailed {
464 detail: String::new(),
465 };
466 let _ = MetadataExchangeError::PieceRejected { piece: 0 };
467 let _ = MetadataExchangeError::PieceTimeout { piece: 0 };
468 let _ = MetadataExchangeError::IncompleteMetadata {
469 expected: 0,
470 received: 0,
471 };
472 let _ = MetadataExchangeError::IoError(String::new());
473 }
474
475 #[test]
476 fn test_fatal_vs_recoverable_errors() {
477 let fatal = MetadataExchangeError::MetadataTooLarge {
478 size: 200_000_000,
479 max: METADATA_MAX_SIZE,
480 };
481 assert!(fatal.is_fatal());
482
483 let fatal2 = MetadataExchangeError::BencodeDecodeFailed {
484 detail: "bad".to_string(),
485 };
486 assert!(fatal2.is_fatal());
487
488 let recoverable = MetadataExchangeError::PeerTimeout {
489 addr: "1.2.3.4:6881".to_string(),
490 };
491 assert!(!recoverable.is_fatal());
492
493 let recoverable2 = MetadataExchangeError::InvalidMetadataSize { size: 0 };
494 assert!(!recoverable2.is_fatal());
495 }
496
497 #[test]
498 fn test_display_impl() {
499 let err = MetadataExchangeError::NoPeersAvailable;
500 assert!(err.to_string().contains("No peers"));
501
502 let err = MetadataExchangeError::MetadataTooLarge {
503 size: 200_000_000,
504 max: METADATA_MAX_SIZE,
505 };
506 let display = err.to_string();
507 assert!(display.contains("too large"));
508 assert!(display.contains("200000000"));
509 }
510
511 #[test]
512 fn test_with_piece_size_builder() {
513 let cfg = MetadataExchangeConfig::default().with_piece_size(8192);
514 assert_eq!(cfg.piece_size, 8192);
515
516 let cfg_clamped_low = MetadataExchangeConfig::default().with_piece_size(512);
517 assert_eq!(cfg_clamped_low.piece_size, PIECE_SIZE_MIN);
518
519 let cfg_clamped_high = MetadataExchangeConfig::default().with_piece_size(128_000);
520 assert_eq!(cfg_clamped_high.piece_size, PIECE_SIZE_MAX);
521 }
522
523 #[test]
524 fn test_with_max_attempts_builder() {
525 let cfg = MetadataExchangeConfig::default().with_max_attempts(5);
526 assert_eq!(cfg.max_attempts, 5);
527
528 let cfg_zero = MetadataExchangeConfig::default().with_max_attempts(0);
529 assert_eq!(cfg_zero.max_attempts, 1);
530 }
531}