Skip to main content

rust_mirror/
resume.rs

1//! Private persistent downloads. A completed tree is assembled only from verified bytes.
2use crate::{
3    selection::{Artifact, Plan},
4    sha256, Result, Source, ARTIFACT_LIMIT, ORIGIN,
5};
6use serde::{Deserialize, Serialize};
7use sha2::{Digest, Sha256};
8use std::{
9    collections::BTreeSet,
10    fs::{self, File},
11    io::{self, Read, Seek, SeekFrom, Write},
12    path::{Path, PathBuf},
13};
14
15#[derive(Debug, PartialEq, Eq)]
16pub enum Transfer {
17    Complete,
18    Restart,
19}
20
21pub(crate) fn http_range(
22    agent: &ureq::Agent,
23    url: &str,
24    out: &mut dyn Write,
25    offset: u64,
26    limit: u64,
27) -> Result<Transfer> {
28    let mut request = agent.get(url).set("Accept-Encoding", "identity");
29    if offset > 0 {
30        request = request.set("Range", &format!("bytes={offset}-"));
31    }
32    let response = match request.call() {
33        Ok(r) => r,
34        Err(ureq::Error::Status(416, _)) if offset > 0 => return Ok(Transfer::Restart),
35        Err(e) => return Err(e.into()),
36    };
37    if offset > 0 && response.status() == 200 {
38        return Ok(Transfer::Restart);
39    }
40    if response
41        .header("Content-Encoding")
42        .is_some_and(|v| v != "identity")
43    {
44        return Err("encoded response cannot be resumed safely".into());
45    }
46    let expected = if offset > 0 {
47        if response.status() != 206 {
48            return Err("resume expected HTTP 206".into());
49        }
50        let header = response
51            .header("Content-Range")
52            .ok_or("missing Content-Range")?;
53        let (span, total) = header
54            .strip_prefix("bytes ")
55            .ok_or("invalid Content-Range unit")?
56            .split_once('/')
57            .ok_or("invalid Content-Range")?;
58        let (start, end) = span.split_once('-').ok_or("invalid Content-Range span")?;
59        let (start, end, total) = (number(start)?, number(end)?, number(total)?);
60        if start != offset
61            || end < start
62            || end.checked_add(1) != Some(total)
63            || total > ARTIFACT_LIMIT
64            || total - offset > limit
65        {
66            return Err("resume range exceeds expected offset or byte budget".into());
67        }
68        Some(total - offset)
69    } else {
70        if response.status() != 200 {
71            return Err("download expected HTTP 200".into());
72        }
73        None
74    };
75    let length = response.header("Content-Length").map(number).transpose()?;
76    if length.is_some_and(|n| n > limit) || expected.zip(length).is_some_and(|(a, b)| a != b) {
77        return Err("response length exceeds budget or disagrees with range".into());
78    }
79    let mut reader = response.into_reader();
80    let mut bytes = 0;
81    let mut buffer = [0; 64 * 1024];
82    loop {
83        let n = reader.read(&mut buffer)?;
84        if n == 0 {
85            break;
86        }
87        bytes += n as u64;
88        if bytes > limit || expected.is_some_and(|max| bytes > max) {
89            return Err("response exceeds download budget".into());
90        }
91        out.write_all(&buffer[..n])?;
92    }
93    if expected.or(length).is_some_and(|n| n != bytes) {
94        return Err("response ended before expected length".into());
95    }
96    Ok(Transfer::Complete)
97}
98fn number(s: &str) -> Result<u64> {
99    if s.is_empty() || !s.bytes().all(|b| b.is_ascii_digit()) {
100        return Err("invalid HTTP byte count".into());
101    }
102    Ok(s.parse()?)
103}
104fn regular(path: &Path) -> Result<Option<u64>> {
105    match fs::symlink_metadata(path) {
106        Ok(m) if m.is_file() && !m.file_type().is_symlink() => Ok(Some(m.len())),
107        Ok(_) => Err("download state must contain regular files, not links/devices".into()),
108        Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(None),
109        Err(e) => Err(e.into()),
110    }
111}
112fn directory(path: &Path) -> Result<()> {
113    let m = fs::symlink_metadata(path)?;
114    if !m.is_dir() || m.file_type().is_symlink() {
115        return Err("download directory must not be a link".into());
116    }
117    Ok(())
118}
119pub(crate) fn hash_file(path: &Path, limit: u64) -> Result<(u64, String)> {
120    if regular(path)?.ok_or("missing download file")? > limit {
121        return Err("stored download exceeds byte budget".into());
122    }
123    let mut f = File::open(path)?;
124    let mut count = 0;
125    let mut hash = Sha256::new();
126    let mut buf = [0; 64 * 1024];
127    loop {
128        let n = f.read(&mut buf)?;
129        if n == 0 {
130            break;
131        }
132        count += n as u64;
133        if count > limit {
134            return Err("stored download grew beyond budget".into());
135        }
136        hash.update(&buf[..n]);
137    }
138    Ok((count, hex::encode(hash.finalize())))
139}
140/// Enumerate only expected native paths. Never reinterpret Unix backslashes as separators.
141pub(crate) fn validate_tree(root: &Path, files: &BTreeSet<PathBuf>) -> Result<()> {
142    directory(root)?;
143    let mut dirs = BTreeSet::new();
144    for file in files {
145        let mut p = file.parent();
146        while let Some(parent) = p {
147            if parent.as_os_str().is_empty() {
148                break;
149            }
150            dirs.insert(parent.to_path_buf());
151            p = parent.parent();
152        }
153    }
154    let mut stack = vec![root.to_path_buf()];
155    let mut visited = 0;
156    while let Some(dir) = stack.pop() {
157        for entry in fs::read_dir(dir)? {
158            let entry = entry?;
159            visited += 1;
160            if visited > files.len() + dirs.len() {
161                return Err("unexpected download entries".into());
162            }
163            let path = entry.path();
164            let rel = path.strip_prefix(root)?;
165            let m = fs::symlink_metadata(&path)?;
166            if m.file_type().is_symlink() {
167                return Err("download trees must not contain links".into());
168            }
169            if m.is_dir() && dirs.contains(rel) {
170                stack.push(path);
171            } else if !m.is_file() || !files.contains(rel) || m.len() > ARTIFACT_LIMIT {
172                return Err("unexpected or oversized download entry".into());
173            }
174        }
175    }
176    Ok(())
177}
178#[derive(Deserialize, Serialize, PartialEq, Eq)]
179#[serde(deny_unknown_fields)]
180struct State {
181    schema: u32,
182    plan_sha256: String,
183}
184pub(crate) struct Session {
185    pub root: PathBuf,
186    files: PathBuf,
187    allowed: BTreeSet<PathBuf>,
188    _lock: File,
189}
190impl Session {
191    pub fn open(output: &Path, plan: &Plan, resume: bool) -> Result<Self> {
192        if !resume {
193            fs::create_dir(output)?;
194        }
195        directory(output)?;
196        for e in fs::read_dir(output)? {
197            let e = e?;
198            if ![".fetch.lock", ".download", "mirror"]
199                .iter()
200                .any(|n| e.file_name() == *n)
201            {
202                return Err("output contains unrelated files".into());
203            }
204        }
205        let lock = output.join(".fetch.lock");
206        if regular(&lock)?.is_some_and(|n| n != 0) {
207            return Err("invalid download lock file".into());
208        }
209        let lock = File::options()
210            .read(true)
211            .write(true)
212            .create(true)
213            .truncate(false)
214            .open(lock)?;
215        lock.try_lock()
216            .map_err(|_| "another fetch owns this output, or filesystem locking is unavailable")?;
217        let root = output.join(".download");
218        let state = State {
219            schema: 1,
220            plan_sha256: sha256(toml::to_string_pretty(plan)?.as_bytes()),
221        };
222        if resume {
223            directory(&root)?;
224            let file = root.join("state.toml");
225            if regular(&file)?.ok_or("missing resume identity")? > 4096 {
226                return Err("oversized resume identity".into());
227            }
228            let saved: State = toml::from_str(&fs::read_to_string(file)?)?;
229            if saved != state {
230                return Err("output belongs to a different frozen plan".into());
231            }
232        } else {
233            fs::create_dir(&root)?;
234            let mut f = File::create_new(root.join("state.toml"))?;
235            f.write_all(toml::to_string(&state)?.as_bytes())?;
236            f.sync_all()?;
237            fs::create_dir(root.join("files"))?;
238        }
239        let files = root.join("files");
240        directory(&files)?;
241        for entry in fs::read_dir(&root)? {
242            if !["state.toml", "files", "tree"]
243                .iter()
244                .any(|n| entry.as_ref().is_ok_and(|e| e.file_name() == *n))
245            {
246                return Err("unexpected download workspace entry".into());
247            }
248        }
249        let allowed = plan
250            .artifacts
251            .iter()
252            .flat_map(|a| {
253                [
254                    PathBuf::from(format!("{}.data", a.sha256)),
255                    PathBuf::from(format!("{}.part", a.sha256)),
256                ]
257            })
258            .collect();
259        validate_tree(&files, &allowed)?;
260        Ok(Self {
261            root,
262            files,
263            allowed,
264            _lock: lock,
265        })
266    }
267    pub fn stage(&self, expected: &BTreeSet<PathBuf>) -> Result<PathBuf> {
268        let tree = self.root.join("tree");
269        if fs::symlink_metadata(&tree).is_ok() {
270            validate_tree(&tree, expected)?;
271            fs::remove_dir_all(&tree)?;
272        }
273        fs::create_dir(&tree)?;
274        Ok(tree)
275    }
276    pub fn cleanup(&self) -> Result<()> {
277        validate_tree(&self.files, &self.allowed)?;
278        for entry in fs::read_dir(&self.files)? {
279            fs::remove_file(entry?.path())?;
280        }
281        Ok(())
282    }
283    pub fn retrieve(
284        &self,
285        a: &Artifact,
286        source: &impl Source,
287        limit: u64,
288    ) -> Result<(PathBuf, u64)> {
289        let ready = self.files.join(format!("{}.data", a.sha256));
290        let part = self.files.join(format!("{}.part", a.sha256));
291        if regular(&ready)?.is_some() {
292            if let Ok((n, h)) = hash_file(&ready, limit) {
293                if h == a.sha256 {
294                    eprintln_status("Reuse", a, n);
295                    return Ok((ready, n));
296                }
297            }
298            fs::remove_file(&ready)?;
299        }
300        let mut offset = regular(&part)?.unwrap_or(0);
301        if offset > limit {
302            fs::remove_file(&part)?;
303            offset = 0;
304        }
305        if offset > 0 {
306            let (n, h) = hash_file(&part, limit)?;
307            if h == a.sha256 {
308                fs::rename(&part, &ready)?;
309                return Ok((ready, n));
310            }
311        }
312        let mut file = File::options()
313            .read(true)
314            .write(true)
315            .create(true)
316            .truncate(false)
317            .open(&part)?;
318        file.seek(SeekFrom::End(0))?;
319        eprintln_status(if offset > 0 { "Resume" } else { "Download" }, a, offset);
320        let transfer = source.fetch_from(
321            &format!("{ORIGIN}{}", a.path),
322            &mut file,
323            offset,
324            limit - offset,
325        );
326        file.sync_all()?;
327        let transfer = transfer?;
328        if transfer == Transfer::Restart {
329            file.set_len(0)?;
330            file.seek(SeekFrom::Start(0))?;
331            eprintln_status("Restart", a, 0);
332            let transfer = source.fetch_from(&format!("{ORIGIN}{}", a.path), &mut file, 0, limit);
333            file.sync_all()?;
334            if transfer? != Transfer::Complete {
335                return Err("source cannot restart download".into());
336            }
337        }
338        drop(file);
339        let (n, h) = hash_file(&part, limit)?;
340        if h != a.sha256 {
341            fs::remove_file(&part)?;
342            return Err(
343                format!("download checksum mismatch; discarded partial: {}", a.path).into(),
344            );
345        }
346        fs::rename(&part, &ready)?;
347        Ok((ready, n))
348    }
349}
350fn eprintln_status(action: &str, a: &Artifact, n: u64) {
351    let _ = writeln!(io::stderr().lock(), "{action}: {} ({n} bytes)", a.path);
352}
353
354#[cfg(test)]
355mod tests {
356    use super::*;
357    use std::net::TcpListener;
358    fn response(
359        status: &str,
360        headers: &str,
361        body: &str,
362        offset: u64,
363        limit: u64,
364    ) -> (Result<Transfer>, Vec<u8>, String) {
365        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
366        let address = listener.local_addr().unwrap();
367        let bytes = format!("HTTP/1.1 {status}\r\nConnection: close\r\n{headers}\r\n{body}");
368        let thread = std::thread::spawn(move || {
369            let (mut socket, _) = listener.accept().unwrap();
370            socket
371                .set_read_timeout(Some(std::time::Duration::from_secs(5)))
372                .unwrap();
373            let mut data = Vec::new();
374            let mut byte = [0];
375            while socket.read(&mut byte).unwrap_or(0) > 0 {
376                data.push(byte[0]);
377                if data.ends_with(b"\r\n\r\n") {
378                    break;
379                }
380            }
381            let _ = socket.write_all(bytes.as_bytes());
382            String::from_utf8(data).unwrap()
383        });
384        let agent = ureq::AgentBuilder::new()
385            .redirects(0)
386            .timeout(std::time::Duration::from_secs(5))
387            .build();
388        let mut out = Vec::new();
389        let result = http_range(
390            &agent,
391            &format!("http://{address}/file"),
392            &mut out,
393            offset,
394            limit,
395        );
396        (result, out, thread.join().unwrap())
397    }
398    #[test]
399    fn valid_range_and_full_responses_preserve_wire_bytes() {
400        let (r, b, h) = response(
401            "206 Partial Content",
402            "Content-Range: bytes 2-5/6\r\nContent-Length: 4\r\n",
403            "cdef",
404            2,
405            4,
406        );
407        assert_eq!(r.unwrap(), Transfer::Complete);
408        assert_eq!(b, b"cdef");
409        assert!(h.to_lowercase().contains("range: bytes=2-"));
410        assert!(h.to_lowercase().contains("accept-encoding: identity"));
411        let (r, b, _) = response("200 OK", "Content-Length: 6\r\n", "abcdef", 0, 6);
412        assert_eq!(r.unwrap(), Transfer::Complete);
413        assert_eq!(b, b"abcdef");
414    }
415    #[test]
416    fn ignored_and_unsatisfiable_ranges_request_restart_without_appending() {
417        for status in ["200 OK", "416 Range Not Satisfiable"] {
418            let (r, b, _) = response(status, "Content-Length: 3\r\n", "bad", 2, 4);
419            assert_eq!(r.unwrap(), Transfer::Restart);
420            assert!(b.is_empty());
421        }
422    }
423    #[test]
424    fn malformed_ranges_encoding_redirects_and_oversize_headers_refuse() {
425        for (status, headers, offset, limit) in [
426            ("206 Partial Content", "Content-Length: 4\r\n", 2, 4),
427            (
428                "206 Partial Content",
429                "Content-Range: bytes 1-5/6\r\n",
430                2,
431                4,
432            ),
433            (
434                "206 Partial Content",
435                "Content-Range: bytes 2-4/6\r\n",
436                2,
437                4,
438            ),
439            (
440                "206 Partial Content",
441                "Content-Range: bytes 2-5/*\r\n",
442                2,
443                4,
444            ),
445            (
446                "206 Partial Content",
447                "Content-Range: bytes 2-5/6\r\nContent-Length: 3\r\n",
448                2,
449                4,
450            ),
451            (
452                "206 Partial Content",
453                "Content-Range: bytes 2-5/6\r\n",
454                2,
455                3,
456            ),
457            ("200 OK", "Content-Length: 6\r\n", 0, 5),
458            (
459                "200 OK",
460                "Content-Encoding: br\r\nContent-Length: 4\r\n",
461                0,
462                4,
463            ),
464            (
465                "302 Found",
466                "Location: http://127.0.0.1:1/elsewhere\r\nContent-Length: 0\r\n",
467                2,
468                4,
469            ),
470            (
471                "206 Partial Content",
472                "Content-Range: bytes 2-18446744073709551615/0\r\n",
473                2,
474                4,
475            ),
476        ] {
477            let (r, b, _) = response(status, headers, "cdef", offset, limit);
478            assert!(r.is_err(), "{headers}");
479            assert!(b.is_empty(), "{headers}");
480        }
481    }
482    #[test]
483    fn truncated_and_unbounded_bodies_never_report_success() {
484        let (r, b, _) = response(
485            "206 Partial Content",
486            "Content-Range: bytes 2-5/6\r\nContent-Length: 4\r\n",
487            "cd",
488            2,
489            4,
490        );
491        assert!(r.is_err());
492        assert!(b.len() <= 2);
493        let (r, b, _) = response("200 OK", "", "abcdef", 0, 3);
494        assert!(r.is_err());
495        assert!(b.len() <= 3);
496    }
497}