feat: optimize decrypt

src/cmd_decrypt.rs

@@ -25,6 +25,7 @@ use crate::consts::{
 };
 use crate::crypto_aes::{aes_gcm_decrypt, try_aes_gcm_decrypt_with_salt};
 use crate::spec::{EncEncryptedMeta, TinyEncryptEnvelop, TinyEncryptEnvelopType, TinyEncryptMeta};
+use crate::util::SecVec;
 use crate::util_progress::Progress;
 use crate::wrap_key::WrapKey;
 
@@ -107,50 +108,52 @@ pub fn decrypt_single(config: &Option<TinyEncryptConfig>,
 
     let selected_envelop = select_envelop(&meta, config)?;
 
-    let key = try_decrypt_key(config, selected_envelop, pin, slot)?;
+    let key = SecVec(try_decrypt_key(config, selected_envelop, pin, slot)?);
-    let nonce = opt_result!(util::decode_base64(&meta.nonce), "Decode nonce failed: {}");
+    let nonce = SecVec(opt_result!(util::decode_base64(&meta.nonce), "Decode nonce failed: {}"));
 
-    // debugging!("Decrypt key: {}", hex::encode(&key));
+    // debugging!("Decrypt key: {}", hex::encode(&key.0));
-    debugging!("Decrypt nonce: {}", hex::encode(&nonce));
+    debugging!("Decrypt nonce: {}", hex::encode(&nonce.0));
 
-    let enc_meta = parse_encrypted_meta(&meta, &key, &nonce)?;
+    let enc_meta = parse_encrypted_meta(&meta, &key.0, &nonce.0)?;
-    parse_encrypted_comment(&meta, &key, &nonce)?;
+    parse_encrypted_comment(&meta, &key.0, &nonce.0)?;
 
     if cmd_decrypt.skip_decrypt_file {
         information!("Decrypt file is skipped.");
-    } else {
+        return Ok(0);
-        let compressed_desc = iff!(meta.compress, " [compressed]", "");
+    }
-        let start = Instant::now();
 
+    // Decrypt to output
     if cmd_decrypt.direct_print {
         if meta.file_length > 10 * 1024 {
             warning!("File too large(more than 10K) cannot direct print on console.");
-            } else {
+            return Ok(0);
+        }
+
         let mut output: Vec<u8> = Vec::with_capacity(10 * 1024);
         let _ = decrypt_file(
-                    &mut file_in, meta.file_length, &mut output, &key, &nonce, meta.compress,
+            &mut file_in, meta.file_length, &mut output, &key.0, &nonce.0, meta.compress,
         )?;
         match String::from_utf8(output) {
             Err(_) => warning!("File is not UTF-8 content."),
             Ok(output) => println!(">>>>> BEGIN CONTENT >>>>>\n{}\n<<<<< END CONTENT <<<<<", &output),
         }
+        return Ok(meta.file_length);
     }
-        } else {
+
+    // Decrypt to file
+    let compressed_desc = iff!(meta.compress, " [compressed]", "");
+    let start = Instant::now();
+    
     let mut file_out = File::create(path_out)?;
     let _ = decrypt_file(
-                &mut file_in, meta.file_length, &mut file_out, &key, &nonce, meta.compress,
+        &mut file_in, meta.file_length, &mut file_out, &key.0, &nonce.0, meta.compress,
     )?;
     drop(file_out);
-        }
-
     util_file::update_out_file_time(enc_meta, path_out);
+
     let encrypt_duration = start.elapsed();
     debugging!("Inner decrypt file{}: {} elapsed: {} ms", compressed_desc, path_display, encrypt_duration.as_millis());
-    }
 
-    util::zeroize(key);
-    util::zeroize(nonce);
-    drop(file_in);
     if do_write_file_out & &cmd_decrypt.remove_file { util::remove_file_with_msg(path); }
     Ok(meta.file_length)
 }

src/cmd_encrypt.rs

@@ -124,12 +124,12 @@ fn encrypt_single(path: &PathBuf, envelops: &[&TinyEncryptConfigEnvelop], cmd_en
     util::require_file_not_exists(path_out.as_str())?;
 
     let (key, nonce) = util::make_key256_and_nonce();
-    let envelops = encrypt_envelops(&key, envelops)?;
+    let envelops = encrypt_envelops(&key.0, envelops)?;
 
     let encrypted_comment = match &cmd_encrypt.encrypted_comment {
         None => None,
         Some(encrypted_comment) => Some(util::encode_base64(
-            &aes_gcm_encrypt_with_salt(&key, &nonce, SALT_COMMENT, encrypted_comment.as_bytes())?))
+            &aes_gcm_encrypt_with_salt(&key.0, &nonce.0, SALT_COMMENT, encrypted_comment.as_bytes())?))
     };
 
     let file_metadata = opt_result!(fs::metadata(path), "Read file: {} meta failed: {}", path.display());
