fluidattacks-core-git 0.36.0

use anyhow::{bail, Context, Result};
use std::path::{Path, PathBuf};
use tokio::io::{AsyncWriteExt, BufWriter};
use tracing::warn;

const MAX_RETRIES: u32 = 3;
const RETRY_SLEEP_SECS: u64 = 60;
const TOTAL_TIMEOUT_SECS: u64 = 60 * 60;
const CONNECT_TIMEOUT_SECS: u64 = 30;
const DEFAULT_DOWNLOAD_BUFFER_SIZE: usize = 64 * 1024;
const MAX_DOWNLOAD_BUFFER_SIZE: usize = 16 * 1024 * 1024;
const GZIP_MAGIC: &[u8] = &[0x1f, 0x8b];
const TMP_PREFIX: &str = "fluidattacks_";
const ARCHIVE_NAME: &str = "repo.tar.gz";

enum Failure {
    Transient(anyhow::Error),
    Permanent(anyhow::Error),
}

fn retryable_status(status: reqwest::StatusCode) -> bool {
    status.is_server_error() || status == reqwest::StatusCode::TOO_MANY_REQUESTS
}

fn resolve_buffer_size(buffer_size: Option<usize>) -> Result<usize> {
    match buffer_size {
        None => Ok(DEFAULT_DOWNLOAD_BUFFER_SIZE),
        Some(0) => bail!("the download buffer size must be greater than zero"),
        Some(size) if size > MAX_DOWNLOAD_BUFFER_SIZE => bail!(
            "the download buffer size must not exceed {MAX_DOWNLOAD_BUFFER_SIZE} bytes, got {size}"
        ),
        Some(size) => Ok(size),
    }
}

/// Writes the body of `url` to `path`.
///
/// `buffer_size` bounds a single write to disk; `None` takes the 64 KiB the
/// python defaults to, and anything above 16 MiB is rejected rather than
/// clamped, since the batch container that clones has 7600 MiB for the whole
/// job. Retries a transport error, a 5xx and a 429; an expired or a wrong URL
/// fails at once. Leaves a partial archive at `path` when the body is not
/// gzip, which the caller must not hand to an extractor.
pub async fn download_to(url: &str, path: &Path, buffer_size: Option<usize>) -> Result<()> {
    let buffer_size = resolve_buffer_size(buffer_size)?;
    let mut last_err = None;

    for attempt in 0..MAX_RETRIES {
        match do_download_to(url, path, buffer_size).await {
            Ok(()) => return Ok(()),
            Err(Failure::Permanent(e)) => return Err(e),
            Err(Failure::Transient(e)) => {
                warn!(
                    attempt = attempt.saturating_add(1),
                    max = MAX_RETRIES,
                    error = %e,
                    "download failed, retrying"
                );
                last_err = Some(e);
                if attempt < MAX_RETRIES - 1 {
                    tokio::time::sleep(std::time::Duration::from_secs(RETRY_SLEEP_SECS)).await;
                }
            }
        }
    }

    Err(last_err.unwrap_or_else(|| anyhow::anyhow!("download failed")))
}

async fn open_stream(url: &str) -> Result<reqwest::Response, Failure> {
    let resp = reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(TOTAL_TIMEOUT_SECS))
        .connect_timeout(std::time::Duration::from_secs(CONNECT_TIMEOUT_SECS))
        .build()
        .context("building the HTTP client")
        .map_err(Failure::Permanent)?
        .get(url)
        .send()
        .await
        .context("HTTP request failed")
        .map_err(Failure::Transient)?;

    let status = resp.status();
    if status == reqwest::StatusCode::OK {
        return Ok(resp);
    }

    let error = anyhow::anyhow!("HTTP {status}");
    Err(if retryable_status(status) {
        Failure::Transient(error)
    } else {
        Failure::Permanent(error)
    })
}

