1use 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}
140pub(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}