aria2_core/filesystem/
control_file.rs

1use crate::error::{Aria2Error, Result};
2use std::path::{Path, PathBuf};
3
4const CONTROL_MAGIC: &[u8; 4] = b"A2CF";
5const CONTROL_VERSION: u16 = 1;
6const FLAG_HAS_CHECKSUM: u8 = 0x01;
7
8#[derive(Debug, Clone)]
9pub struct ControlFile {
10    path: PathBuf,
11    total_length: u64,
12    completed_length: u64,
13    upload_length: u64,
14    bitfield: Vec<u8>,
15    num_pieces: usize,
16    checksum_algo: u8,
17    checksum_value: Vec<u8>,
18}
19
20impl ControlFile {
21    pub fn path(&self) -> &Path {
22        &self.path
23    }
24    pub fn total_length(&self) -> u64 {
25        self.total_length
26    }
27    pub fn completed_length(&self) -> u64 {
28        self.completed_length
29    }
30    pub fn bitfield(&self) -> &[u8] {
31        &self.bitfield
32    }
33    pub fn set_checksum(&mut self, algo: u8, value: Vec<u8>) {
34        self.checksum_algo = algo;
35        self.checksum_value = value;
36    }
37    pub fn checksum_algo(&self) -> u8 {
38        self.checksum_algo
39    }
40
41    pub async fn open_or_create(
42        ctrl_path: &Path,
43        total_length: u64,
44        num_pieces: usize,
45    ) -> Result<Self> {
46        if ctrl_path.exists() {
47            Self::load(ctrl_path)
48                .await?
49                .ok_or_else(|| Aria2Error::Io(format!("无法加载控制文件: {}", ctrl_path.display())))
50        } else {
51            let bitfield_len = num_pieces.div_ceil(8);
52            Ok(Self {
53                path: ctrl_path.to_path_buf(),
54                total_length,
55                completed_length: 0,
56                upload_length: 0,
57                bitfield: vec![0u8; bitfield_len],
58                num_pieces,
59                checksum_algo: 0,
60                checksum_value: Vec::new(),
61            })
62        }
63    }
64
65    pub async fn load(path: &Path) -> Result<Option<Self>> {
66        if !path.exists() {
67            return Ok(None);
68        }
69
70        let data = tokio::fs::read(path)
71            .await
72            .map_err(|e| Aria2Error::Io(e.to_string()))?;
73
74        if data.len() < 8 {
75            return Ok(None);
76        }
77
78        if &data[0..4] != CONTROL_MAGIC {
79            return Err(Aria2Error::Io("无效的控制文件magic".to_string()));
80        }
81
82        let version = u16_from_le(&data[4..6]);
83        if version > CONTROL_VERSION {
84            return Err(Aria2Error::Io(format!("不支持的版本: {}", version)));
85        }
86
87        let flags = data[6];
88        let total_length = u64_from_le(&data[7..15]);
89        let completed_length = u64_from_le(&data[15..23]);
90        let upload_length = u64_from_le(&data[23..31]);
91        let _bitfield_length = u64_from_le(&data[31..39]);
92
93        let mut offset = 39usize;
94        let checksum_algo = if flags & FLAG_HAS_CHECKSUM != 0 {
95            let algo = data[offset];
96            offset += 1;
97            algo
98        } else {
99            0
100        };
101
102        let checksum_value = if flags & FLAG_HAS_CHECKSUM != 0 && checksum_algo > 0 {
103            let len = match checksum_algo {
104                1 => 16,
105                2 => 20,
106                3 => 32,
107                4 => 8,
108                _ => 0,
109            };
110            if len > 0 && offset + len <= data.len() {
111                let val = data[offset..offset + len].to_vec();
112                offset += len;
113                val
114            } else {
115                Vec::new()
116            }
117        } else {
118            Vec::new()
119        };
120
121        let bitfield = data[offset..].to_vec();
122        let num_pieces = bitfield.len() * 8;
123
124        Ok(Some(Self {
125            path: path.to_path_buf(),
126            total_length,
127            completed_length,
128            upload_length,
129            bitfield,
130            num_pieces,
131            checksum_algo,
132            checksum_value,
133        }))
134    }
135
136    pub async fn save(&self) -> Result<()> {
137        let mut buf = Vec::with_capacity(64 + self.bitfield.len());
138
139        buf.extend_from_slice(CONTROL_MAGIC);
140        buf.extend_from_slice(&CONTROL_VERSION.to_le_bytes());
141        let mut flags: u8 = 0;
142        if self.checksum_algo > 0 && !self.checksum_value.is_empty() {
143            flags |= FLAG_HAS_CHECKSUM;
144        }
145        buf.push(flags);
146        buf.extend_from_slice(&self.total_length.to_le_bytes());
147        buf.extend_from_slice(&self.completed_length.to_le_bytes());
148        buf.extend_from_slice(&self.upload_length.to_le_bytes());
149        buf.extend_from_slice(&(self.bitfield.len() as u64).to_le_bytes());
150
151        if flags & FLAG_HAS_CHECKSUM != 0 {
152            buf.push(self.checksum_algo);
153            buf.extend_from_slice(&self.checksum_value);
154        }
155
156        buf.extend_from_slice(&self.bitfield);
157
158        let tmp_path = self.path.with_extension("aria2.tmp");
159        {
160            tokio::fs::write(&tmp_path, &buf)
161                .await
162                .map_err(|e| Aria2Error::Io(e.to_string()))?;
163            if let Ok(f) = tokio::fs::File::open(&tmp_path).await {
164                let _ = f.sync_all().await;
165            }
166        }
167        tokio::fs::rename(&tmp_path, &self.path)
168            .await
169            .map_err(|e| Aria2Error::Io(e.to_string()))?;
170        Ok(())
171    }
172
173    pub fn mark_piece_done(&mut self, index: usize) {
174        let byte_index = index / 8;
175        let bit_index = index % 8;
176        if byte_index < self.bitfield.len() {
177            self.bitfield[byte_index] |= 1 << (7 - bit_index);
178            self.completed_length = self.calculate_completed();
179        }
180    }
181
182    pub fn is_piece_done(&self, index: usize) -> bool {
183        let byte_index = index / 8;
184        let bit_index = index % 8;
185        if byte_index < self.bitfield.len() {
186            (self.bitfield[byte_index] & (1 << (7 - bit_index))) != 0
187        } else {
188            false
189        }
190    }
191
192    pub fn completed_pieces(&self) -> usize {
193        self.bitfield.iter().map(|b| b.count_ones() as usize).sum()
194    }
195
196    fn calculate_completed(&self) -> u64 {
197        let bits = self.completed_pieces() as u64;
198        if self.total_length == 0 || self.num_pieces == 0 {
199            return 0;
200        }
201        let piece_size = self.total_length / self.num_pieces as u64;
202        bits * piece_size
203    }
204
205    pub fn update_completed_length(&mut self, length: u64) {
206        self.completed_length = length.min(self.total_length);
207    }
208
209    pub fn control_path_for(output_path: &Path) -> PathBuf {
210        let mut p = output_path.to_path_buf();
211        p.set_extension(".aria2");
212        p
213    }
214}
215
216fn u16_from_le(b: &[u8]) -> u16 {
217    u16::from_le_bytes([b[0], b[1]])
218}
219
220fn u64_from_le(b: &[u8]) -> u64 {
221    u64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]])
222}
223
224#[cfg(test)]
225mod tests {
226    use super::*;
227
228    #[tokio::test]
229    async fn test_control_file_new_and_save() {
230        let dir = tempfile::tempdir().unwrap();
231        let path = dir.path().join("test.aria2");
232
233        let cf = ControlFile::open_or_create(&path, 10000, 10).await.unwrap();
234        assert_eq!(cf.total_length(), 10000);
235        assert_eq!(cf.completed_length(), 0);
236        assert!(!cf.is_piece_done(0));
237
238        cf.save().await.unwrap();
239
240        assert!(path.exists());
241        let data = tokio::fs::read(&path).await.unwrap();
242        assert_eq!(&data[0..4], b"A2CF");
243    }
244
245    #[tokio::test]
246    async fn test_control_file_mark_and_check_pieces() {
247        let dir = tempfile::tempdir().unwrap();
248        let path = dir.path().join("test.aria2");
249
250        let mut cf = ControlFile::open_or_create(&path, 1000, 8).await.unwrap();
251
252        cf.mark_piece_done(0);
253        cf.mark_piece_done(3);
254        cf.mark_piece_done(7);
255
256        assert!(cf.is_piece_done(0));
257        assert!(!cf.is_piece_done(1));
258        assert!(cf.is_piece_done(3));
259        assert!(!cf.is_piece_done(5));
260        assert!(cf.is_piece_done(7));
261        assert_eq!(cf.completed_pieces(), 3);
262
263        cf.save().await.unwrap();
264
265        let loaded = ControlFile::load(&path).await.unwrap().unwrap();
266        assert_eq!(loaded.completed_pieces(), 3);
267        assert!(loaded.is_piece_done(0));
268        assert!(loaded.is_piece_done(7));
269    }
270
271    #[tokio::test]
272    async fn test_control_file_roundtrip_with_checksum() {
273        let dir = tempfile::tempdir().unwrap();
274        let path = dir.path().join("test_hash.aria2");
275
276        let mut cf = ControlFile::open_or_create(&path, 5000, 5).await.unwrap();
277        cf.checksum_algo = 2;
278        cf.checksum_value = vec![0xAB; 20];
279        cf.mark_piece_done(0);
280        cf.mark_piece_done(2);
281        cf.save().await.unwrap();
282
283        let loaded = ControlFile::load(&path).await.unwrap().unwrap();
284        assert_eq!(loaded.total_length(), 5000);
285        assert_eq!(loaded.checksum_algo, 2);
286        assert_eq!(loaded.completed_pieces(), 2);
287    }
288
289    #[tokio::test]
290    async fn test_control_file_atomic_save() {
291        let dir = tempfile::tempdir().unwrap();
292        let path = dir.path().join("test_atomic.aria2");
293
294        let mut cf = ControlFile::open_or_create(&path, 999, 4).await.unwrap();
295        cf.mark_piece_done(1);
296        cf.save().await.unwrap();
297
298        let tmp_path = path.with_extension("aria2.tmp");
299        assert!(!tmp_path.exists());
300        assert!(path.exists());
301    }
302
303    #[tokio::test]
304    async fn test_control_file_load_nonexistent() {
305        let dir = tempfile::tempdir().unwrap();
306        let path = dir.path().join("nonexistent.aria2");
307        let result = ControlFile::load(&path).await.unwrap();
308        assert!(result.is_none());
309    }
310
311    #[tokio::test]
312    async fn test_control_file_load_invalid_magic() {
313        let dir = tempfile::tempdir().unwrap();
314        let path = dir.path().join("bad.aria2");
315        tokio::fs::write(&path, b"NOT_A2CF_DATA").await.unwrap();
316
317        let result = ControlFile::load(&path).await;
318        assert!(result.is_err());
319    }
320
321    #[tokio::test]
322    async fn test_control_path_for_output() {
323        let out = Path::new("/downloads/file.iso");
324        let ctrl = ControlFile::control_path_for(out);
325        assert_eq!(ctrl.extension().unwrap().to_str().unwrap(), "aria2");
326        assert!(ctrl.to_str().unwrap().ends_with(".aria2"));
327    }
328
329    #[tokio::test]
330    async fn test_control_file_update_completed_length() {
331        let dir = tempfile::tempdir().unwrap();
332        let path = dir.path().join("test_len.aria2");
333
334        let mut cf = ControlFile::open_or_create(&path, 8000, 8).await.unwrap();
335        cf.update_completed_length(3500);
336        assert_eq!(cf.completed_length(), 3500);
337
338        cf.update_completed_length(9000);
339        assert_eq!(cf.completed_length(), 8000);
340    }
341}
aria2_core/filesystem/control_file.rs

aria2_core/filesystem/
control_file.rs