async fn do_download_to(url: &str, path: &Path, buffer_size: usize) -> Result<(), Failure> {
    let mut resp = open_stream(url).await?;

    let file = tokio::fs::File::create(path)
        .await
        .with_context(|| format!("creating {}", path.display()))
        .map_err(Failure::Permanent)?;
    let mut writer = BufWriter::with_capacity(buffer_size, file);
    let mut header = Vec::with_capacity(GZIP_MAGIC.len());
    let mut written: usize = 0;

    while let Some(chunk) = resp
        .chunk()
        .await
        .context("reading response body")
        .map_err(Failure::Transient)?
    {
        take_header(&mut header, &chunk).map_err(Failure::Permanent)?;
        for slice in chunk.chunks(buffer_size) {
            writer
                .write_all(slice)
                .await
                .context("writing the archive")
                .map_err(Failure::Permanent)?;
        }
        written = written.saturating_add(chunk.len());
    }
    writer
        .flush()
        .await
        .context("flushing the archive")
        .map_err(Failure::Permanent)?;

    if !header.starts_with(GZIP_MAGIC) {
        return Err(Failure::Permanent(anyhow::anyhow!(
            "response is not a valid gzip file (got {written} bytes)"
        )));
    }

    Ok(())
}

fn take_header(header: &mut Vec<u8>, chunk: &[u8]) -> Result<()> {
    let missing = GZIP_MAGIC.len().saturating_sub(header.len());
    if missing == 0 {
        return Ok(());
    }

    header.extend(chunk.iter().take(missing).copied());
    if header.len() == GZIP_MAGIC.len() && !header.starts_with(GZIP_MAGIC) {
        bail!("response is not a valid gzip file");
    }

    Ok(())
}

fn archive_parent(dest: &Path) -> Option<&Path> {
    match dest.parent() {
        Some(parent) if parent.as_os_str().is_empty() => Some(Path::new(".")),
        parent => parent,
    }
}

fn extracted_dir(tmp_path: &Path) -> Result<PathBuf> {
    let mut dirs: Vec<PathBuf> = std::fs::read_dir(tmp_path)
        .with_context(|| format!("reading {}", tmp_path.display()))?
        .filter_map(std::result::Result::ok)
        .filter(|entry| entry.file_type().is_ok_and(|kind| kind.is_dir()))
        .map(|entry| entry.path())
        .collect();
    dirs.sort();

    if dirs.len() > 1 {
        warn!(
            count = dirs.len(),
            "multiple directories found in archive, using the first one"
        );
    }

    dirs.into_iter()
        .next()
        .context("no directory found in the extracted archive")
}

fn move_or_copy(from: &Path, to: &Path) -> Result<()> {
    if let Err(rename_error) = std::fs::rename(from, to) {
        super::fs::copy_dir_all(from, to).with_context(|| {
            format!(
                "moving {} to {} failed with {rename_error}, and so did copying it",
                from.display(),
                to.display()
            )
        })?;
    }
    Ok(())
}

fn extract_and_replace(archive_path: &Path, tmp_path: &Path, dest: &Path) -> Result<()> {
    let archive = std::fs::File::open(archive_path)
        .with_context(|| format!("opening {}", archive_path.display()))?;
    super::tar::extract(archive, tmp_path)?;
    let _ = std::fs::remove_file(archive_path);

    let extracted = extracted_dir(tmp_path)?;

    if dest.exists() {
        std::fs::remove_dir_all(dest).with_context(|| format!("removing {}", dest.display()))?;
    }

    move_or_copy(&extracted, dest)
}

