From b33dfac11794ee2f41236395587325d5bd017e4b Mon Sep 17 00:00:00 2001 From: Edward Shen Date: Wed, 24 Aug 2022 00:58:20 -0700 Subject: [PATCH] Add progress bar to CLI --- Cargo.lock | 55 ++++++++++++++++++++++++++++++++++++++++++++++ cli/Cargo.toml | 2 ++ cli/src/main.rs | 58 ++++++++++++++++++++++++++++++++++++++++++++----- 3 files changed, 110 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 79b708d..1bb1db7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -351,6 +351,20 @@ dependencies = [ "os_str_bytes", ] +[[package]] +name = "console" +version = "0.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89eab4d20ce20cea182308bca13088fecea9c05f6776cf287205d41a0ed3c847" +dependencies = [ + "encode_unicode", + "libc", + "once_cell", + "terminal_size", + "unicode-width", + "winapi", +] + [[package]] name = "console_error_panic_hook" version = "0.1.7" @@ -426,6 +440,12 @@ dependencies = [ "subtle", ] +[[package]] +name = "encode_unicode" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a357d28ed41a50f9c765dbfe56cbc04a64e53e5fc58ba79fbc34c10ef3df831f" + [[package]] name = "encoding_rs" version = "0.8.31" @@ -840,6 +860,17 @@ dependencies = [ "hashbrown", ] +[[package]] +name = "indicatif" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fcc42b206e70d86ec03285b123e65a5458c92027d1fb2ae3555878b8113b3ddf" +dependencies = [ + "console", + "number_prefix", + "unicode-width", +] + [[package]] name = "instant" version = "0.1.12" @@ -1060,13 +1091,21 @@ dependencies = [ "libc", ] +[[package]] +name = "number_prefix" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" + [[package]] name = "omegaupload" version = "0.1.1" dependencies = [ "anyhow", "atty", + "bytes", "clap", + "indicatif", "omegaupload-common", "reqwest", "rpassword", @@ -1764,6 +1803,16 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "terminal_size" +version = "0.1.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "633c1a546cee861a1a6d0dc69ebeca693bf4296661ba7852b9d21d159e0506df" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "textwrap" version = "0.15.0" @@ -2065,6 +2114,12 @@ dependencies = [ "tinyvec", ] +[[package]] +name = "unicode-width" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ed742d4ea2bd1176e236172c8429aaf54486e7ac098db29ffe6529e0ce50973" + [[package]] name = "universal-hash" version = "0.4.1" diff --git a/cli/Cargo.toml b/cli/Cargo.toml index 157e672..03b4da2 100644 --- a/cli/Cargo.toml +++ b/cli/Cargo.toml @@ -12,6 +12,8 @@ license = "GPL-3.0-or-later" omegaupload-common = { path = "../common" } anyhow = "1.0.58" atty = "0.2.14" +bytes = "1" clap = { version = "3.2.15", features = ["derive"] } +indicatif = "0.17" reqwest = { version = "0.11.11", default-features = false, features = ["rustls-tls", "blocking"] } rpassword = "7.0.0" diff --git a/cli/src/main.rs b/cli/src/main.rs index cdac41a..8cfbf73 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -17,19 +17,21 @@ // You should have received a copy of the GNU General Public License // along with this program. If not, see . -use std::io::{Read, Write}; +use std::io::{Cursor, Read, Write}; use std::path::PathBuf; use anyhow::{anyhow, bail, Context, Result}; use atty::Stream; +use bytes::Bytes; use clap::Parser; +use indicatif::{ProgressBar, ProgressStyle}; use omegaupload_common::crypto::{open_in_place, seal_in_place}; use omegaupload_common::fragment::Builder; use omegaupload_common::secrecy::{ExposeSecret, SecretString, SecretVec}; use omegaupload_common::{ base64, Expiration, ParsedUrl, Url, API_ENDPOINT, EXPIRATION_HEADER_NAME, }; -use reqwest::blocking::Client; +use reqwest::blocking::{Body, Client}; use reqwest::header::EXPIRES; use reqwest::StatusCode; use rpassword::prompt_password; @@ -128,13 +130,35 @@ fn handle_upload( (container, key) }; - let mut res = Client::new().post(url.as_ref()); + let mut req = Client::new().post(url.as_ref()); if let Some(duration) = duration { - res = res.header(&*EXPIRATION_HEADER_NAME, duration); + req = req.header(&*EXPIRATION_HEADER_NAME, duration); } - let res = res.body(data).send().context("Request to server failed")?; + let data_size = data.len() as u64; + let progress_style = ProgressStyle::with_template( + "[{elapsed_precise}] {bar:40} {bytes}/{total_bytes} {eta_precise}", + ) + .unwrap(); + let progress_bar = ProgressBar::new(data_size).with_style(progress_style); + let res = req + .body(Body::sized( + WrappedBody::new( + move |amt| { + progress_bar.inc(amt as u64); + }, + data, + ), + data_size, + )) + .build() + .expect("Failed to build body"); + let res = reqwest::blocking::ClientBuilder::new() + .timeout(None) + .build()? + .execute(res) + .context("Request to server failed")?; if res.status() != StatusCode::OK { bail!("Upload failed. Got HTTP error {}", res.status()); @@ -170,6 +194,30 @@ fn handle_upload( Ok(()) } +struct WrappedBody { + callback: Callback, + inner: Cursor, +} + +impl WrappedBody { + fn new(callback: Callback, data: Vec) -> Self { + Self { + callback, + inner: Cursor::new(Bytes::from(data)), + } + } +} + +impl Read for WrappedBody { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + let res = self.inner.read(buf); + if let Ok(size) = res { + (self.callback)(size); + } + res + } +} + fn handle_download(mut url: ParsedUrl) -> Result<()> { url.sanitized_url .set_path(&format!("{API_ENDPOINT}{}", url.sanitized_url.path()));