fluidattacks_core_git/
download.rs

1use anyhow::{bail, Context, Result};
2use std::path::{Path, PathBuf};
3use tokio::io::{AsyncWriteExt, BufWriter};
4use tracing::warn;
5
6const MAX_RETRIES: u32 = 3;
7const RETRY_SLEEP_SECS: u64 = 60;
8const TOTAL_TIMEOUT_SECS: u64 = 60 * 60;
9const CONNECT_TIMEOUT_SECS: u64 = 30;
10const DEFAULT_DOWNLOAD_BUFFER_SIZE: usize = 64 * 1024;
11const MAX_DOWNLOAD_BUFFER_SIZE: usize = 16 * 1024 * 1024;
12const GZIP_MAGIC: &[u8] = &[0x1f, 0x8b];
13const TMP_PREFIX: &str = "fluidattacks_";
14const ARCHIVE_NAME: &str = "repo.tar.gz";
15
16enum Failure {
17    Transient(anyhow::Error),
18    Permanent(anyhow::Error),
19}
20
21fn retryable_status(status: reqwest::StatusCode) -> bool {
22    status.is_server_error() || status == reqwest::StatusCode::TOO_MANY_REQUESTS
23}
24
25fn resolve_buffer_size(buffer_size: Option<usize>) -> Result<usize> {
26    match buffer_size {
27        None => Ok(DEFAULT_DOWNLOAD_BUFFER_SIZE),
28        Some(0) => bail!("the download buffer size must be greater than zero"),
29        Some(size) if size > MAX_DOWNLOAD_BUFFER_SIZE => bail!(
30            "the download buffer size must not exceed {MAX_DOWNLOAD_BUFFER_SIZE} bytes, got {size}"
31        ),
32        Some(size) => Ok(size),
33    }
34}
35
36/// Writes the body of `url` to `path`.
37///
38/// `buffer_size` bounds a single write to disk; `None` takes the 64 KiB the
39/// python defaults to, and anything above 16 MiB is rejected rather than
40/// clamped, since the batch container that clones has 7600 MiB for the whole
41/// job. Retries a transport error, a 5xx and a 429; an expired or a wrong URL
42/// fails at once. Leaves a partial archive at `path` when the body is not
43/// gzip, which the caller must not hand to an extractor.
44pub async fn download_to(url: &str, path: &Path, buffer_size: Option<usize>) -> Result<()> {
45    let buffer_size = resolve_buffer_size(buffer_size)?;
46    let mut last_err = None;
47
48    for attempt in 0..MAX_RETRIES {
49        match do_download_to(url, path, buffer_size).await {
50            Ok(()) => return Ok(()),
51            Err(Failure::Permanent(e)) => return Err(e),
52            Err(Failure::Transient(e)) => {
53                warn!(
54                    attempt = attempt.saturating_add(1),
55                    max = MAX_RETRIES,
56                    error = %e,
57                    "download failed, retrying"
58                );
59                last_err = Some(e);
60                if attempt < MAX_RETRIES - 1 {
61                    tokio::time::sleep(std::time::Duration::from_secs(RETRY_SLEEP_SECS)).await;
62                }
63            }
64        }
65    }
66
67    Err(last_err.unwrap_or_else(|| anyhow::anyhow!("download failed")))
68}
69
70async fn open_stream(url: &str) -> Result<reqwest::Response, Failure> {
71    let resp = reqwest::Client::builder()
72        .timeout(std::time::Duration::from_secs(TOTAL_TIMEOUT_SECS))
73        .connect_timeout(std::time::Duration::from_secs(CONNECT_TIMEOUT_SECS))
74        .build()
75        .context("building the HTTP client")
76        .map_err(Failure::Permanent)?
77        .get(url)
78        .send()
79        .await
80        .context("HTTP request failed")
81        .map_err(Failure::Transient)?;
82
83    let status = resp.status();
84    if status == reqwest::StatusCode::OK {
85        return Ok(resp);
86    }
87
88    let error = anyhow::anyhow!("HTTP {status}");
89    Err(if retryable_status(status) {
90        Failure::Transient(error)
91    } else {
92        Failure::Permanent(error)
93    })
94}
95
96async fn do_download_to(url: &str, path: &Path, buffer_size: usize) -> Result<(), Failure> {
97    let mut resp = open_stream(url).await?;
98
99    let file = tokio::fs::File::create(path)
100        .await
101        .with_context(|| format!("creating {}", path.display()))
102        .map_err(Failure::Permanent)?;
103    let mut writer = BufWriter::with_capacity(buffer_size, file);
104    let mut header = Vec::with_capacity(GZIP_MAGIC.len());
105    let mut written: usize = 0;
106
107    while let Some(chunk) = resp
108        .chunk()
109        .await
110        .context("reading response body")
111        .map_err(Failure::Transient)?
112    {
113        take_header(&mut header, &chunk).map_err(Failure::Permanent)?;
114        for slice in chunk.chunks(buffer_size) {
115            writer
116                .write_all(slice)
117                .await
118                .context("writing the archive")
119                .map_err(Failure::Permanent)?;
120        }
121        written = written.saturating_add(chunk.len());
122    }
123    writer
124        .flush()
125        .await
126        .context("flushing the archive")
127        .map_err(Failure::Permanent)?;
128
129    if !header.starts_with(GZIP_MAGIC) {
130        return Err(Failure::Permanent(anyhow::anyhow!(
131            "response is not a valid gzip file (got {written} bytes)"
132        )));
133    }
134
135    Ok(())
136}
137
138fn take_header(header: &mut Vec<u8>, chunk: &[u8]) -> Result<()> {
139    let missing = GZIP_MAGIC.len().saturating_sub(header.len());
140    if missing == 0 {
141        return Ok(());
142    }
143
144    header.extend(chunk.iter().take(missing).copied());
145    if header.len() == GZIP_MAGIC.len() && !header.starts_with(GZIP_MAGIC) {
146        bail!("response is not a valid gzip file");
147    }
148
149    Ok(())
150}
151
152fn archive_parent(dest: &Path) -> Option<&Path> {
153    match dest.parent() {
154        Some(parent) if parent.as_os_str().is_empty() => Some(Path::new(".")),
155        parent => parent,
156    }
157}
158
159fn extracted_dir(tmp_path: &Path) -> Result<PathBuf> {
160    let mut dirs: Vec<PathBuf> = std::fs::read_dir(tmp_path)
161        .with_context(|| format!("reading {}", tmp_path.display()))?
162        .filter_map(std::result::Result::ok)
163        .filter(|entry| entry.file_type().is_ok_and(|kind| kind.is_dir()))
164        .map(|entry| entry.path())
165        .collect();
166    dirs.sort();
167
168    if dirs.len() > 1 {
169        warn!(
170            count = dirs.len(),
171            "multiple directories found in archive, using the first one"
172        );
173    }
174
175    dirs.into_iter()
176        .next()
177        .context("no directory found in the extracted archive")
178}
179
180fn move_or_copy(from: &Path, to: &Path) -> Result<()> {
181    if let Err(rename_error) = std::fs::rename(from, to) {
182        super::fs::copy_dir_all(from, to).with_context(|| {
183            format!(
184                "moving {} to {} failed with {rename_error}, and so did copying it",
185                from.display(),
186                to.display()
187            )
188        })?;
189    }
190    Ok(())
191}
192
193fn extract_and_replace(archive_path: &Path, tmp_path: &Path, dest: &Path) -> Result<()> {
194    let archive = std::fs::File::open(archive_path)
195        .with_context(|| format!("opening {}", archive_path.display()))?;
196    super::tar::extract(archive, tmp_path)?;
197    let _ = std::fs::remove_file(archive_path);
198
199    let extracted = extracted_dir(tmp_path)?;
200
201    if dest.exists() {
202        std::fs::remove_dir_all(dest).with_context(|| format!("removing {}", dest.display()))?;
203    }
204
205    move_or_copy(&extracted, dest)
206}
207
208async fn place_repo(url: &str, dest: &Path, buffer_size: Option<usize>) -> bool {
209    let Some(parent) = archive_parent(dest) else {
210        warn!("the destination has no parent directory");
211        return false;
212    };
213    if let Err(e) = std::fs::create_dir_all(parent) {
214        warn!(error = %e, "failed to create parent dir");
215        return false;
216    }
217
218    let tmp_dir = match tempfile::Builder::new().prefix(TMP_PREFIX).tempdir() {
219        Ok(d) => d,
220        Err(e) => {
221            warn!(error = %e, "failed to create temp dir");
222            return false;
223        }
224    };
225    let tmp_path = tmp_dir.path().to_path_buf();
226    let archive_path = tmp_path.join(ARCHIVE_NAME);
227
228    if let Err(e) = download_to(url, &archive_path, buffer_size).await {
229        warn!(error = %e, "failed to download");
230        return false;
231    }
232
233    let dest_path = dest.to_path_buf();
234    let extracted = tokio::task::spawn_blocking(move || {
235        extract_and_replace(&archive_path, &tmp_path, &dest_path)
236    })
237    .await
238    .unwrap_or_else(|e| Err(anyhow::Error::new(e)));
239
240    if let Err(e) = extracted {
241        warn!(error = %e, "failed to extract archive");
242        return false;
243    }
244
245    true
246}
247
248/// Leaves the repository behind `url` at `dest`.
249///
250/// The archive holds only `.git`, so the working tree is restored with a
251/// `git reset --hard`. Pass `checkout = false` to keep just the history: peak
252/// disk stays at the size of the archive plus `.git`, which is what a consumer
253/// that walks commits needs. `git_ignore` prunes working-tree files and has no
254/// effect on a history-only download. `buffer_size` bounds a single write of
255/// the archive to disk, as it does in the python, and `None` takes the library
256/// default. Anything already at `dest` is replaced.
257pub async fn download_and_extract_repo(
258    url: &str,
259    dest: &Path,
260    git_ignore: &[String],
261    buffer_size: Option<usize>,
262    checkout: bool,
263) -> bool {
264    if !place_repo(url, dest, buffer_size).await {
265        return false;
266    }
267
268    let _ = super::cli::set_safe_directory().await;
269    if !checkout {
270        return true;
271    }
272
273    if super::cli::reset_hard(dest).await.is_err() {
274        let _ = std::fs::remove_dir_all(dest);
275        return false;
276    }
277
278    let dest_path = dest.to_path_buf();
279    let patterns = git_ignore.to_vec();
280    let pruned = tokio::task::spawn_blocking(move || {
281        super::fs::remove_symlinks(&dest_path)?;
282        super::fs::delete_gitignored_files(&dest_path, &patterns)
283    })
284    .await;
285    match pruned {
286        Ok(Ok(())) => true,
287        Ok(Err(error)) => {
288            warn!(error = %error, "failed to prune the working tree");
289            false
290        }
291        Err(error) => {
292            warn!(error = %error, "failed to prune the working tree");
293            false
294        }
295    }
296}
297
298#[cfg(test)]
299mod tests {
300    use super::{
301        extract_and_replace, extracted_dir, resolve_buffer_size, retryable_status, take_header,
302        DEFAULT_DOWNLOAD_BUFFER_SIZE, GZIP_MAGIC, MAX_DOWNLOAD_BUFFER_SIZE,
303    };
304    use flate2::write::GzEncoder;
305    use flate2::Compression;
306    use std::path::Path;
307
308    fn write_tar_gz(path: &Path, name: &str, contents: &[u8]) {
309        let mut builder = tar::Builder::new(GzEncoder::new(Vec::new(), Compression::default()));
310        let mut header = tar::Header::new_gnu();
311        header.set_size(u64::try_from(contents.len()).unwrap());
312        header.set_mode(0o644);
313        builder.append_data(&mut header, name, contents).unwrap();
314        let archive = builder.into_inner().unwrap().finish().unwrap();
315        std::fs::write(path, archive).unwrap();
316    }
317
318    #[test]
319    fn the_magic_bytes_can_arrive_split_across_chunks() {
320        let mut header = Vec::new();
321        assert!(take_header(&mut header, &[0x1f]).is_ok());
322        assert!(take_header(&mut header, &[0x8b, 0x08]).is_ok());
323        assert_eq!(header, GZIP_MAGIC);
324    }
325
326    #[test]
327    fn a_body_that_is_not_gzip_is_rejected_on_the_first_chunk() {
328        let mut header = Vec::new();
329        assert!(take_header(&mut header, b"<?xml version=\"1.0\"?>").is_err());
330    }
331
332    #[test]
333    fn later_chunks_are_not_inspected_once_the_magic_is_known() {
334        let mut header = Vec::new();
335        assert!(take_header(&mut header, GZIP_MAGIC).is_ok());
336        assert!(take_header(&mut header, b"anything at all").is_ok());
337        assert_eq!(header, GZIP_MAGIC);
338    }
339
340    #[test]
341    fn a_body_shorter_than_the_magic_never_completes_it() {
342        let mut header = Vec::new();
343        assert!(take_header(&mut header, &[0x1f]).is_ok());
344        assert!(!header.starts_with(GZIP_MAGIC));
345    }
346
347    #[test]
348    fn the_archive_replaces_an_existing_destination_and_is_deleted() {
349        let tmp = tempfile::tempdir().unwrap();
350        let work = tmp.path().join("work");
351        std::fs::create_dir_all(&work).unwrap();
352        let archive_path = work.join("repo.tar.gz");
353        write_tar_gz(&archive_path, "repo/HEAD", b"ref: refs/heads/trunk\n");
354
355        let dest = tmp.path().join("dest");
356        std::fs::create_dir_all(&dest).unwrap();
357        std::fs::write(dest.join("stale"), b"old").unwrap();
358
359        extract_and_replace(&archive_path, &work, &dest).unwrap();
360
361        assert_eq!(
362            std::fs::read_to_string(dest.join("HEAD")).unwrap(),
363            "ref: refs/heads/trunk\n"
364        );
365        assert!(!dest.join("stale").exists());
366        assert!(!archive_path.exists());
367    }
368
369    #[test]
370    fn an_archive_without_a_directory_is_rejected() {
371        let tmp = tempfile::tempdir().unwrap();
372        let work = tmp.path().join("work");
373        std::fs::create_dir_all(&work).unwrap();
374        let archive_path = work.join("repo.tar.gz");
375        write_tar_gz(&archive_path, "HEAD", b"ref: refs/heads/trunk\n");
376
377        let dest = tmp.path().join("dest");
378
379        assert!(extract_and_replace(&archive_path, &work, &dest).is_err());
380        assert!(!dest.exists());
381    }
382
383    #[test]
384    fn a_url_that_is_expired_or_wrong_is_not_retried() {
385        assert!(!retryable_status(reqwest::StatusCode::FORBIDDEN));
386        assert!(!retryable_status(reqwest::StatusCode::NOT_FOUND));
387        assert!(retryable_status(reqwest::StatusCode::TOO_MANY_REQUESTS));
388        assert!(retryable_status(reqwest::StatusCode::INTERNAL_SERVER_ERROR));
389        assert!(retryable_status(reqwest::StatusCode::SERVICE_UNAVAILABLE));
390    }
391
392    #[test]
393    fn an_absent_buffer_size_takes_the_default_of_the_python() {
394        assert_eq!(
395            resolve_buffer_size(None).unwrap(),
396            DEFAULT_DOWNLOAD_BUFFER_SIZE
397        );
398        assert_eq!(resolve_buffer_size(Some(8 * 1024)).unwrap(), 8 * 1024);
399    }
400
401    #[test]
402    fn a_buffer_size_of_zero_is_rejected_instead_of_writing_an_empty_archive() {
403        assert!(resolve_buffer_size(Some(0)).is_err());
404    }
405
406    #[test]
407    fn a_buffer_size_above_the_maximum_is_rejected_instead_of_being_clamped() {
408        assert_eq!(
409            resolve_buffer_size(Some(MAX_DOWNLOAD_BUFFER_SIZE)).unwrap(),
410            MAX_DOWNLOAD_BUFFER_SIZE
411        );
412        assert!(resolve_buffer_size(Some(MAX_DOWNLOAD_BUFFER_SIZE + 1)).is_err());
413    }
414
415    #[test]
416    fn the_first_directory_is_taken_in_a_stable_order() {
417        let tmp = tempfile::tempdir().unwrap();
418        std::fs::write(tmp.path().join("repo.tar.gz"), b"").unwrap();
419        std::fs::create_dir_all(tmp.path().join("b_repo")).unwrap();
420        std::fs::create_dir_all(tmp.path().join("a_repo")).unwrap();
421
422        assert_eq!(
423            extracted_dir(tmp.path()).unwrap(),
424            tmp.path().join("a_repo")
425        );
426    }
427}
fluidattacks_core_git/download.rs

fluidattacks_core_git/
download.rs
