VirtWeb/virtweb_backend/src/utils/file_disks_utils.rs

use crate::app_config::AppConfig;
use crate::constants;
use crate::utils::file_size_utils::FileSize;
use std::fs::File;
use std::os::linux::fs::MetadataExt;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::time::UNIX_EPOCH;

#[derive(thiserror::Error, Debug)]
enum DisksError {
    #[error("DiskParseError: {0}")]
    Parse(&'static str),
    #[error("DiskCreateError")]
    Create,
    #[error("DiskConvertError: {0}")]
    Convert(String),
}

#[derive(Debug, serde::Serialize, serde::Deserialize, Copy, Clone, PartialEq, Eq)]
#[serde(tag = "format")]
pub enum DiskFileFormat {
    Raw {
        #[serde(default)]
        is_sparse: bool,
    },
    QCow2 {
        #[serde(default)]
        virtual_size: FileSize,
    },
    GzCompressedRaw,
    GzCompressedQCow2,
    XzCompressedRaw,
    XzCompressedQCow2,
}

impl DiskFileFormat {
    pub fn ext(&self) -> &'static [&'static str] {
        match self {
            DiskFileFormat::Raw { .. } => &["raw", ""],
            DiskFileFormat::QCow2 { .. } => &["qcow2"],
            DiskFileFormat::GzCompressedRaw => &["raw.gz"],
            DiskFileFormat::GzCompressedQCow2 => &["qcow2.gz"],
            DiskFileFormat::XzCompressedRaw => &["raw.xz"],
            DiskFileFormat::XzCompressedQCow2 => &["qcow2.xz"],
        }
    }
}

/// Disk file information
#[derive(serde::Serialize)]
pub struct DiskFileInfo {
    pub file_path: PathBuf,
    pub file_size: FileSize,
    #[serde(flatten)]
    pub format: DiskFileFormat,
    pub file_name: String,
    pub name: String,
    pub created: u64,
}

impl DiskFileInfo {
    /// Get disk image file information
    pub fn load_file(file: &Path) -> anyhow::Result<Self> {
        if !file.is_file() {
            return Err(DisksError::Parse("Path is not a file!").into());
        }

        // Get file metadata
        let metadata = file.metadata()?;
        let mut name = file
            .file_stem()
            .and_then(|s| s.to_str())
            .unwrap_or("disk")
            .to_string();
        let ext = file.extension().and_then(|s| s.to_str()).unwrap_or("raw");

        // Determine file format
        let format = match ext {
            "qcow2" => DiskFileFormat::QCow2 {
                virtual_size: qcow_virt_size(file)?,
            },
            "raw" => DiskFileFormat::Raw {
                is_sparse: metadata.len() / 512 >= metadata.st_blocks(),
            },
            "gz" if name.ends_with(".qcow2") => {
                name = name.strip_suffix(".qcow2").unwrap_or(&name).to_string();
                DiskFileFormat::GzCompressedQCow2
            }
            "gz" => DiskFileFormat::GzCompressedRaw,
            "xz" if name.ends_with(".qcow2") => {
                name = name.strip_suffix(".qcow2").unwrap_or(&name).to_string();
                DiskFileFormat::XzCompressedQCow2
            }
            "xz" => DiskFileFormat::XzCompressedRaw,
            _ => anyhow::bail!("Unsupported disk extension: {ext}!"),
        };

        Ok(Self {
            file_path: file.to_path_buf(),
            name,
            file_size: FileSize::from_bytes(metadata.len() as usize),
            format,
            file_name: file
                .file_name()
                .and_then(|s| s.to_str())
                .unwrap_or("")
                .to_string(),
            created: metadata
                .created()?
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_secs(),
        })
    }

