Merge pull request #297 from marawan31/add-zip-download

Add zip download functionality for windows users

5 files changed, 184 insertions, 14 deletions

diff --git a/src/archive.rs b/src/archive.rs
index 268bb47..c96814f 100644
--- a/src/archive.rs
+++ b/src/archive.rs
@@ -1,9 +1,13 @@
 use actix_web::http::ContentEncoding;
 use libflate::gzip::Encoder;
+use zip::{ZipWriter, write};
+use std::fs::File;
+use std::path::PathBuf;
 use serde::Deserialize;
 use std::path::Path;
 use strum_macros::{Display, EnumIter, EnumString};
 use tar::Builder;
+use std::io::{Cursor, Write, Read};
 
 use crate::errors::ContextualError;
 
@@ -17,6 +21,9 @@ pub enum CompressionMethod {
 
     /// Regular tarball
     Tar,
+
+    /// Regular zip
+    Zip,
 }
 
 impl CompressionMethod {
@@ -24,6 +31,7 @@ impl CompressionMethod {
         match self {
             CompressionMethod::TarGz => "tar.gz",
             CompressionMethod::Tar => "tar",
+            CompressionMethod::Zip => "zip",
         }
         .to_string()
     }
@@ -32,6 +40,7 @@ impl CompressionMethod {
         match self {
             CompressionMethod::TarGz => "application/gzip",
             CompressionMethod::Tar => "application/tar",
+            CompressionMethod::Zip => "application/zip",
         }
         .to_string()
     }
@@ -40,6 +49,15 @@ impl CompressionMethod {
         match self {
             CompressionMethod::TarGz => ContentEncoding::Gzip,
             CompressionMethod::Tar => ContentEncoding::Identity,
+            CompressionMethod::Zip => ContentEncoding::Identity,
+        }
+    }
+
+    pub fn is_enabled(self, tar_enabled: bool, zip_enabled: bool,) -> bool {
+        match self {
+            CompressionMethod::TarGz => tar_enabled,
+            CompressionMethod::Tar => tar_enabled,
+            CompressionMethod::Zip => zip_enabled,
         }
     }
 
@@ -62,6 +80,7 @@ impl CompressionMethod {
         match self {
             CompressionMethod::TarGz => tar_gz(dir, skip_symlinks, out),
             CompressionMethod::Tar => tar_dir(dir, skip_symlinks, out),
+            CompressionMethod::Zip => zip_dir(dir, skip_symlinks, out),
         }
     }
 }
@@ -159,3 +178,137 @@ where
 
     Ok(())
 }
