Download compressed gz image from remote source

bmap-rs/Cargo.toml

@@ -11,5 +11,10 @@ bmap = { path = "../bmap" }
 anyhow = "1.0.66"
 nix = "0.26.1"
 flate2 = "1.0.24"
-clap = { version = "4.0.18", features = ["derive"] }
-indicatif = "0.17.1"
+clap = { version = "4.0.18", features = ["cargo"] }
+indicatif = { version = "0.17.1", features = ["tokio"] }
+async-compression = { version = "0.3.15", features = ["gzip", "futures-io"] }
+tokio = { version = "1.21.2", features = ["rt", "macros", "fs", "rt-multi-thread"] }
+reqwest = { version = "0.11.12", features = ["stream"] }
+tokio-util = { version = "0.7.4", features = ["compat"] }
+futures = "0.3.25"

bmap-rs/src/main.rs

@@ -1,32 +1,75 @@
-use anyhow::{anyhow, bail, Context, Result};
-use bmap::{Bmap, Discarder, SeekForward};
-use clap::Parser;
+use anyhow::{anyhow, bail, ensure, Context, Result};
+use async_compression::futures::bufread::GzipDecoder;
+use bmap::{AsyncDiscarder, Bmap, Discarder, SeekForward};
+use clap::{arg, command, Command};
 use flate2::read::GzDecoder;
+use futures::TryStreamExt;
 use indicatif::{ProgressBar, ProgressState, ProgressStyle};
 use nix::unistd::ftruncate;
+use reqwest::{Response, Url};
 use std::ffi::OsStr;
 use std::fmt::Write;
 use std::fs::File;
 use std::io::Read;
 use std::os::unix::io::AsRawFd;
 use std::path::{Path, PathBuf};
+use tokio_util::compat::TokioAsyncReadCompatExt;
 
-#[derive(Parser, Debug)]
+#[derive(Debug)]
+enum Image {
+    Path(PathBuf),
+    Url(Url),
+}
+
+#[derive(Debug)]
 struct Copy {
-    image: PathBuf,
+    image: Image,
     dest: PathBuf,
 }
 
-#[derive(Parser, Debug)]
+#[derive(Debug)]
 
-enum Command {
+enum Subcommand {
     Copy(Copy),
 }
 
-#[derive(Parser, Debug)]
+#[derive(Debug)]
 struct Opts {
-    #[command(subcommand)]
-    command: Command,
+    command: Subcommand,
+}
+
+impl Opts {
+    fn parser() -> Opts {
+        let matches = command!()
+            .propagate_version(true)
+            .subcommand_required(true)
+            .arg_required_else_help(true)
+            .subcommand(
+                Command::new("copy")
+                    .about("Copy image to block device or file")
+                    .arg(arg!([IMAGE]).required(true))
+                    .arg(arg!([DESTINATION]).required(true)),
+            )
+            .get_matches();
+        match matches.subcommand() {
+            Some(("copy", sub_matches)) => Opts {
+                command: Subcommand::Copy({
+                    Copy {
+                        image: match Url::parse(sub_matches.get_one::<String>("IMAGE").unwrap()) {
+                            Ok(url) => Image::Url(url),
+                            Err(_) => Image::Path(PathBuf::from(
+                                sub_matches.get_one::<String>("IMAGE").unwrap(),
+                            )),
+                        },
+                        dest: PathBuf::from(sub_matches.get_one::<String>("DESTINATION").unwrap()),
+                    }
+                }),
+            },
+            _ => unreachable!(
+                "Exhausted list of subcommands and subcommand_required prevents `None`"
+            ),
+        }
+    }
 }
 
 fn append(path: PathBuf) -> PathBuf {
@@ -51,6 +94,13 @@ fn find_bmap(img: &Path) -> Option<PathBuf> {
     }
 }
 
+fn find_remote_bmap(mut url: Url) -> Result<Url> {
+    let mut path = PathBuf::from(url.path());
+    path.set_extension("bmap");
+    url.set_path(path.to_str().unwrap());
+    Ok(url)
+}
+
 trait ReadSeekForward: SeekForward + Read {}
 impl<T: Read + SeekForward> ReadSeekForward for T {}
 
@@ -78,7 +128,7 @@ impl SeekForward for Decoder {
     }
 }
 
