ambient-ci 0.14.0

//! Virtual drive handling for ambient-run.

use std::{
    fs::{File, OpenOptions},
    path::{Path, PathBuf},
    process::Command,
};

/// A virtual drive.
#[derive(Debug, Clone)]
pub struct VirtualDrive {
    filename: PathBuf,
}

impl VirtualDrive {
    /// Path to file containing virtual drive.
    pub fn filename(&self) -> &Path {
        self.filename.as_path()
    }

    /// Extract files in the virtual drive into a directory. Create
    /// the directory if it doesn't exist.
    pub fn extract_to(&self, dirname: &Path) -> Result<(), VirtualDriveError> {
        if !dirname.exists() {
            std::fs::create_dir(dirname)
                .map_err(|e| VirtualDriveError::Extract(dirname.into(), e))?;
        }
        tar_extract(&self.filename, dirname)?;
        Ok(())
    }
}

/// Builder for [`VirtualDrive`].
#[derive(Debug, Default)]
pub struct VirtualDriveBuilder {
    filename: Option<PathBuf>,
    root: Option<PathBuf>,
    size: Option<u64>,
}

impl VirtualDriveBuilder {
    /// Set filename for virtual drive.
    pub fn filename(mut self, filename: &Path) -> Self {
        self.filename = Some(filename.into());
        self
    }

    /// Set directory of tree to copy into virtual drive.
    pub fn root_directory(mut self, dirname: &Path) -> Self {
        self.root = Some(dirname.into());
        self
    }

    /// Set size of new drive. This is important when the build VM
    /// writes to the drive.
    pub fn size(mut self, size: u64) -> Self {
        self.size = Some(size);
        self
    }

    /// Create a virtual drive.
    pub fn create(self) -> Result<VirtualDrive, VirtualDriveError> {
        let filename = self.filename.expect("filename has been set");

        // Create the file, either empty or to the desired size. If we
        // don't have self.root set, the file created (and maybe
        // truncated) will be.
        {
            let archive = File::create(&filename)
                .map_err(|e| VirtualDriveError::Create(filename.clone(), e))?;
            if let Some(size) = self.size {
                archive
                    .set_len(size)
                    .map_err(|e| VirtualDriveError::Create(filename.clone(), e))?;
            }
        }

        if let Some(root) = self.root {
            match tar_create(&filename, &root) {
                Ok(_) => (),
                Err(VirtualDriveError::TarFailed(cmd, filename, exit, stderr)) => {
                    Err(VirtualDriveError::TarFailed(cmd, filename, exit, stderr))?;
                }
                Err(err) => {
                    Err(err)?;
                }
            }
        }

        Ok(VirtualDrive { filename })
    }

    /// Open an existing virtual drive.
    pub fn open(self) -> Result<VirtualDrive, VirtualDriveError> {
        let filename = self.filename.expect("filename has been set");
        Ok(VirtualDrive { filename })
    }
}

/// Create a tar archive out of a directory.
pub fn create_tar(
    tar_filename: PathBuf,
    dirname: &Path,
) -> Result<VirtualDrive, VirtualDriveError> {
    assert!(!tar_filename.starts_with(dirname));
    VirtualDriveBuilder::default()
        .filename(&tar_filename)
        .root_directory(dirname)
        .create()
}

/// Create a tar archive with a fixed length, out of a directory.
pub fn create_tar_with_size(
    tar_filename: PathBuf,
    dirname: &Path,
    size: u64,
) -> Result<VirtualDrive, VirtualDriveError> {
    let tar = VirtualDriveBuilder::default()
        .filename(&tar_filename)
        .root_directory(dirname)
        .size(size)
        .create()?;

    let metadata = std::fs::metadata(&tar_filename)
        .map_err(|err| VirtualDriveError::Metadata(tar_filename.clone(), err))?;
    if metadata.len() < size {
        let file = OpenOptions::new()
            .write(true)
            .truncate(false)
            .open(&tar_filename)
            .map_err(|err| VirtualDriveError::CreateTar(tar_filename.clone(), err))?;
        file.set_len(size)
            .map_err(|err| VirtualDriveError::SetLen(size, tar_filename.clone(), err))?;
    }

    let metadata = std::fs::metadata(&tar_filename)
        .map_err(|err| VirtualDriveError::Metadata(tar_filename.clone(), err))?;
    if metadata.len() > size {
        return Err(VirtualDriveError::DriveTooBig(metadata.len(), size));
    }

    Ok(tar)
}

fn tar_create(tar_filename: &Path, dirname: &Path) -> Result<(), VirtualDriveError> {
    let output = Command::new("tar")
        .arg("-cvf")
        .arg(tar_filename)
        .arg("-C")
        .arg(dirname)
        .arg(".")
        .output()
        .map_err(|err| VirtualDriveError::Tar("create", dirname.into(), err))?;

    if let Some(exit) = output.status.code() {
        if exit != 0 {
            let stderr = String::from_utf8_lossy(&output.stderr).to_string();
            return Err(VirtualDriveError::TarFailed(
                "create",
                dirname.into(),
                exit,
                stderr,
            ));
        }
    }

    Ok(())
}

fn tar_extract(tar_filename: &Path, dirname: &Path) -> Result<(), VirtualDriveError> {
    let output = Command::new("tar")
        .arg("-xvvvf")
        .arg(tar_filename)
        .arg("-C")
        .arg(dirname)
        .arg("--no-same-owner")
        .output()
        .map_err(|err| VirtualDriveError::Tar("extract", dirname.into(), err))?;

    if let Some(exit) = output.status.code() {
        if exit != 0 {
            let stderr = String::from_utf8_lossy(&output.stderr).to_string();
            return Err(VirtualDriveError::TarFailed(
                "extract",
                dirname.into(),
                exit,
                stderr,
            ));
        }
    }

    Ok(())
}

/// Errors that may be returned from [`VirtualDrive`] use.
#[allow(missing_docs)]
#[derive(Debug, thiserror::Error)]
pub enum VirtualDriveError {
    #[error("failed to create virtual drive {0}")]
    Create(PathBuf, #[source] std::io::Error),

    #[error("failed to create tar archive for virtual drive from {0}")]
    CreateTar(PathBuf, #[source] std::io::Error),

    #[error("failed to open virtual drive {0}")]
    Open(PathBuf, #[source] std::io::Error),

    #[error("failed to list files in virtual drive {0}")]
    List(PathBuf, #[source] std::io::Error),

    #[error("failed to create directory {0}")]
    Extract(PathBuf, #[source] std::io::Error),

    #[error("failed to extract {0} to {1}")]
    ExtractEntry(PathBuf, PathBuf, #[source] std::io::Error),

    #[error(transparent)]
    Util(#[from] crate::util::UtilError),

    #[error("failed to get length of file {0}")]
    Metadata(PathBuf, #[source] std::io::Error),

    #[error("failed to set length of file to {0}: {1}")]
    SetLen(u64, PathBuf, #[source] std::io::Error),

    #[error("virtual drive is too big: {0} > {1}")]
    DriveTooBig(u64, u64),

    /// Can't run tar.
    #[error("failed to run system tar command: {0}: {1}")]
    Tar(&'static str, PathBuf, #[source] std::io::Error),

    /// Tar failed.
    #[error("failed to run system tar command: {0}: {1}")]
    TarFailed(&'static str, PathBuf, i32, String),
}