+
+/// Write a zip of `dir` in `out`.
+///
+/// The target directory will be saved as a top-level directory in the archive.
+///
+/// For example, consider this directory structure:
+///
+/// ```
+/// a
+/// └── b
+///     └── c
+///         ├── e
+///         ├── f
+///         └── g
+/// ```
+///
+/// Making a zip out of `"a/b/c"` will result in this archive content:
+///
+/// ```
+/// c
+/// ├── e
+/// ├── f
+/// └── g
+/// ```
+fn create_zip_from_directory<W>(out: W, directory: &PathBuf, skip_symlinks: bool,) -> Result<(), ContextualError> 
+where 
+    W: std::io::Write + std::io::Seek
+{
+    let options = write::FileOptions::default().compression_method(zip::CompressionMethod::Stored);
+    let mut paths_queue: Vec<PathBuf> = vec![];
+    paths_queue.push(directory.clone());
+    let zip_root_folder_name = directory.file_name().ok_or_else(|| {
+        ContextualError::InvalidPathError("Directory name terminates in \"..\"".to_string())
+    })?;
+
+    let mut zip_writer = ZipWriter::new(out);
+    let mut buffer = Vec::new();
+    while paths_queue.len() > 0 {
+        let next = paths_queue.pop().ok_or(ContextualError::CustomError("Could not get path from queue".to_string()))?;
+        let current_dir = next.as_path();
+        let directory_entry_iterator = std::fs::read_dir(current_dir).map_err(|e|{
+            ContextualError::IOError("Could not read directory".to_string(), e)
+        })?;
+        let zip_directory = Path::new(zip_root_folder_name)
+            .join(current_dir.strip_prefix(directory).map_err(|_|{
+                ContextualError::CustomError("Could not append base directory".to_string())
+            })?);
+
+        for entry in directory_entry_iterator {
+            let entry_path = entry.ok().ok_or(
+                ContextualError::InvalidPathError("Directory name terminates in \"..\"".to_string())
+            )?.path();
+            let entry_metadata = std::fs::metadata(entry_path.clone()).map_err(|e|{
+                ContextualError::IOError("Could not get file metadata".to_string(), e)
+            })?;
+
+            if entry_metadata.file_type().is_symlink() && skip_symlinks {
+                continue;
+            }
+            let current_entry_name = entry_path.file_name().ok_or_else(|| {
+                ContextualError::InvalidPathError("Invalid file or direcotory name".to_string())
+            })?;
+            if entry_metadata.is_file() {
+                let mut f = File::open(&entry_path).map_err(|e| {
+                     ContextualError::IOError("Could not open file".to_string(), e)
+                })?;
+                f.read_to_end(&mut buffer).map_err(|e| {
+                    ContextualError::IOError("Could not read from file".to_string(), e)
+                })?;
+                let relative_path = zip_directory.join(current_entry_name);
+                zip_writer.start_file_from_path(Path::new(&relative_path), options).map_err(|_| {
+                    ContextualError::CustomError("Could not add file path to ZIP".to_string())
+                })?;
+                zip_writer.write(buffer.as_ref()).map_err(|_| {
+                    ContextualError::CustomError("Could not write file to ZIP".to_string())
+                })?;
+                buffer.clear();
+            } else if entry_metadata.is_dir() {
+                let relative_path = zip_directory.join(current_entry_name);
+                zip_writer.add_directory_from_path(Path::new(&relative_path), options).map_err(|_| {
+                    ContextualError::CustomError("Could not add directory path to ZIP".to_string())
+                })?;
+                paths_queue.push(entry_path.clone());
+            }
+        }
+    }
+
+    zip_writer.finish().map_err(|_| {
+        ContextualError::CustomError("Could not finish writing ZIP archive".to_string())
+    })?;
+    Ok(())
+}
+
+/// Writes a zip of `dir` in `out`.
+///
+/// The content of `src_dir` will be saved in the archive as the  folder named .
+fn zip_data<W>(
+    src_dir: &Path,
+    skip_symlinks: bool,
+    mut out: W,
+) -> Result<(), ContextualError>
+where
+    W: std::io::Write,
+{
+    let mut data = Vec::new();
+    let memory_file = Cursor::new(&mut data);
+    create_zip_from_directory(memory_file, &src_dir.to_path_buf(), skip_symlinks).map_err(|e| {
+        ContextualError::ArchiveCreationError("Failed to create the ZIP archive".to_string(), Box::new(e))
+    })?;
+
+    out.write_all(data.as_mut_slice()).map_err(|e| {
+        ContextualError::IOError("Failed to write the ZIP archive".to_string(), e)
+    })?;
+
+    Ok(())
+}
+
+fn zip_dir<W>(dir: &Path, skip_symlinks: bool, out: W) -> Result<(), ContextualError>
+where
+    W: std::io::Write,
+{
+    let inner_folder = dir.file_name().ok_or_else(|| {
+        ContextualError::InvalidPathError("Directory name terminates in \"..\"".to_string())
+    })?;
+
+    inner_folder.to_str().ok_or_else(|| {
+        ContextualError::InvalidPathError(
+            "Directory name contains invalid UTF-8 characters".to_string(),
+        )
+    })?;
+
+    zip_data(dir, skip_symlinks, out)
+        .map_err(|e| ContextualError::ArchiveCreationError("zip".to_string(), Box::new(e)))
+}
\ No newline at end of file
diff --git a/src/args.rs b/src/args.rs
index 49fe276..ea4e90f 100644
--- a/src/args.rs
+++ b/src/args.rs
@@ -84,9 +84,15 @@ struct CLIArgs {
     #[structopt(short = "o", long = "overwrite-files")]
     overwrite_files: bool,
 
-    /// Disable archive generation
-    #[structopt(short = "r", long = "disable-archives")]
-    disable_archives: bool,
+    /// Enable tar archive generation
+    #[structopt(short = "r", long = "enable-tar")]
+    enable_tar: bool,
+
+    /// Enable zip archive generation
+    /// Zipping large directories can result in out-of-memory exception
+    /// because zip generation is done in memory and cannot be sent on the fly
+    #[structopt(short = "z", long = "enable-zip")]
+    enable_zip: bool,
 }
 
 /// Checks wether an interface is valid, i.e. it can be parsed into an IP address