async fn place_repo(url: &str, dest: &Path, buffer_size: Option<usize>) -> bool {
    let Some(parent) = archive_parent(dest) else {
        warn!("the destination has no parent directory");
        return false;
    };
    if let Err(e) = std::fs::create_dir_all(parent) {
        warn!(error = %e, "failed to create parent dir");
        return false;
    }

    let tmp_dir = match tempfile::Builder::new().prefix(TMP_PREFIX).tempdir() {
        Ok(d) => d,
        Err(e) => {
            warn!(error = %e, "failed to create temp dir");
            return false;
        }
    };
    let tmp_path = tmp_dir.path().to_path_buf();
    let archive_path = tmp_path.join(ARCHIVE_NAME);

    if let Err(e) = download_to(url, &archive_path, buffer_size).await {
        warn!(error = %e, "failed to download");
        return false;
    }

    let dest_path = dest.to_path_buf();
    let extracted = tokio::task::spawn_blocking(move || {
        extract_and_replace(&archive_path, &tmp_path, &dest_path)
    })
    .await
    .unwrap_or_else(|e| Err(anyhow::Error::new(e)));

    if let Err(e) = extracted {
        warn!(error = %e, "failed to extract archive");
        return false;
    }

    true
}

/// Leaves the repository behind `url` at `dest`.
///
/// The archive holds only `.git`, so the working tree is restored with a
/// `git reset --hard`. Pass `checkout = false` to keep just the history: peak
/// disk stays at the size of the archive plus `.git`, which is what a consumer
/// that walks commits needs. `git_ignore` prunes working-tree files and has no
/// effect on a history-only download. `buffer_size` bounds a single write of
/// the archive to disk, as it does in the python, and `None` takes the library
/// default. Anything already at `dest` is replaced.
pub async fn download_and_extract_repo(
    url: &str,
    dest: &Path,
    git_ignore: &[String],
    buffer_size: Option<usize>,
    checkout: bool,
) -> bool {
    if !place_repo(url, dest, buffer_size).await {
        return false;
    }

    let _ = super::cli::set_safe_directory().await;
    if !checkout {
        return true;
    }

    if super::cli::reset_hard(dest).await.is_err() {
        let _ = std::fs::remove_dir_all(dest);
        return false;
    }

    let dest_path = dest.to_path_buf();
    let patterns = git_ignore.to_vec();
    let pruned = tokio::task::spawn_blocking(move || {
        super::fs::remove_symlinks(&dest_path)?;
        super::fs::delete_gitignored_files(&dest_path, &patterns)
    })
    .await;
    match pruned {
        Ok(Ok(())) => true,
        Ok(Err(error)) => {
            warn!(error = %error, "failed to prune the working tree");
            false
        }
        Err(error) => {
            warn!(error = %error, "failed to prune the working tree");
            false
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{
        extract_and_replace, extracted_dir, resolve_buffer_size, retryable_status, take_header,
        DEFAULT_DOWNLOAD_BUFFER_SIZE, GZIP_MAGIC, MAX_DOWNLOAD_BUFFER_SIZE,
    };
    use flate2::write::GzEncoder;
    use flate2::Compression;
    use std::path::Path;

    fn write_tar_gz(path: &Path, name: &str, contents: &[u8]) {
        let mut builder = tar::Builder::new(GzEncoder::new(Vec::new(), Compression::default()));
        let mut header = tar::Header::new_gnu();
        header.set_size(u64::try_from(contents.len()).unwrap());
        header.set_mode(0o644);
        builder.append_data(&mut header, name, contents).unwrap();
        let archive = builder.into_inner().unwrap().finish().unwrap();
        std::fs::write(path, archive).unwrap();
    }

    #[test]
    fn the_magic_bytes_can_arrive_split_across_chunks() {
        let mut header = Vec::new();
        assert!(take_header(&mut header, &[0x1f]).is_ok());
        assert!(take_header(&mut header, &[0x8b, 0x08]).is_ok());
        assert_eq!(header, GZIP_MAGIC);
    }

    #[test]
    fn a_body_that_is_not_gzip_is_rejected_on_the_first_chunk() {
        let mut header = Vec::new();
        assert!(take_header(&mut header, b"<?xml version=\"1.0\"?>").is_err());
    }

    #[test]
    fn later_chunks_are_not_inspected_once_the_magic_is_known() {
        let mut header = Vec::new();
        assert!(take_header(&mut header, GZIP_MAGIC).is_ok());
        assert!(take_header(&mut header, b"anything at all").is_ok());
        assert_eq!(header, GZIP_MAGIC);
    }

    #[test]
    fn a_body_shorter_than_the_magic_never_completes_it() {
        let mut header = Vec::new();
        assert!(take_header(&mut header, &[0x1f]).is_ok());
        assert!(!header.starts_with(GZIP_MAGIC));
    }

    #[test]
    fn the_archive_replaces_an_existing_destination_and_is_deleted() {
        let tmp = tempfile::tempdir().unwrap();
        let work = tmp.path().join("work");
        std::fs::create_dir_all(&work).unwrap();
        let archive_path = work.join("repo.tar.gz");
        write_tar_gz(&archive_path, "repo/HEAD", b"ref: refs/heads/trunk\n");

        let dest = tmp.path().join("dest");
        std::fs::create_dir_all(&dest).unwrap();
        std::fs::write(dest.join("stale"), b"old").unwrap();

        extract_and_replace(&archive_path, &work, &dest).unwrap();

        assert_eq!(
            std::fs::read_to_string(dest.join("HEAD")).unwrap(),
            "ref: refs/heads/trunk\n"
        );
        assert!(!dest.join("stale").exists());
        assert!(!archive_path.exists());
    }

    #[test]
    fn an_archive_without_a_directory_is_rejected() {
        let tmp = tempfile::tempdir().unwrap();
        let work = tmp.path().join("work");
        std::fs::create_dir_all(&work).unwrap();
        let archive_path = work.join("repo.tar.gz");
        write_tar_gz(&archive_path, "HEAD", b"ref: refs/heads/trunk\n");

        let dest = tmp.path().join("dest");

        assert!(extract_and_replace(&archive_path, &work, &dest).is_err());
        assert!(!dest.exists());
    }

    #[test]
    fn a_url_that_is_expired_or_wrong_is_not_retried() {
        assert!(!retryable_status(reqwest::StatusCode::FORBIDDEN));
        assert!(!retryable_status(reqwest::StatusCode::NOT_FOUND));
        assert!(retryable_status(reqwest::StatusCode::TOO_MANY_REQUESTS));
        assert!(retryable_status(reqwest::StatusCode::INTERNAL_SERVER_ERROR));
        assert!(retryable_status(reqwest::StatusCode::SERVICE_UNAVAILABLE));
    }

    #[test]
    fn an_absent_buffer_size_takes_the_default_of_the_python() {
        assert_eq!(
            resolve_buffer_size(None).unwrap(),
            DEFAULT_DOWNLOAD_BUFFER_SIZE
        );
        assert_eq!(resolve_buffer_size(Some(8 * 1024)).unwrap(), 8 * 1024);
    }

    #[test]
    fn a_buffer_size_of_zero_is_rejected_instead_of_writing_an_empty_archive() {
        assert!(resolve_buffer_size(Some(0)).is_err());
    }

    #[test]
    fn a_buffer_size_above_the_maximum_is_rejected_instead_of_being_clamped() {
        assert_eq!(
            resolve_buffer_size(Some(MAX_DOWNLOAD_BUFFER_SIZE)).unwrap(),
            MAX_DOWNLOAD_BUFFER_SIZE
        );
        assert!(resolve_buffer_size(Some(MAX_DOWNLOAD_BUFFER_SIZE + 1)).is_err());
    }

    #[test]
    fn the_first_directory_is_taken_in_a_stable_order() {
        let tmp = tempfile::tempdir().unwrap();
        std::fs::write(tmp.path().join("repo.tar.gz"), b"").unwrap();
        std::fs::create_dir_all(tmp.path().join("b_repo")).unwrap();
        std::fs::create_dir_all(tmp.path().join("a_repo")).unwrap();

        assert_eq!(
            extracted_dir(tmp.path()).unwrap(),
            tmp.path().join("a_repo")
        );
    }
}