avml 0.18.0

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

use avml::{Result, Snapshot, Source, iomem};
use clap::Parser;
#[cfg(feature = "blobstore")]
use core::num::NonZeroUsize;
use core::{num::NonZeroU64, ops::Range};
use std::path::PathBuf;
#[cfg(any(feature = "blobstore", feature = "put"))]
use {avml::Error, tokio::fs::remove_file, url::Url};

#[derive(Parser)]
/// A portable volatile memory acquisition tool for Linux
#[command(author, version, about, long_about = None)]
struct Config {
    /// compress via snappy
    #[arg(long)]
    compress: bool,

    /// specify input source
    #[arg(long, value_enum)]
    source: Option<Source>,

    /// Specify the maximum estimated disk usage (in MB)
    #[arg(long)]
    max_disk_usage: Option<NonZeroU64>,

    /// Specify the maximum estimated disk usage to stay under
    #[arg(long, value_parser = disk_usage_percentage)]
    max_disk_usage_percentage: Option<f64>,

    /// upload via HTTP PUT upon acquisition
    #[cfg(feature = "put")]
    #[arg(long)]
    url: Option<Url>,

    /// delete upon successful upload
    #[cfg(any(feature = "blobstore", feature = "put"))]
    #[arg(long)]
    delete: bool,

    /// upload via Azure Blob Store upon acquisition
    #[cfg(feature = "blobstore")]
    #[arg(long)]
    sas_url: Option<Url>,

    /// specify maximum block size in MiB; must be greater than 0
    #[cfg(feature = "blobstore")]
    #[arg(long)]
    sas_block_size: Option<NonZeroU64>,

    /// specify blob upload concurrency; must be greater than 0
    #[cfg(feature = "blobstore")]
    #[arg(long)]
    sas_block_concurrency: Option<NonZeroUsize>,

    /// name of the file to write to on local system
    filename: PathBuf,
}

const PERCENTAGE: Range<f64> = 0.01..100.0;

fn disk_usage_percentage(s: &str) -> core::result::Result<f64, String> {
    let value = s
        .parse()
        .map_err(|_| format!("`{s}` isn't a valid value"))?;
    if PERCENTAGE.contains(&value) {
        Ok(value)
    } else {
        Err(format!(
            "value is not a valid percentage in range {}-{}",
            PERCENTAGE.start, PERCENTAGE.end
        ))
    }
}

#[cfg(any(feature = "blobstore", feature = "put"))]
async fn upload(config: &Config) -> Result<()> {
    let mut delete = false;

    #[cfg(feature = "put")]
    {
        if let Some(ref url) = config.url {
            avml::put(&config.filename, url).await?;
            delete = true;
        }
    }

    #[cfg(feature = "blobstore")]
    {
        if let Some(ref sas_url) = config.sas_url {
            let uploader = avml::BlobUploader::new(sas_url)?
                .block_size(config.sas_block_size)
                .concurrency(config.sas_block_concurrency);
            uploader.upload_file(&config.filename).await?;
            delete = true;
        }
    }

    if delete && config.delete {
        remove_file(&config.filename)
            .await
            .map_err(|source| Error::Io {
                context: "unable to remove snapshot",
                source,
            })?;
    }

    Ok(())
}

fn acquire(config: &Config) -> Result<()> {
    let version = if config.compress { 2 } else { 1 };

    let ranges = iomem::parse()?;
    let snapshot = Snapshot::new(&config.filename, ranges)
        .source(config.source.as_ref())
        .max_disk_usage_percentage(config.max_disk_usage_percentage)
        .max_disk_usage(config.max_disk_usage)
        .version(version);
    snapshot.create()?;
    Ok(())
}

#[cfg(not(any(feature = "blobstore", feature = "put")))]
fn main() -> Result<()> {
    let config = Config::parse();
    acquire(&config)
}

#[cfg(any(feature = "blobstore", feature = "put"))]
#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<()> {
    let config = Config::parse();
    acquire(&config)?;
    upload(&config).await
}