@@ -176,7 +182,8 @@ pub fn parse_args() -> crate::MiniserveConfig {
         index: args.index,
         overwrite_files: args.overwrite_files,
         file_upload: args.file_upload,
-        archives: !args.disable_archives,
+        tar_enabled: args.enable_tar,
+        zip_enabled: args.enable_zip,
     }
 }
 
diff --git a/src/listing.rs b/src/listing.rs
index d28824c..45b3732 100644
--- a/src/listing.rs
+++ b/src/listing.rs
@@ -136,7 +136,8 @@ pub fn directory_listing<S>(
     random_route: Option<String>,
     default_color_scheme: ColorScheme,
     upload_route: String,
-    archives_enabled: bool,
+    tar_enabled: bool,
+    zip_enabled: bool,
 ) -> Result<HttpResponse, io::Error> {
     let serve_path = req.path();
 
@@ -250,7 +251,7 @@ pub fn directory_listing<S>(
     let color_scheme = query_params.theme.unwrap_or(default_color_scheme);
 
     if let Some(compression_method) = query_params.download {
-        if !archives_enabled {
+        if !compression_method.is_enabled(tar_enabled, zip_enabled) {
             return Ok(HttpResponse::Forbidden()
                 .content_type("text/html; charset=utf-8")
                 .body(
@@ -332,7 +333,8 @@ pub fn directory_listing<S>(
                     file_upload,
                     &upload_route,
                     &current_dir.display().to_string(),
-                    archives_enabled,
+                    tar_enabled,
+                    zip_enabled,
                 )
                 .into_string(),
             ))
diff --git a/src/main.rs b/src/main.rs
index a1ee303..3ff35c8 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -63,8 +63,11 @@ pub struct MiniserveConfig {
     /// Enable upload to override existing files
     pub overwrite_files: bool,
 
-    /// If false, creation of archives is disabled
-    pub archives: bool,
+    /// If false, creation of tar archives is disabled
+    pub tar_enabled: bool,
+
+    /// If false, creation of zip archives is disabled
+    pub zip_enabled: bool,
 }
 
 fn main() {
@@ -256,7 +259,8 @@ fn configure_app(app: App<MiniserveConfig>) -> App<MiniserveConfig> {
         let random_route = app.state().random_route.clone();
         let default_color_scheme = app.state().default_color_scheme;
         let file_upload = app.state().file_upload;
-        let archives_enabled = app.state().archives;
+        let tar_enabled = app.state().tar_enabled;
+        let zip_enabled = app.state().zip_enabled;
         upload_route = if let Some(random_route) = app.state().random_route.clone() {
             format!("/{}/upload", random_route)
         } else {
@@ -285,7 +289,8 @@ fn configure_app(app: App<MiniserveConfig>) -> App<MiniserveConfig> {
                             random_route.clone(),
                             default_color_scheme,
                             u_r.clone(),
-                            archives_enabled,
+                            tar_enabled,
+                            zip_enabled,
                         )
                     })
                     .default_handler(error_404),
diff --git a/src/renderer.rs b/src/renderer.rs
index face6ff..accb49b 100644
--- a/src/renderer.rs
+++ b/src/renderer.rs
@@ -22,7 +22,8 @@ pub fn page(
     file_upload: bool,
     upload_route: &str,
     current_dir: &str,
-    archives: bool,
+    tar_enabled: bool,
+    zip_enabled: bool,
 ) -> Markup {
     let upload_action = build_upload_action(
         upload_route,
@@ -50,10 +51,12 @@ pub fn page(
                     span#top { }
                     h1.title { "Index of " (serve_path) }
                     div.toolbar {
-                        @if archives {
+                        @if tar_enabled || zip_enabled {
                             div.download {
                                 @for compression_method in CompressionMethod::iter() {
-                                    (archive_button(compression_method, sort_method, sort_order, color_scheme, default_color_scheme))
+                                    @if compression_method.is_enabled(tar_enabled, zip_enabled) {
+                                        (archive_button(compression_method, sort_method, sort_order, color_scheme, default_color_scheme))
+                                    }
                                 }
                             }
                         }