    /// Create a new empty disk
    pub fn create(file: &Path, format: DiskFileFormat, size: FileSize) -> anyhow::Result<()> {
        // Prepare command to create file
        let res = match format {
            DiskFileFormat::Raw { is_sparse } => {
                let mut cmd = Command::new("/usr/bin/dd");
                cmd.arg("if=/dev/zero")
                    .arg(format!("of={}", file.to_string_lossy()))
                    .arg("bs=1M");

                match is_sparse {
                    false => cmd.arg(format!("count={}", size.as_mb())),
                    true => cmd.arg(format!("seek={}", size.as_mb())).arg("count=0"),
                };

                cmd.output()?
            }

            DiskFileFormat::QCow2 { virtual_size } => {
                let mut cmd = Command::new(constants::PROGRAM_QEMU_IMAGE);
                cmd.arg("create")
                    .arg("-f")
                    .arg("qcow2")
                    .arg(file)
                    .arg(format!("{}M", virtual_size.as_mb()));

                cmd.output()?
            }
            _ => anyhow::bail!("Cannot create disk file image of this format: {format:?}!"),
        };

        // Execute Linux command
        if !res.status.success() {
            log::error!(
                "Failed to create disk! stderr={} stdout={}",
                String::from_utf8_lossy(&res.stderr),
                String::from_utf8_lossy(&res.stdout)
            );
            return Err(DisksError::Create.into());
        }

        Ok(())
    }

