Pārlūkot izejas kodu

Implement progress bar both for upload and download

timvisee 7 gadi atpakaļ
vecāks
revīzija
d382a015c5

+ 20 - 5
api/src/action/download.rs

@@ -1,5 +1,6 @@
 use std::fs::File;
 use std::io;
+use std::sync::{Arc, Mutex};
 
 use openssl::symm::decrypt_aead;
 use reqwest::{
@@ -15,7 +16,7 @@ use crypto::key_set::KeySet;
 use crypto::sign::signature_encoded;
 use file::file::DownloadFile;
 use file::metadata::Metadata;
-use reader::EncryptedFileWriter;
+use reader::{EncryptedFileWriter, ProgressReporter, ProgressWriter};
 
 pub type Result<T> = ::std::result::Result<T, DownloadError>;
 
@@ -42,6 +43,7 @@ impl<'a> Download<'a> {
     pub fn invoke(
         self,
         client: &Client,
+        reporter: Arc<Mutex<ProgressReporter>>,
     ) -> Result<()> {
         // Create a key set for the file
         let mut key = KeySet::from(self.file);
@@ -166,25 +168,38 @@ impl<'a> Download<'a> {
             .expect("failed to fetch file, missing content length header")
             .0;
 
-        // Open a file to write to
+        // Open a file to write to, and build an encrypted writer
         // TODO: this should become a temporary file first
-        let out = File::create("downloaded.toml")
+        let out = File::create("downloaded.zip")
             .expect("failed to open file");
-        let mut writer = EncryptedFileWriter::new(
+        let writer = EncryptedFileWriter::new(
             out,
             response_len as usize,
             KeySet::cipher(),
             key.file_key().unwrap(),
             key.iv(),
         ).expect("failed to create encrypted writer");
+        let mut writer = ProgressWriter::new(writer)
+            .expect("failed to create encrypted writer");
+        writer.set_reporter(reporter.clone());
+
+        // Start the writer
+        reporter.lock()
+            .expect("unable to start progress, failed to get lock")
+            .start(response_len);
 
         // Write to the output file
         io::copy(&mut response, &mut writer)
             .expect("failed to download and decrypt file");
 
+        // Finish
+        reporter.lock()
+            .expect("unable to finish progress, failed to get lock")
+            .finish();
+
         // Verify the writer
         // TODO: delete the file if verification failed, show a proper error
-        assert!(writer.verified(), "downloaded and decrypted file could not be verified");
+        assert!(writer.unwrap().verified(), "downloaded and decrypted file could not be verified");
 
         // // Crpate metadata and a file reader
         // let metadata = self.create_metadata(&key, &file)?;

+ 122 - 5
api/src/reader.rs

@@ -304,9 +304,8 @@ impl<R: Read> Read for ProgressReader<R> {
 
         // Report
         if let Some(reporter) = self.reporter.as_mut() {
-            reporter.lock()
-                .expect("failed to update progress, unable to lock reproter")
-                .progress(self.progress);
+            let progress = self.progress;
+            let _ = reporter.lock().map(|mut r| r.progress(progress));
         }
 
         Ok(len)
@@ -333,12 +332,12 @@ pub trait ProgressReporter: Send {
 }
 
 /// A trait for readers, to get the exact length of a reader.
-pub trait ExactLengthReader: Read {
+pub trait ExactLengthReader {
     /// Get the exact length of the reader in bytes.
     fn len(&self) -> Result<u64, io::Error>;
 }
 
-impl<R: ExactLengthReader> ExactLengthReader for BufReader<R> {
+impl<R: ExactLengthReader + Read> ExactLengthReader for BufReader<R> {
     fn len(&self) -> Result<u64, io::Error> {
         self.get_ref().len()
     }
@@ -444,6 +443,12 @@ impl EncryptedFileWriter {
     }
 }
 
+impl ExactLengthReader for EncryptedFileWriter {
+    fn len(&self) -> Result<u64, IoError> {
+        Ok(self.len as u64)
+    }
+}
+
 /// The writer trait implementation.
 impl Write for EncryptedFileWriter {
     fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
@@ -508,3 +513,115 @@ impl Write for EncryptedFileWriter {
         self.file.flush()
     }
 }
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+/// A writer wrapper, that measures the reading process for a writer with a
+/// known length.
+///
+/// If the writer exceeds the initially specified length,
+/// the writer will continue to allow reads.
+/// The length property will grow accordingly.
+///
+/// The writer will only start producing `None` if the wrapped writer is doing
+/// so.
+pub struct ProgressWriter<W> {
+    /// The wrapped writer.
+    inner: W,
+
+    /// The total length of the writer.
+    len: u64,
+
+    /// The current reading progress.
+    progress: u64,
+
+    /// A reporter, to report the progress status to.
+    reporter: Option<Arc<Mutex<ProgressReporter>>>,
+}
+
+impl<W: Write> ProgressWriter<W> {
+    /// Wrap the given writer with an exact length, in a progress writer.
+    pub fn new(inner: W) -> Result<Self, IoError>
+        where
+            W: ExactLengthReader
+    {
+        Ok(
+            Self {
+                len: inner.len()?,
+                inner,
+                progress: 0,
+                reporter: None,
+            }
+        )
+    }
+
+    /// Wrap the given writer with the given length in a progress writer.
+    pub fn from(inner: W, len: u64) -> Self {
+        Self {
+            inner,
+            len,
+            progress: 0,
+            reporter: None,
+        }
+    }
+
+    /// Set the reporter to report the status to.
+    pub fn set_reporter(&mut self, reporter: Arc<Mutex<ProgressReporter>>) {
+        self.reporter = Some(reporter);
+    }
+
+    /// Get the current progress.
+    pub fn progress(&self) -> u64 {
+        self.progress
+    }
+
+    /// Unwrap the inner from the progress writer.
+    pub fn unwrap(self) -> W {
+        self.inner
+    }
+}
+
+impl<W: Write> Write for ProgressWriter<W> {
+    fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
+        // Write from the wrapped writer, increase the progress
+        let len = self.inner.write(buf)?;
+        self.progress += len as u64;
+
+        // Keep the specified length in-bound
+        if self.progress > self.len {
+            self.len = self.progress;
+        }
+
+        // Report
+        if let Some(reporter) = self.reporter.as_mut() {
+            let progress = self.progress;
+            let _ = reporter.lock().map(|mut r| r.progress(progress));
+        }
+
+        Ok(len)
+    }
+
+    fn flush(&mut self) -> Result<(), IoError> {
+        self.inner.flush()
+    }
+}
+
+impl<W: Write> ExactLengthReader for ProgressWriter<W> {
+    // Return the specified length.
+    fn len(&self) -> Result<u64, io::Error> {
+        Ok(self.len)
+    }
+}

+ 7 - 3
cli/src/action/download.rs

@@ -1,8 +1,11 @@
+use std::sync::{Arc, Mutex};
+
 use ffsend_api::action::download::Download as ApiDownload;
 use ffsend_api::file::file::DownloadFile;
 use ffsend_api::reqwest::Client;
 
 use cmd::cmd_download::CmdDownload;
+use progress::ProgressBar;
 
 /// A file download action.
 pub struct Download<'a> {
@@ -30,13 +33,14 @@ impl<'a> Download<'a> {
         let file = DownloadFile::parse_url(url)
             .expect("invalid download URL, could not parse file data");
 
+        // Create a progress bar reporter
+        let bar = Arc::new(Mutex::new(ProgressBar::new_download()));
+
         // Execute an download action
         // TODO: do not unwrap, but return an error
-        ApiDownload::new(&file).invoke(&client).unwrap();
+        ApiDownload::new(&file).invoke(&client, bar).unwrap();
 
         // TODO: open the file, or it's location
         // TODO: copy the file location
-
-        println!("Download complete");
     }
 }

+ 1 - 1
cli/src/action/upload.rs

@@ -34,7 +34,7 @@ impl<'a> Upload<'a> {
         let client = Client::new();
 
         // Create a progress bar reporter
-        let bar = Arc::new(Mutex::new(ProgressBar::new()));
+        let bar = Arc::new(Mutex::new(ProgressBar::new_upload()));
 
         // Execute an upload action
         // TODO: do not unwrap, but return an error

+ 28 - 6
cli/src/progress.rs

@@ -1,6 +1,7 @@
 extern crate pbr;
 
 use std::io::Stdout;
+use std::time::Duration;
 
 use ffsend_api::reader::ProgressReporter;
 use self::pbr::{
@@ -8,26 +9,47 @@ use self::pbr::{
     Units,
 };
 
+/// The refresh rate of the progress bar, in milliseconds.
+const PROGRESS_BAR_FPS_MILLIS: u64 = 500;
+
 /// A progress bar reporter.
-pub struct ProgressBar {
+pub struct ProgressBar<'a> {
     bar: Option<Pbr<Stdout>>,
+    msg_progress: &'a str,
+    msg_finish: &'a str,
 }
 
-impl ProgressBar {
-    pub fn new() -> Self {
+impl<'a> ProgressBar<'a> {
+    /// Construct a new progress bar, with the given messages.
+    pub fn new(msg_progress: &'a str, msg_finish: &'a str) -> ProgressBar<'a> {
         Self {
             bar: None,
+            msg_progress,
+            msg_finish,
         }
     }
+
+    /// Construct a new progress bar for uploading.
+    pub fn new_upload() -> ProgressBar<'a> {
+        Self::new("Encrypt & Upload ", "Upload complete")
+    }
+
+    /// Construct a new progress bar for downloading.
+    pub fn new_download() -> ProgressBar<'a> {
+        Self::new("Download & Decrypt ", "Download complete")
+    }
 }
 
-impl ProgressReporter for ProgressBar {
+impl<'a> ProgressReporter for ProgressBar<'a> {
     /// Start the progress with the given total.
     fn start(&mut self, total: u64) {
         // Initialize the progress bar
         let mut bar = Pbr::new(total);
+        bar.set_max_refresh_rate(
+            Some(Duration::from_millis(PROGRESS_BAR_FPS_MILLIS))
+        );
         bar.set_units(Units::Bytes);
-        bar.message("Encrypt & Upload ");
+        bar.message(self.msg_progress);
 
         self.bar = Some(bar);
     }
@@ -43,6 +65,6 @@ impl ProgressReporter for ProgressBar {
     fn finish(&mut self) {
         self.bar.as_mut()
             .expect("progress bar not yet instantiated")
-            .finish_print("Upload complete");
+            .finish_print(self.msg_finish);
     }
 }