From ddd3ac3b2da4637d4114c501a1eb3eb8b047bb75 Mon Sep 17 00:00:00 2001 From: Hatter Jiang Date: Tue, 17 Oct 2023 00:40:41 +0800 Subject: [PATCH] feat: optimize decrypt --- src/cmd_decrypt.rs | 77 ++++++++++++++++++++++++---------------------- src/cmd_encrypt.rs | 19 +++++------- src/util.rs | 15 +++++++-- 3 files changed, 60 insertions(+), 51 deletions(-) diff --git a/src/cmd_decrypt.rs b/src/cmd_decrypt.rs index 2a289b7..f4dc65e 100644 --- a/src/cmd_decrypt.rs +++ b/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, let selected_envelop = select_envelop(&meta, config)?; - let key = try_decrypt_key(config, selected_envelop, pin, slot)?; - let nonce = opt_result!(util::decode_base64(&meta.nonce), "Decode nonce failed: {}"); + let key = SecVec(try_decrypt_key(config, selected_envelop, pin, slot)?); + let nonce = SecVec(opt_result!(util::decode_base64(&meta.nonce), "Decode nonce failed: {}")); - // debugging!("Decrypt key: {}", hex::encode(&key)); - debugging!("Decrypt nonce: {}", hex::encode(&nonce)); + // debugging!("Decrypt key: {}", hex::encode(&key.0)); + debugging!("Decrypt nonce: {}", hex::encode(&nonce.0)); - let enc_meta = parse_encrypted_meta(&meta, &key, &nonce)?; - parse_encrypted_comment(&meta, &key, &nonce)?; + let enc_meta = parse_encrypted_meta(&meta, &key.0, &nonce.0)?; + parse_encrypted_comment(&meta, &key.0, &nonce.0)?; if cmd_decrypt.skip_decrypt_file { information!("Decrypt file is skipped."); - } else { - let compressed_desc = iff!(meta.compress, " [compressed]", ""); - let start = Instant::now(); - - if cmd_decrypt.direct_print { - if meta.file_length > 10 * 1024 { - warning!("File too large(more than 10K) cannot direct print on console."); - } else { - let mut output: Vec = Vec::with_capacity(10 * 1024); - let _ = decrypt_file( - &mut file_in, meta.file_length, &mut output, &key, &nonce, 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), - } - } - } else { - let mut file_out = File::create(path_out)?; - let _ = decrypt_file( - &mut file_in, meta.file_length, &mut file_out, &key, &nonce, 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()); + return Ok(0); } - util::zeroize(key); - util::zeroize(nonce); - drop(file_in); + // 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."); + return Ok(0); + } + + let mut output: Vec = Vec::with_capacity(10 * 1024); + let _ = decrypt_file( + &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); + } + + // 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.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()); + if do_write_file_out & &cmd_decrypt.remove_file { util::remove_file_with_msg(path); } Ok(meta.file_length) } diff --git a/src/cmd_encrypt.rs b/src/cmd_encrypt.rs index 22c673c..20bf6e7 100644 --- a/src/cmd_encrypt.rs +++ b/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()?; diff --git a/src/util.rs b/src/util.rs index d0a7518..402fd4e 100644 --- a/src/util.rs +++ b/src/util.rs @@ -10,6 +10,14 @@ use zeroize::Zeroize; use crate::consts::TINY_ENC_FILE_EXT; +pub struct SecVec(pub Vec); + +impl Drop for SecVec { + fn drop(&mut self) { + self.0.zeroize() + } +} + pub fn read_pin(pin: &Option) -> String { match pin { Some(pin) => pin.to_string(), @@ -68,14 +76,15 @@ pub fn require_file_not_exists(path: impl AsRef) -> XResult<()> { } } -pub fn make_key256_and_nonce() -> (Vec, Vec) { +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 = key.into(); + let nonce_vec: Vec = 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 {