    /// Copy / convert file disk image into a new destination with optionally a new file format
    pub fn convert(&self, dest_file: &Path, dest_format: DiskFileFormat) -> anyhow::Result<()> {
        // Create a temporary directory to perform the operation
        let temp_dir = tempfile::tempdir_in(&AppConfig::get().temp_dir)?;
        let temp_file = temp_dir
            .path()
            .join(format!("temp_file.{}", dest_format.ext()[0]));

        // Prepare the conversion
        let mut cmd = match (self.format, dest_format) {
            // Decompress QCow2 (GZIP)
            (DiskFileFormat::GzCompressedQCow2, DiskFileFormat::QCow2 { .. }) => {
                let mut cmd = Command::new(constants::PROGRAM_GZIP);
                cmd.arg("--keep")
                    .arg("--decompress")
                    .arg("--to-stdout")
                    .arg(&self.file_path)
                    .stdout(File::create(&temp_file)?);
                cmd
            }

            // Decompress QCow2 (XZ)
            (DiskFileFormat::XzCompressedQCow2, DiskFileFormat::QCow2 { .. }) => {
                let mut cmd = Command::new(constants::PROGRAM_XZ);
                cmd.arg("--stdout")
                    .arg("--keep")
                    .arg("--decompress")
                    .arg(&self.file_path)
                    .stdout(File::create(&temp_file)?);
                cmd
            }

            // Compress QCow2 (Gzip)
            (DiskFileFormat::QCow2 { .. }, DiskFileFormat::GzCompressedQCow2) => {
                let mut cmd = Command::new(constants::PROGRAM_GZIP);
                cmd.arg("--keep")
                    .arg("--to-stdout")
                    .arg(&self.file_path)
                    .stdout(File::create(&temp_file)?);
                cmd
            }

            // Compress QCow2 (Xz)
            (DiskFileFormat::QCow2 { .. }, DiskFileFormat::XzCompressedQCow2) => {
                let mut cmd = Command::new(constants::PROGRAM_XZ);
                cmd.arg("--keep")
                    .arg("--to-stdout")
                    .arg(&self.file_path)
                    .stdout(File::create(&temp_file)?);
                cmd
            }

            // Convert QCow2 to Raw file
            (DiskFileFormat::QCow2 { .. }, DiskFileFormat::Raw { is_sparse }) => {
                let mut cmd = Command::new(constants::PROGRAM_QEMU_IMAGE);
                cmd.arg("convert")
                    .arg("-f")
                    .arg("qcow2")
                    .arg("-O")
                    .arg("raw")
                    .arg(&self.file_path)
                    .arg(&temp_file);

                if !is_sparse {
                    cmd.args(["-S", "0"]);
                }

                cmd
            }

            // Clone a QCow file, using qemu-image instead of cp might improve "sparsification" of
            // file
            (DiskFileFormat::QCow2 { .. }, DiskFileFormat::QCow2 { .. }) => {
                let mut cmd = Command::new(constants::PROGRAM_QEMU_IMAGE);
                cmd.arg("convert")
                    .arg("-f")
                    .arg("qcow2")
                    .arg("-O")
                    .arg("qcow2")
                    .arg(&self.file_path)
                    .arg(&temp_file);
                cmd
            }

            // Convert Raw to QCow2 file
            (DiskFileFormat::Raw { .. }, DiskFileFormat::QCow2 { .. }) => {
                let mut cmd = Command::new(constants::PROGRAM_QEMU_IMAGE);
                cmd.arg("convert")
                    .arg("-f")
                    .arg("raw")
                    .arg("-O")
                    .arg("qcow2")
                    .arg(&self.file_path)
                    .arg(&temp_file);

                cmd
            }

            // Render raw file non sparse
            (DiskFileFormat::Raw { is_sparse: true }, DiskFileFormat::Raw { is_sparse: false }) => {
                let mut cmd = Command::new(constants::PROGRAM_COPY);
                cmd.arg("--sparse=never")
                    .arg(&self.file_path)
                    .arg(&temp_file);
                cmd
            }

            // Render raw file sparse
            (DiskFileFormat::Raw { is_sparse: false }, DiskFileFormat::Raw { is_sparse: true }) => {
                let mut cmd = Command::new(constants::PROGRAM_DD);
                cmd.arg("conv=sparse")
                    .arg(format!("if={}", self.file_path.display()))
                    .arg(format!("of={}", temp_file.display()));
                cmd
            }

            // Compress Raw (Gz)
            (DiskFileFormat::Raw { .. }, DiskFileFormat::GzCompressedRaw) => {
                let mut cmd = Command::new(constants::PROGRAM_GZIP);
                cmd.arg("--keep")
                    .arg("--to-stdout")
                    .arg(&self.file_path)
                    .stdout(File::create(&temp_file)?);
                cmd
            }

            // Compress Raw (Xz)
            (DiskFileFormat::Raw { .. }, DiskFileFormat::XzCompressedRaw) => {
                let mut cmd = Command::new(constants::PROGRAM_XZ);
                cmd.arg("--keep")
                    .arg("--to-stdout")
                    .arg(&self.file_path)
                    .stdout(File::create(&temp_file)?);
                cmd
            }

            // Decompress Raw (Gz) to not sparse file
            (DiskFileFormat::GzCompressedRaw, DiskFileFormat::Raw { is_sparse: false }) => {
                let mut cmd = Command::new(constants::PROGRAM_GZIP);
                cmd.arg("--keep")
                    .arg("--decompress")
                    .arg("--to-stdout")
                    .arg(&self.file_path)
                    .stdout(File::create(&temp_file)?);
                cmd
            }
            // Decompress Raw (Xz) to not sparse file
            (DiskFileFormat::XzCompressedRaw, DiskFileFormat::Raw { is_sparse: false }) => {
                let mut cmd = Command::new(constants::PROGRAM_XZ);
                cmd.arg("--keep")
                    .arg("--decompress")
                    .arg("--to-stdout")
                    .arg(&self.file_path)
                    .stdout(File::create(&temp_file)?);
                cmd
            }

            // Decompress Raw (Gz) to sparse file
            // https://benou.fr/www/ben/decompressing-sparse-files.html
            (DiskFileFormat::GzCompressedRaw, DiskFileFormat::Raw { is_sparse: true }) => {
                let mut cmd = Command::new(constants::PROGRAM_BASH);
                cmd.arg("-c").arg(format!(
                    "{} --decompress --to-stdout {} | {} conv=sparse of={}",
                    constants::PROGRAM_GZIP,
                    self.file_path.display(),
                    constants::PROGRAM_DD,
                    temp_file.display()
                ));
                cmd
            }

            // Decompress Raw (XZ) to sparse file
            // https://benou.fr/www/ben/decompressing-sparse-files.html
            (DiskFileFormat::XzCompressedRaw, DiskFileFormat::Raw { is_sparse: true }) => {
                let mut cmd = Command::new(constants::PROGRAM_BASH);
                cmd.arg("-c").arg(format!(
                    "{} --decompress --to-stdout {} | {} conv=sparse of={}",
                    constants::PROGRAM_XZ,
                    self.file_path.display(),
                    constants::PROGRAM_DD,
                    temp_file.display()
                ));
                cmd
            }

            // Dumb copy of file
            (a, b) if a == b => {
                let mut cmd = Command::new(constants::PROGRAM_COPY);
                cmd.arg("--sparse=auto")
                    .arg(&self.file_path)
                    .arg(&temp_file);
                cmd
            }

            // By default, conversion is unsupported
            (src, dest) => {
                return Err(DisksError::Convert(format!(
                    "Conversion from {src:?} to {dest:?} is not supported!"
                ))
                .into());
            }
        };

        // Execute the conversion
        let command_s = format!(
            "{} {}",
            cmd.get_program().display(),
            cmd.get_args()
                .map(|a| format!("'{}'", a.display()))
                .collect::<Vec<String>>()
                .join(" ")
        );
        let cmd_output = cmd.output()?;
        if !cmd_output.status.success() {
            return Err(DisksError::Convert(format!(
                "Command failed:\n{command_s}\nStatus: {}\nstdout: {}\nstderr: {}",
                cmd_output.status,
                String::from_utf8_lossy(&cmd_output.stdout),
                String::from_utf8_lossy(&cmd_output.stderr)
            ))
            .into());
        }

        // Check the file was created
        if !temp_file.is_file() {
            return Err(DisksError::Convert(
                "Temporary was not created after execution of command!".to_string(),
            )
            .into());
        }

        // Move the file to its final location
        std::fs::rename(temp_file, dest_file)?;

        Ok(())
    }

