1use std::path::{Path, PathBuf};
2
3use serde::{Deserialize, Serialize};
4use tracing::{debug, info};
5
6use aria2_protocol::bittorrent::torrent::parser::InfoDict;
7
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct TorrentFileEntry {
10 pub index: usize,
11 pub path: String,
12 pub length: u64,
13 pub completed_length: u64,
14}
15
16#[derive(Debug, Clone)]
17pub struct FileInfo {
18 pub path: Vec<String>,
19 pub length: u64,
20 pub start_piece: u32,
21 pub end_piece: u32,
22 pub start_offset_in_piece: u32,
23 pub end_offset_in_piece: u32,
24 pub absolute_path: PathBuf,
25}
26
27#[derive(Clone)]
28pub struct MultiFileLayout {
29 #[allow(dead_code)] base_dir: PathBuf,
31 files: Vec<FileInfo>,
32 piece_length: u32,
33 total_pieces: u32,
34 total_size: u64,
35 is_single_file: bool,
36}
37
38impl MultiFileLayout {
39 pub fn from_info_dict(info: &InfoDict, base_dir: &Path) -> Result<Self, String> {
40 let piece_length = info.piece_length;
41 let total_pieces = info.pieces.len() as u32;
42
43 if let Some(length) = info.length {
44 let name = &info.name;
45 let absolute_path = base_dir.join(name);
46
47 let total_size = length;
48
49 let start_piece = 0u32;
50 let end_piece = if total_size == 0 {
51 0
52 } else if total_pieces > 0 {
53 total_pieces - 1
54 } else {
55 0
56 };
57 let start_offset_in_piece = 0u32;
58 let end_offset_in_piece = if total_size == 0 {
59 0
60 } else {
61 ((total_size - 1) % piece_length as u64 + 1) as u32
62 };
63
64 let file_info = FileInfo {
65 path: vec![name.clone()],
66 length,
67 start_piece,
68 end_piece,
69 start_offset_in_piece,
70 end_offset_in_piece,
71 absolute_path,
72 };
73
74 info!(
75 "Single-file layout: name={}, length={}, pieces={}",
76 name, length, total_pieces
77 );
78
79 Ok(Self {
80 base_dir: base_dir.to_path_buf(),
81 files: vec![file_info],
82 piece_length,
83 total_pieces,
84 total_size,
85 is_single_file: true,
86 })
87 } else if let Some(ref files) = info.files {
88 if files.is_empty() {
89 return Err("files list is empty".to_string());
90 }
91
92 let mut file_infos = Vec::with_capacity(files.len());
93 let mut running_offset: u64 = 0;
94 let mut computed_total_size: u64 = 0;
95
96 for (i, entry) in files.iter().enumerate() {
97 let start_byte = running_offset;
98 let end_byte = start_byte + entry.length;
99
100 let pl = piece_length as u64;
101
102 let start_piece = (start_byte / pl) as u32;
103 let start_offset_in_piece = (start_byte % pl) as u32;
104
105 let end_piece = if entry.length == 0 && start_byte == 0 || end_byte == 0 {
106 0
107 } else {
108 ((end_byte - 1) / pl) as u32
109 };
110 let end_offset_in_piece = if entry.length == 0 {
111 0
112 } else {
113 ((end_byte - 1) % pl + 1) as u32
114 };
115
116 let mut path_buf = base_dir.to_path_buf();
118 for component in &entry.path {
119 path_buf.push(component);
120 }
121 let abs_path = path_buf;
122 debug!(
123 "File[{}]: path={:?}, bytes=[{}..{}), pieces=[{}..{}] offsets=[{}..{})",
124 i,
125 entry.path,
126 start_byte,
127 end_byte,
128 start_piece,
129 end_piece,
130 start_offset_in_piece,
131 end_offset_in_piece
132 );
133
134 file_infos.push(FileInfo {
135 path: entry.path.clone(),
136 length: entry.length,
137 start_piece,
138 end_piece,
139 start_offset_in_piece,
140 end_offset_in_piece,
141 absolute_path: abs_path,
142 });
143
144 running_offset = end_byte;
145 computed_total_size += entry.length;
146 }
147
148 info!(
149 "Multi-file layout: {} files, total_size={}, pieces={}",
150 file_infos.len(),
151 computed_total_size,
152 total_pieces
153 );
154
155 Ok(Self {
156 base_dir: base_dir.to_path_buf(),
157 files: file_infos,
158 piece_length,
159 total_pieces,
160 total_size: computed_total_size,
161 is_single_file: false,
162 })
163 } else {
164 Err("InfoDict has neither length nor files field".to_string())
165 }
166 }
167
168 pub fn create_directories(&self) -> Result<(), String> {
169 for (i, file) in self.files.iter().enumerate() {
170 if let Some(parent) = file.absolute_path.parent() {
171 std::fs::create_dir_all(parent).map_err(|e| {
172 format!(
173 "Failed to create directory {:?} for file[{}] {:?}: {}",
174 parent, i, file.path, e
175 )
176 })?;
177 debug!("Created directory: {:?}", parent);
178 }
179 }
180 Ok(())
181 }
182
183 pub fn resolve_file_offset(
184 &self,
185 piece_idx: u32,
186 offset_in_piece: u32,
187 ) -> Option<(usize, u64)> {
188 let global_byte = piece_idx as u64 * self.piece_length as u64 + offset_in_piece as u64;
189
190 if global_byte >= self.total_size {
191 return None;
192 }
193
194 for (i, file) in self.files.iter().enumerate() {
195 let file_start = file.start_piece as u64 * self.piece_length as u64
196 + file.start_offset_in_piece as u64;
197 let file_end = file_start + file.length;
198
199 if global_byte >= file_start && global_byte < file_end {
200 return Some((i, global_byte - file_start));
201 }
202 }
203
204 None
205 }
206
207 pub fn file_absolute_path(&self, file_index: usize) -> Option<&Path> {
208 self.files
209 .get(file_index)
210 .map(|f| f.absolute_path.as_path())
211 }
212
213 pub fn file_completed_bytes(&self, file_idx: usize, bitfield: &[u8]) -> u64 {
214 let file = match self.files.get(file_idx) {
215 Some(f) => f,
216 None => return 0,
217 };
218
219 if file.length == 0 {
220 return 0;
221 }
222
223 let pl = self.piece_length as u64;
224 let mut completed: u64 = 0;
225
226 for piece_idx in file.start_piece..=file.end_piece {
227 let byte_index = piece_idx as usize / 8;
228 let bit_index = 7 - (piece_idx as usize % 8);
229
230 if byte_index >= bitfield.len() {
231 break;
232 }
233
234 let is_complete = (bitfield[byte_index] >> bit_index) & 1 == 1;
235
236 if !is_complete {
237 continue;
238 }
239
240 if piece_idx == file.start_piece && piece_idx == file.end_piece {
241 completed += file.length;
242 } else if piece_idx == file.start_piece {
243 let bytes_in_this_piece = pl - file.start_offset_in_piece as u64;
244 completed += bytes_in_this_piece.min(file.length);
245 } else if piece_idx == file.end_piece {
246 completed += file.end_offset_in_piece as u64;
247 } else {
248 completed += pl;
249 }
250 }
251
252 completed.min(file.length)
253 }
254
255 pub fn file_list(&self) -> Vec<TorrentFileEntry> {
256 self.files
257 .iter()
258 .enumerate()
259 .map(|(i, f)| TorrentFileEntry {
260 index: i,
261 path: f.path.join("/"),
262 length: f.length,
263 completed_length: 0,
264 })
265 .collect()
266 }
267
268 pub fn is_multi_file(&self) -> bool {
269 !self.is_single_file
270 }
271
272 pub fn num_files(&self) -> usize {
273 self.files.len()
274 }
275
276 pub fn total_size(&self) -> u64 {
277 self.total_size
278 }
279
280 pub fn piece_length(&self) -> u32 {
281 self.piece_length
282 }
283
284 pub fn total_pieces(&self) -> u32 {
285 self.total_pieces
286 }
287
288 pub fn get_file_info(&self, index: usize) -> Option<&FileInfo> {
289 self.files.get(index)
290 }
291}
292
293#[cfg(test)]
294mod tests {
295 use super::*;
296 use aria2_protocol::bittorrent::torrent::parser::FileEntry;
297
298 fn make_single_file_info_dict() -> InfoDict {
299 InfoDict {
300 name: "single_file.bin".to_string(),
301 piece_length: 512,
302 pieces: vec![[0u8; 20], [1u8; 20]],
303 length: Some(1024),
304 files: None,
305 private: None,
306 }
307 }
308
309 fn make_multi_file_info_dict() -> InfoDict {
310 InfoDict {
311 name: "multi_dir".to_string(),
312 piece_length: 512,
313 pieces: vec![[0u8; 20], [1u8; 20], [2u8; 20]],
314 length: None,
315 files: Some(vec![
316 FileEntry {
317 length: 500,
318 path: vec!["dir1".to_string(), "file1.txt".to_string()],
319 },
320 FileEntry {
321 length: 524,
322 path: vec!["dir2".to_string(), "file2.dat".to_string()],
323 },
324 FileEntry {
325 length: 300,
326 path: vec!["dir3".to_string(), "file3.log".to_string()],
327 },
328 ]),
329 private: None,
330 }
331 }
332
333 #[test]
334 fn test_from_info_dict_single_file() {
335 let info = make_single_file_info_dict();
336 let base = Path::new("/tmp/download");
337 let layout = MultiFileLayout::from_info_dict(&info, base).unwrap();
338
339 assert_eq!(layout.num_files(), 1);
340 assert!(!layout.is_multi_file());
341 assert_eq!(layout.total_size(), 1024);
342 assert_eq!(layout.piece_length(), 512);
343 assert_eq!(layout.total_pieces(), 2);
344
345 let file = layout.get_file_info(0).unwrap();
346 assert_eq!(file.path, vec!["single_file.bin"]);
347 assert_eq!(file.length, 1024);
348 assert_eq!(file.start_piece, 0);
349 assert_eq!(file.end_piece, 1);
350 assert_eq!(file.start_offset_in_piece, 0);
351 assert_eq!(file.end_offset_in_piece, 512);
352 }
353
354 #[test]
355 fn test_from_info_dict_multi_file() {
356 let info = make_multi_file_info_dict();
357 let base = Path::new("/tmp/download");
358 let layout = MultiFileLayout::from_info_dict(&info, base).unwrap();
359
360 assert_eq!(layout.num_files(), 3);
361 assert!(layout.is_multi_file());
362 assert_eq!(layout.total_size(), 1324);
363
364 let f0 = layout.get_file_info(0).unwrap();
365 assert_eq!(f0.length, 500);
366 assert_eq!(f0.start_piece, 0);
367 assert_eq!(f0.end_piece, 0);
368 assert_eq!(f0.start_offset_in_piece, 0);
369 assert_eq!(f0.end_offset_in_piece, 500);
370
371 let f1 = layout.get_file_info(1).unwrap();
372 assert_eq!(f1.length, 524);
373 assert_eq!(f1.start_piece, 0);
374 assert_eq!(f1.end_piece, 1);
375 assert_eq!(f1.start_offset_in_piece, 500);
376 assert_eq!(f1.end_offset_in_piece, 512);
377
378 let f2 = layout.get_file_info(2).unwrap();
379 assert_eq!(f2.length, 300);
380 assert_eq!(f2.start_piece, 2);
381 assert_eq!(f2.end_piece, 2);
382 assert_eq!(f2.start_offset_in_piece, 0);
383 assert_eq!(f2.end_offset_in_piece, 300);
384 }
385
386 #[test]
387 fn test_from_info_dict_empty_files() {
388 let info = InfoDict {
389 name: "empty".to_string(),
390 piece_length: 512,
391 pieces: vec![],
392 length: None,
393 files: Some(vec![]),
394 private: None,
395 };
396 let base = Path::new("/tmp/download");
397 let result = MultiFileLayout::from_info_dict(&info, base);
398 assert!(result.is_err());
399 }
400
401 #[test]
402 fn test_create_directories() {
403 let info = make_multi_file_info_dict();
404 let temp_dir = tempfile::tempdir().unwrap();
405 let base = temp_dir.path();
406 let layout = MultiFileLayout::from_info_dict(&info, base).unwrap();
407
408 let result = layout.create_directories();
409 assert!(result.is_ok());
410
411 let dir1 = base.join("dir1");
412 let dir2 = base.join("dir2");
413 let dir3 = base.join("dir3");
414
415 assert!(dir1.exists());
416 assert!(dir2.exists());
417 assert!(dir3.exists());
418
419 }
421
422 #[test]
423 fn test_resolve_file_offset_single_file() {
424 let info = make_single_file_info_dict();
425 let base = Path::new("/tmp/download");
426 let layout = MultiFileLayout::from_info_dict(&info, base).unwrap();
427
428 let result = layout.resolve_file_offset(0, 0);
429 assert_eq!(result, Some((0, 0)));
430
431 let result = layout.resolve_file_offset(0, 256);
432 assert_eq!(result, Some((0, 256)));
433
434 let result = layout.resolve_file_offset(1, 0);
435 assert_eq!(result, Some((0, 512)));
436
437 let result = layout.resolve_file_offset(1, 511);
438 assert_eq!(result, Some((0, 1023)));
439 }
440
441 #[test]
442 fn test_resolve_file_offset_multi_file() {
443 let info = make_multi_file_info_dict();
444 let base = Path::new("/tmp/download");
445 let layout = MultiFileLayout::from_info_dict(&info, base).unwrap();
446
447 let result = layout.resolve_file_offset(0, 0);
448 assert_eq!(result, Some((0, 0)));
449
450 let result = layout.resolve_file_offset(0, 499);
451 assert_eq!(result, Some((0, 499)));
452
453 let result = layout.resolve_file_offset(0, 500);
454 assert_eq!(result, Some((1, 0)));
455
456 let result = layout.resolve_file_offset(1, 0);
457 assert_eq!(result, Some((1, 12)));
458
459 let result = layout.resolve_file_offset(1, 200);
460 assert_eq!(result, Some((1, 212)));
461
462 let result = layout.resolve_file_offset(2, 50);
463 assert_eq!(result, Some((2, 50)));
464 }
465
466 #[test]
467 fn test_resolve_file_offset_boundary() {
468 let info = make_multi_file_info_dict();
469 let base = Path::new("/tmp/download");
470 let layout = MultiFileLayout::from_info_dict(&info, base).unwrap();
471
472 let result = layout.resolve_file_offset(0, 499);
473 assert_eq!(result, Some((0, 499)));
474
475 let result = layout.resolve_file_offset(0, 500);
476 assert_eq!(result, Some((1, 0)));
477
478 let result = layout.resolve_file_offset(1, 511);
479 assert_eq!(result, Some((1, 523)));
480
481 let result = layout.resolve_file_offset(2, 0);
482 assert_eq!(result, Some((2, 0)));
483 }
484
485 #[test]
486 fn test_resolve_file_offset_out_of_range() {
487 let info = make_single_file_info_dict();
488 let base = Path::new("/tmp/download");
489 let layout = MultiFileLayout::from_info_dict(&info, base).unwrap();
490
491 let result = layout.resolve_file_offset(1, 512);
492 assert_eq!(result, None);
493
494 let result = layout.resolve_file_offset(2, 0);
495 assert_eq!(result, None);
496
497 let result = layout.resolve_file_offset(u32::MAX, 0);
498 assert_eq!(result, None);
499 }
500
501 #[test]
502 fn test_file_completed_bytes() {
503 let info = make_single_file_info_dict();
504 let base = Path::new("/tmp/download");
505 let layout = MultiFileLayout::from_info_dict(&info, base).unwrap();
506
507 let no_pieces = [0u8; 1];
508 assert_eq!(layout.file_completed_bytes(0, &no_pieces), 0);
509
510 let piece0_only = [0b10000000u8];
511 assert_eq!(layout.file_completed_bytes(0, &piece0_only), 512);
512
513 let both_pieces = [0b11000000u8];
514 assert_eq!(layout.file_completed_bytes(0, &both_pieces), 1024);
515 }
516
517 #[test]
518 fn test_file_list_returns_correct_entries() {
519 let info = make_multi_file_info_dict();
520 let base = Path::new("/tmp/download");
521 let layout = MultiFileLayout::from_info_dict(&info, base).unwrap();
522
523 let list = layout.file_list();
524 assert_eq!(list.len(), 3);
525
526 assert_eq!(list[0].index, 0);
527 assert_eq!(list[0].path, "dir1/file1.txt");
528 assert_eq!(list[0].length, 500);
529 assert_eq!(list[0].completed_length, 0);
530
531 assert_eq!(list[1].index, 1);
532 assert_eq!(list[1].path, "dir2/file2.dat");
533 assert_eq!(list[1].length, 524);
534
535 assert_eq!(list[2].index, 2);
536 assert_eq!(list[2].path, "dir3/file3.log");
537 assert_eq!(list[2].length, 300);
538 }
539
540 #[test]
541 fn test_is_multi_file_flags() {
542 let single = make_single_file_info_dict();
543 let single_layout = MultiFileLayout::from_info_dict(&single, Path::new("/tmp")).unwrap();
544 assert!(!single_layout.is_multi_file());
545
546 let multi = make_multi_file_info_dict();
547 let multi_layout = MultiFileLayout::from_info_dict(&multi, Path::new("/tmp")).unwrap();
548 assert!(multi_layout.is_multi_file());
549 }
550
551 #[test]
552 fn test_total_size_matches() {
553 let single = make_single_file_info_dict();
554 let single_layout = MultiFileLayout::from_info_dict(&single, Path::new("/tmp")).unwrap();
555 assert_eq!(single_layout.total_size(), 1024);
556
557 let multi = make_multi_file_info_dict();
558 let multi_layout = MultiFileLayout::from_info_dict(&multi, Path::new("/tmp")).unwrap();
559 assert_eq!(multi_layout.total_size(), 1324);
560 assert_eq!(multi_layout.total_size(), 500 + 524 + 300);
561 }
562
563 #[test]
564 fn test_file_absolute_path() {
565 let info = make_multi_file_info_dict();
566 let base = Path::new("/base/dir");
567 let layout = MultiFileLayout::from_info_dict(&info, base).unwrap();
568
569 let p0 = layout.file_absolute_path(0).unwrap();
570 assert_eq!(p0, Path::new("/base/dir/dir1/file1.txt"));
571
572 let p1 = layout.file_absolute_path(1).unwrap();
573 assert_eq!(p1, Path::new("/base/dir/dir2/file2.dat"));
574
575 let p2 = layout.file_absolute_path(2).unwrap();
576 assert_eq!(p2, Path::new("/base/dir/dir3/file3.log"));
577
578 assert!(layout.file_absolute_path(3).is_none());
579 }
580
581 #[test]
582 fn test_file_completed_bytes_multi_file_partial() {
583 let info = make_multi_file_info_dict();
584 let base = Path::new("/tmp/download");
585 let layout = MultiFileLayout::from_info_dict(&info, base).unwrap();
586
587 let only_piece0 = [0b10000000u8];
588 assert_eq!(layout.file_completed_bytes(0, &only_piece0), 500);
589 assert_eq!(layout.file_completed_bytes(1, &only_piece0), 12);
590 assert_eq!(layout.file_completed_bytes(2, &only_piece0), 0);
591
592 let piece0_and_1 = [0b11000000u8];
593 assert_eq!(layout.file_completed_bytes(0, &piece0_and_1), 500);
594 assert_eq!(layout.file_completed_bytes(1, &piece0_and_1), 524);
595 assert_eq!(layout.file_completed_bytes(2, &piece0_and_1), 0);
596
597 let all_pieces = [0b11100000u8];
598 assert_eq!(layout.file_completed_bytes(0, &all_pieces), 500);
599 assert_eq!(layout.file_completed_bytes(1, &all_pieces), 524);
600 assert_eq!(layout.file_completed_bytes(2, &all_pieces), 300);
601 }
602}