feat: gz supports assigned compress level

src/cmd_encrypt.rs

@@ -5,17 +5,17 @@ use std::path::PathBuf;
 use std::time::Instant;
 
 use clap::Args;
+use flate2::Compression;
 use rsa::Pkcs1v15Encrypt;
 use rust_util::{debugging, failure, information, opt_result, simple_error, success, warning, XResult};
 
+use crate::{util, util_ecdh};
 use crate::compress::GzStreamEncoder;
 use crate::config::{TinyEncryptConfig, TinyEncryptConfigEnvelop};
 use crate::crypto_aes::aes_gcm_encrypt;
 use crate::crypto_rsa::parse_spki;
 use crate::spec::{EncMetadata, TINY_ENCRYPT_VERSION_10, TinyEncryptEnvelop, TinyEncryptEnvelopType, TinyEncryptMeta};
-use crate::util;
-use crate::util::{ENC_AES256_GCM_P256, encode_base64, encode_base64_url_no_pad, make_key256_and_nonce, simple_kdf, TINY_ENC_CONFIG_FILE, zeroize};
-use crate::util_ecdh::compute_shared_secret;
+use crate::util::{ENC_AES256_GCM_P256, TINY_ENC_CONFIG_FILE};
 use crate::wrap_key::{WrapKey, WrapKeyHeader};
 
 #[derive(Debug, Args)]
@@ -34,6 +34,9 @@ pub struct CmdEncrypt {
     /// Compress before encrypt
     #[arg(long, short = 'x')]
     pub compress: bool,
+    /// Compress level (from 0[none], 1[fast] .. 6[default] .. to 9[best])
+    #[arg(long, short = 'L')]
+    pub compress_level: Option<u32>,
     /// Compatible with 1.0
     #[arg(long, short = '1')]
     pub compatible_with_1_0: bool,
@@ -68,12 +71,12 @@ fn encrypt_single(path: &PathBuf, envelops: &[&TinyEncryptConfigEnvelop], cmd_en
     let path_out = format!("{}{}", path_display, util::TINY_ENC_FILE_EXT);
     util::require_file_not_exists(path_out.as_str())?;
 
-    let (key, nonce) = make_key256_and_nonce();
+    let (key, nonce) = util::make_key256_and_nonce();
     let envelops = encrypt_envelops(&key, &envelops)?;
 
     let encrypted_comment = match &cmd_encrypt.encrypted_comment {
         None => None,
-        Some(encrypted_comment) => Some(encode_base64(
+        Some(encrypted_comment) => Some(util::encode_base64(
             &aes_gcm_encrypt(&key, &nonce, encrypted_comment.as_bytes())?))
     };
 
@@ -117,12 +120,12 @@ fn encrypt_single(path: &PathBuf, envelops: &[&TinyEncryptConfigEnvelop], cmd_en
     opt_result!(file_out.write_all(&encrypted_meta_bytes), "Write meta failed: {}");
 
     let start = Instant::now();
-    encrypt_file(&mut file_in, &mut file_out, &key, &nonce, cmd_encrypt.compress)?;
+    encrypt_file(&mut file_in, &mut file_out, &key, &nonce, cmd_encrypt.compress, &cmd_encrypt.compress_level)?;
     let encrypt_duration = start.elapsed();
     debugging!("Encrypt file: {} elapsed: {} ms", path_display, encrypt_duration.as_millis());
 
-    zeroize(key);
-    zeroize(nonce);
+    util::zeroize(key);
+    util::zeroize(nonce);
     drop(file_in);
     drop(file_out);
     if cmd_encrypt.remove_source_file {
@@ -135,11 +138,19 @@ fn encrypt_single(path: &PathBuf, envelops: &[&TinyEncryptConfigEnvelop], cmd_en
 }
 
 
-fn encrypt_file(file_in: &mut File, file_out: &mut File, key: &[u8], nonce: &[u8], compress: bool) -> XResult<usize> {
+fn encrypt_file(file_in: &mut File, file_out: &mut File, key: &[u8], nonce: &[u8], compress: bool, compress_level: &Option<u32>) -> XResult<usize> {
     let mut total_len = 0;
     let mut buffer = [0u8; 1024 * 8];
     let key = opt_result!(key.try_into(), "Key is not 32 bytes: {}");
-    let mut gz_encoder = GzStreamEncoder::new_default();
+    let mut gz_encoder = match compress_level {
+        None => GzStreamEncoder::new_default(),
+        Some(compress_level) => {
+            if *compress_level > 9 {
+                return simple_error!("Compress level must in range [0, 9]");
+            }
+            GzStreamEncoder::new(Compression::new(*compress_level))
+        }
+    };
     let mut encryptor = aes_gcm_stream::Aes256GcmStreamEncryptor::new(key, &nonce);
     loop {
         let len = opt_result!(file_in.read(&mut buffer), "Read file failed: {}");
@@ -170,6 +181,7 @@ fn encrypt_file(file_in: &mut File, file_out: &mut File, key: &[u8], nonce: &[u8
             opt_result!(file_out.write_all(&encrypted), "Write file failed: {}");
         }
     }
+    util::zeroize(key);
     Ok(total_len)
 }
 
@@ -191,9 +203,9 @@ fn encrypt_envelops(key: &[u8], envelops: &[&TinyEncryptConfigEnvelop]) -> XResu
 
 fn encrypt_envelop_ecdh(key: &[u8], envelop: &TinyEncryptConfigEnvelop) -> XResult<TinyEncryptEnvelop> {
     let public_key_point_hex = &envelop.public_part;
-    let (shared_secret, ephemeral_spki) = compute_shared_secret(public_key_point_hex)?;
-    let shared_key = simple_kdf(shared_secret.as_slice());
-    let (_, nonce) = make_key256_and_nonce();
+    let (shared_secret, ephemeral_spki) = util_ecdh::compute_shared_secret(public_key_point_hex)?;
+    let shared_key = util::simple_kdf(shared_secret.as_slice());
+    let (_, nonce) = util::make_key256_and_nonce();
 
     let encrypted_key = aes_gcm_encrypt(&shared_key, &nonce, key)?;
 
@@ -201,7 +213,7 @@ fn encrypt_envelop_ecdh(key: &[u8], envelop: &TinyEncryptConfigEnvelop) -> XResu
         header: WrapKeyHeader {
             kid: Some(envelop.kid.clone()),
             enc: ENC_AES256_GCM_P256.to_string(),
-            e_pub_key: encode_base64_url_no_pad(&ephemeral_spki),
+            e_pub_key: util::encode_base64_url_no_pad(&ephemeral_spki),
         },
         nonce,
         encrypted_data: encrypted_key,
@@ -224,6 +236,6 @@ fn encrypt_envelop_pgp(key: &[u8], envelop: &TinyEncryptConfigEnvelop) -> XResul
         r#type: envelop.r#type,
         kid: envelop.kid.clone(),
         desc: envelop.desc.clone(),
-        encrypted_key: encode_base64(&encrypted_key),
+        encrypted_key: util::encode_base64(&encrypted_key),
     })
 }