    /// Get disk virtual size, if available
    pub fn virtual_size(&self) -> Option<FileSize> {
        match self.format {
            DiskFileFormat::Raw { .. } => Some(self.file_size),
            DiskFileFormat::QCow2 { virtual_size } => Some(virtual_size),
            _ => None,
        }
    }

    /// Resize disk
    pub fn resize(&self, new_size: FileSize) -> anyhow::Result<()> {
        if new_size <= self.virtual_size().unwrap_or(new_size) {
            anyhow::bail!("Shrinking disk image file is not supported!");
        }

        let mut cmd = Command::new(constants::PROGRAM_QEMU_IMAGE);
        cmd.arg("resize")
            .arg("-f")
            .arg(match self.format {
                DiskFileFormat::QCow2 { .. } => "qcow2",
                DiskFileFormat::Raw { .. } => "raw",
                f => anyhow::bail!("Unsupported disk format for resize: {f:?}"),
            })
            .arg(&self.file_path)
            .arg(new_size.as_bytes().to_string());

        let output = cmd.output()?;
        if !output.status.success() {
            anyhow::bail!(
                "{} info failed, status: {}, stderr: {}",
                constants::PROGRAM_QEMU_IMAGE,
                output.status,
                String::from_utf8_lossy(&output.stderr)
            );
        }
        Ok(())
    }
}

#[derive(serde::Deserialize)]
struct QCowInfoOutput {
    #[serde(rename = "virtual-size")]
    virtual_size: usize,
}

/// Get QCow2 virtual size
fn qcow_virt_size(path: &Path) -> anyhow::Result<FileSize> {
    // Run qemu-img
    let mut cmd = Command::new(constants::PROGRAM_QEMU_IMAGE);
    cmd.args([
        "info",
        path.to_str().unwrap_or(""),
        "--output",
        "json",
        "--force-share",
    ]);
    let output = cmd.output()?;
    if !output.status.success() {
        anyhow::bail!(
            "{} info failed, status: {}, stderr: {}",
            constants::PROGRAM_QEMU_IMAGE,
            output.status,
            String::from_utf8_lossy(&output.stderr)
        );
    }
    let res_json = String::from_utf8(output.stdout)?;

    // Decode JSON
    let decoded: QCowInfoOutput = serde_json::from_str(&res_json)?;
    Ok(FileSize::from_bytes(decoded.virtual_size))
}