-fn setup_input(path: &Path) -> Result<Decoder> {
+fn setup_local_input(path: &Path) -> Result<Decoder> {
     let f = File::open(path)?;
     match path.extension().and_then(OsStr::to_str) {
         Some("gz") => {
@@ -89,12 +139,44 @@ fn setup_input(path: &Path) -> Result<Decoder> {
     }
 }
 
-fn copy(c: Copy) -> Result<()> {
-    if !c.image.exists() {
-        bail!("Image file doesn't exist")
+async fn setup_remote_input(url: Url) -> Result<Response> {
+    match PathBuf::from(url.path())
+        .extension()
+        .and_then(OsStr::to_str)
+    {
+        Some("gz") => reqwest::get(url).await.map_err(anyhow::Error::new),
+        None => bail!("No file extension found"),
+        _ => bail!("Image file format not implemented"),
     }
+}
 
-    let bmap = find_bmap(&c.image).ok_or_else(|| anyhow!("Couldn't find bmap file"))?;
+fn setup_progress_bar(bmap: &Bmap) -> ProgressBar {
+    let pb = ProgressBar::new(bmap.total_mapped_size());
+    pb.set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
+        .unwrap()
+        .with_key("eta", |state: &ProgressState, w: &mut dyn Write| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap())
+        .progress_chars("#>-"));
+    pb
+}
+
+fn setup_output<T: AsRawFd>(output: &T, bmap: &Bmap, metadata: std::fs::Metadata) -> Result<()> {
+    if metadata.is_file() {
+        ftruncate(output.as_raw_fd(), bmap.image_size() as i64)
+            .context("Failed to truncate file")?;
+    }
+    Ok(())
+}
+
+async fn copy(c: Copy) -> Result<()> {
+    match c.image {
+        Image::Path(path) => copy_local_input(path, c.dest),
+        Image::Url(url) => copy_remote_input(url, c.dest).await,
+    }
+}
+
+fn copy_local_input(source: PathBuf, destination: PathBuf) -> Result<()> {
+    ensure!(source.exists(), "Image file doesn't exist");
+    let bmap = find_bmap(&source).ok_or_else(|| anyhow!("Couldn't find bmap file"))?;
     println!("Found bmap file: {}", bmap.display());
 
     let mut b = File::open(&bmap).context("Failed to open bmap file")?;
@@ -105,32 +187,63 @@ fn copy(c: Copy) -> Result<()> {
     let output = std::fs::OpenOptions::new()
         .write(true)
         .create(true)
-        .open(c.dest)?;
+        .open(destination)?;
 
-    if output.metadata()?.is_file() {
-        ftruncate(output.as_raw_fd(), bmap.image_size() as i64)
-            .context("Failed to truncate file")?;
-    }
+    setup_output(&output, &bmap, output.metadata()?)?;
 
-    let mut input = setup_input(&c.image)?;
-    let pb = ProgressBar::new(bmap.total_mapped_size());
-    pb.set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
-        .unwrap()
-        .with_key("eta", |state: &ProgressState, w: &mut dyn Write| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap())
-        .progress_chars("#>-"));
+    let mut input = setup_local_input(&source)?;
+    let pb = setup_progress_bar(&bmap);
     bmap::copy(&mut input, &mut pb.wrap_write(&output), &bmap)?;
     pb.finish_and_clear();
 
     println!("Done: Syncing...");
-    output.sync_all().expect("Sync failure");
+    output.sync_all()?;
+
+    Ok(())
+}
+
+async fn copy_remote_input(source: Url, destination: PathBuf) -> Result<()> {
+    let bmap_url = find_remote_bmap(source.clone())?;
+
+    let xml = reqwest::get(bmap_url.clone()).await?.text().await?;
+    println!("Found bmap file: {}", bmap_url);
+
+    let bmap = Bmap::from_xml(&xml)?;
+    let mut output = tokio::fs::OpenOptions::new()
+        .write(true)
+        .create(true)
+        .open(destination)
+        .await?;
+
+    setup_output(&output, &bmap, output.metadata().await?)?;
+
+    let res = setup_remote_input(source).await?;
+    let stream = res
+        .bytes_stream()
+        .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))
+        .into_async_read();
+    let reader = GzipDecoder::new(stream);
+    let mut input = AsyncDiscarder::new(reader);
+    let pb = setup_progress_bar(&bmap);
+    bmap::copy_async(
+        &mut input,
+        &mut pb.wrap_async_write(&mut output).compat(),
+        &bmap,
+    )
+    .await?;
+    pb.finish_and_clear();
+
+    println!("Done: Syncing...");
+    output.sync_all().await?;
 
     Ok(())
 }
 
-fn main() -> Result<()> {
-    let opts = Opts::parse();
+#[tokio::main]
+async fn main() -> Result<()> {
+    let opts = Opts::parser();
 
     match opts.command {
-        Command::Copy(c) => copy(c),
+        Subcommand::Copy(c) => copy(c).await,
     }
 }

bmap/Cargo.toml

@@ -15,3 +15,5 @@ sha2 = { version = "0.10.6", features = [ "asm" ] }
 strum = { version = "0.24.1", features = [ "derive"] }
 digest = "0.10.5"
 flate2 = "1.0.20"
+async-trait = "0.1.58"
+futures = "0.3.25"

bmap/src/discarder.rs

@@ -1,6 +1,10 @@
-use crate::SeekForward;
+use crate::{AsyncSeekForward, SeekForward};
+use async_trait::async_trait;
+use futures::io::{AsyncRead, AsyncReadExt};
 use std::io::Read;
 use std::io::Result as IOResult;
+use std::pin::Pin;
+use std::task::{Context, Poll};
 
 /// Adaptor that implements SeekForward on types only implementing Read by discarding data
 pub struct Discarder<R: Read> {
@@ -36,6 +40,44 @@ impl<R: Read> SeekForward for Discarder<R> {
     }
 }
 
+pub struct AsyncDiscarder<R: AsyncRead> {
+    reader: R,
+}
+
+impl<R: AsyncRead> AsyncDiscarder<R> {
+    pub fn new(reader: R) -> Self {
+        Self { reader }
+    }
+
+    pub fn into_inner(self) -> R {
+        self.reader
+    }
+}
+
+impl<R: AsyncRead + Unpin> AsyncRead for AsyncDiscarder<R> {
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut [u8],
+    ) -> Poll<IOResult<usize>> {
+        Pin::new(&mut self.reader).poll_read(cx, buf)
+    }
+}
+
+#[async_trait(?Send)]
+impl<R: AsyncRead + Unpin> AsyncSeekForward for AsyncDiscarder<R> {
+    async fn async_seek_forward(&mut self, forward: u64) -> IOResult<()> {
+        let mut buf = [0; 4096];
+        let mut left = forward as usize;
+        while left > 0 {
+            let toread = left.min(buf.len());
+            let r = self.read(&mut buf[0..toread]).await?;
+            left -= r;
+        }
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;