From 6c5c82ee1cd273361ba075dfa194e46e9eb546db Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=B2=81=E6=A0=91=E4=BA=BA?= <lu.shuren@um-react.app>
Date: Thu, 5 Sep 2024 02:21:34 +0100
Subject: [PATCH] feat: allow override of buffer size

---
 um_cli/src/cmd/qmc1.rs |  8 +++-----
 um_cli/src/main.rs     | 18 ++++++++++++------
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/um_cli/src/cmd/qmc1.rs b/um_cli/src/cmd/qmc1.rs
index 11007cd..ea68f9a 100644
--- a/um_cli/src/cmd/qmc1.rs
+++ b/um_cli/src/cmd/qmc1.rs
@@ -3,6 +3,7 @@ use clap::Args;
 use std::fs::File;
 use std::io::{Read, Write};
 use std::path::PathBuf;
+use crate::Cli;
 
 /// Decrypt a QMCv1 file
 #[derive(Args)]
@@ -16,16 +17,13 @@ pub struct ArgsQMCv1 {
     input: PathBuf,
 }
 
-// 4MiB buffer is working well on my machine.
-const BUFFER_SIZE: usize = 4 * 1024 * 1024;
-
 impl ArgsQMCv1 {
-    pub fn run(&self) -> Result<i32> {
+    pub fn run(&self, cli: &Cli) -> Result<i32> {
         let mut file_input = File::open(&self.input)?;
         let mut file_output = File::create(&self.output)?;
 
         let mut offset = 0usize;
-        let mut buffer = vec![0u8; BUFFER_SIZE].into_boxed_slice();
+        let mut buffer = vec![0u8; cli.buffer_size].into_boxed_slice();
         while let Ok(n) = file_input.read(&mut buffer) {
             if n == 0 {
                 break;
diff --git a/um_cli/src/main.rs b/um_cli/src/main.rs
index fc034f7..8d436c7 100644
--- a/um_cli/src/main.rs
+++ b/um_cli/src/main.rs
@@ -6,22 +6,28 @@ use std::time::Instant;
 
 mod cmd;
 
+/// um_cli (rust ver.)
+/// A cli-tool to unlock encrypted audio files.
 #[derive(Parser)]
 #[command(name = "um_cli")]
 #[command(version = "0.1")]
-#[command(about = "um_cli (rust ver.)", long_about = None)]
-struct Cli {
+pub struct Cli {
     #[clap(subcommand)]
     command: Option<Commands>,
 
-    /// Time command
+    /// Be more verbose about what is going on.
     #[clap(long, short, action=clap::ArgAction::SetTrue)]
-    time: Option<bool>,
+    verbose: Option<bool>,
+
+    /// Preferred buffer size when reading file, in bytes.
+    /// Default to 4MiB.
+    #[clap(long, short = 'B', default_value_t=4*1024*1024)]
+    buffer_size: usize
 }
 
 fn run_command(cli: &Cli) -> Result<i32> {
     match &cli.command {
-        Some(Commands::QMCv1(cmd)) => cmd.run(),
+        Some(Commands::QMCv1(cmd)) => cmd.run(&cli),
         None => {
             // https://github.com/clap-rs/clap/issues/3857#issuecomment-1161796261
             todo!("implement a sensible default command, similar to um/cli");
@@ -37,7 +43,7 @@ fn main() {
         -1
     });
     let duration = start.elapsed();
-    if let Some(true) = cli.time {
+    if let Some(true) = cli.verbose {
         eprintln!("time: {:?}", duration);
     };
     exit(code);