@@ -138,7 +138,7 @@ fn encrypt_single(path: &PathBuf, envelops: &[&TinyEncryptConfigEnvelop], cmd_en
         c_time: file_metadata.created().ok().and_then(|t| t.to_millis()),
         m_time: file_metadata.modified().ok().and_then(|t| t.to_millis()),
     };
-    let enc_encrypted_meta_bytes = opt_result!(enc_encrypted_meta.seal(&key, &nonce), "Seal enc-encrypted-meta failed: {}");
+    let enc_encrypted_meta_bytes = opt_result!(enc_encrypted_meta.seal(&key.0, &nonce.0), "Seal enc-encrypted-meta failed: {}");
 
     let enc_metadata = EncMetadata {
         comment: cmd_encrypt.comment.clone(),
@@ -147,7 +147,7 @@ fn encrypt_single(path: &PathBuf, envelops: &[&TinyEncryptConfigEnvelop], cmd_en
         compress: cmd_encrypt.compress,
     };
 
-    let mut encrypt_meta = TinyEncryptMeta::new(&file_metadata, &enc_metadata, &nonce, envelops);
+    let mut encrypt_meta = TinyEncryptMeta::new(&file_metadata, &enc_metadata, &nonce.0, envelops);
     debugging!("Encrypted meta: {:?}", encrypt_meta);
 
     if cmd_encrypt.compatible_with_1_0 {
@@ -162,15 +162,12 @@ fn encrypt_single(path: &PathBuf, envelops: &[&TinyEncryptConfigEnvelop], cmd_en
     let start = Instant::now();
     encrypt_file(
         &mut file_in, file_metadata.len(), &mut file_out,
-        &key, &nonce, cmd_encrypt.compress, &cmd_encrypt.compress_level,
+        &key.0, &nonce.0, cmd_encrypt.compress, &cmd_encrypt.compress_level,
     )?;
+    drop(file_out);
     let encrypt_duration = start.elapsed();
     debugging!("Inner encrypt file{}: {} elapsed: {} ms", compress_desc, path_display, encrypt_duration.as_millis());
 
-    util::zeroize(key);
-    util::zeroize(nonce);
-    drop(file_in);
-    drop(file_out);
     if cmd_encrypt.remove_file { util::remove_file_with_msg(path); }
     Ok(file_metadata.len())
 }
@@ -303,7 +300,7 @@ fn encrypt_envelop_shared_secret(key: &[u8],
     let shared_key = util::simple_kdf(shared_secret);
     let (_, nonce) = util::make_key256_and_nonce();
 
-    let encrypted_key = aes_gcm_encrypt(&shared_key, &nonce, key)?;
+    let encrypted_key = aes_gcm_encrypt(&shared_key, &nonce.0, key)?;
 
     let wrap_key = WrapKey {
         header: WrapKeyHeader {
@@ -311,7 +308,7 @@ fn encrypt_envelop_shared_secret(key: &[u8],
             enc: enc_type.to_string(),
             e_pub_key: util::encode_base64_url_no_pad(ephemeral_spki),
         },
-        nonce,
+        nonce: nonce.0.clone(),
         encrypted_data: encrypted_key,
     };
     let encoded_wrap_key = wrap_key.encode()?;

src/util.rs

@@ -10,6 +10,14 @@ use zeroize::Zeroize;
 
 use crate::consts::TINY_ENC_FILE_EXT;
 
+pub struct SecVec(pub Vec<u8>);
+
+impl Drop for SecVec {
+    fn drop(&mut self) {
+        self.0.zeroize()
+    }
+}
+
 pub fn read_pin(pin: &Option<String>) -> String {
     match pin {
         Some(pin) => pin.to_string(),
@@ -68,14 +76,15 @@ pub fn require_file_not_exists(path: impl AsRef<Path>) -> XResult<()> {
     }
 }
 
-pub fn make_key256_and_nonce() -> (Vec<u8>, Vec<u8>) {
+pub fn make_key256_and_nonce() -> (SecVec, SecVec) {
     let key: [u8; 32] = random();
     let nonce: [u8; 12] = random();
-    let result = (key.into(), nonce.into());
+    let key_vec: Vec<u8> = key.into();
+    let nonce_vec: Vec<u8> = nonce.into();
     let (mut key, mut nonce) = (key, nonce);
     key.zeroize();
     nonce.zeroize();
-    result
+    (SecVec(key_vec), SecVec(nonce_vec))
 }
 
 pub fn simple_kdf(input: &[u8]) -> Vec<u8> {