calculate hash only once, improve performance

Author: lxl66566

Cargo.lock

@@ -5,7 +5,7 @@ description = "Encrypt/decrypt files in git repo using one password"
 homepage = "https://github.com/lxl66566/git-simple-encrypt"
 repository = "https://github.com/lxl66566/git-simple-encrypt"
 license = "MIT"
-version = "1.1.0"
+version = "1.1.1"
 edition = "2021"
 readme = "./README.md"
 categories = ["cryptography"]
@@ -44,6 +44,7 @@ openssl = { version = "0.10", features = ["vendored"] }
 libz-sys = { version = "1.1.16", features = ["static"] }
 
 [dev-dependencies]
+rand = "0.8.5"
 temp_testdir = "0.2.3"
 
 [[bin]]

Cargo.toml

@@ -12,7 +12,6 @@ use anyhow::{anyhow, Context, Result};
 use colored::Colorize;
 use die_exit::die;
 use log::{debug, info};
-use sha3::{Digest, Sha3_224};
 use tap::Tap;
 
 #[cfg(any(test, debug_assertions))]
@@ -27,22 +26,7 @@ pub static ENCRYPTED_EXTENSION: &str = "enc";
 pub static COMPRESSED_EXTENSION: &str = "zst";
 
 pub fn encrypt(key: &[u8], text: &[u8]) -> std::result::Result<Vec<u8>, aes_gcm_siv::Error> {
-    #[cfg(any(test, debug_assertions))]
-    println!("Key: {}", String::from_utf8_lossy(key).green());
-    let mut hasher = Sha3_224::default();
-    hasher.update(key);
-    let hash_result = hasher.finalize();
-    let hash_result_slice = hash_result.as_slice();
-    #[cfg(any(test, debug_assertions))]
-    {
-        println!("Hash result: {}", format_hex(hash_result_slice).green());
-        println!(
-            "Hash Cut result: {}",
-            format_hex(&hash_result_slice[..16]).green()
-        );
-    }
-    let cipher =
-        Aes128GcmSiv::new_from_slice(&hash_result_slice[..16]).expect("cipher key length error.");
+    let cipher = Aes128GcmSiv::new_from_slice(key).expect("cipher key length error.");
     let encrypted = cipher.encrypt(NONCE.deref(), text)?;
 
     #[cfg(any(test, debug_assertions))]
@@ -52,10 +36,7 @@ pub fn encrypt(key: &[u8], text: &[u8]) -> std::result::Result<Vec<u8>, aes_gcm_
 }
 
 pub fn decrypt(key: &[u8], text: &[u8]) -> std::result::Result<Vec<u8>, aes_gcm_siv::Error> {
-    let mut hasher = Sha3_224::default();
-    hasher.update(key);
-    let cipher = Aes128GcmSiv::new_from_slice(&hasher.finalize().as_slice()[..16])
-        .expect("cipher key length error.");
+    let cipher = Aes128GcmSiv::new_from_slice(key).expect("cipher key length error.");
     let plaintext = cipher.decrypt(NONCE.deref(), text)?;
     Ok(plaintext)
 }
@@ -154,7 +135,7 @@ pub async fn encrypt_file(
 
     let (encrypted, new_file) = tokio::task::spawn_blocking(move || {
         let (compressed, new_file) = try_compress(&bytes, new_file, repo.conf.zstd_level)?;
-        encrypt_change_path(repo.get_key().as_bytes(), &compressed, new_file)
+        encrypt_change_path(repo.get_key_sha(), &compressed, new_file)
     })
     .await??;
 
@@ -179,8 +160,7 @@ pub async fn decrypt_file(
         .with_context(|| format!("{:?}", file.as_ref()))?;
 
     let (decompressed, new_file) = tokio::task::spawn_blocking(move || {
-        let (decrypted, new_file) =
-            try_decrypt_change_path(repo.get_key().as_bytes(), &bytes, new_file)?;
+        let (decrypted, new_file) = try_decrypt_change_path(repo.get_key_sha(), &bytes, new_file)?;
         try_decompress(&decrypted, new_file)
     })
     .await??;
@@ -253,7 +233,7 @@ mod test {
 
     #[test]
     fn test_encrypt_decrypt() -> Result<()> {
-        let key = b"123456";
+        let key = b"602bdc204140db0a";
         let content = b"456789";
         let encrypted_content = encrypt(key, content).unwrap();
         assert_ne!(content.to_vec(), encrypted_content);

src/crypt.rs

@@ -1,9 +1,15 @@
-use std::path::{Path, PathBuf};
+use std::{
+    path::{Path, PathBuf},
+    sync::OnceLock,
+};
 
 use anyhow::{anyhow, Result};
 use assert2::assert;
+#[cfg(any(test, debug_assertions))]
+use colored::Colorize;
 use die_exit::Die;
 use log::debug;
+use sha3::{Digest, Sha3_224};
 use tap::Tap;
 
 use crate::{
@@ -14,10 +20,11 @@ use crate::{
 pub const GIT_CONFIG_PREFIX: &str =
     const_str::replace!(concat!(env!("CARGO_CRATE_NAME"), "."), "_", "-");
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Default)]
 pub struct Repo {
     pub path: PathBuf,
     pub conf: Config,
+    pub key_sha: OnceLock<Box<[u8]>>,
 }
 
 impl Repo {
@@ -35,7 +42,11 @@ impl Repo {
         }
         println!("Open repo: {}", path.display());
         let conf = Config::load_or_create_from(path.join(CONFIG_FILE));
-        Ok(Self { path, conf })
+        Ok(Self {
+            path,
+            conf,
+            key_sha: OnceLock::new(),
+        })
     }
     pub fn path(&self) -> &Path {
         &self.path
@@ -60,6 +71,32 @@ impl Repo {
         self.get_config("key")
             .die("Key not found, please exec `git-se set key <KEY>` first.")
     }
+
+    /// returns the first 16 bytes of sha3-224 of the key.
+    /// The sha result will only be calculated once in the lifetime of the
+    /// object.
+    pub fn get_key_sha(&self) -> &[u8] {
+        self.key_sha.get_or_init(|| {
+            let key = self.get_key();
+            #[cfg(any(test, debug_assertions))]
+            println!("Key: {}", key.green());
+            let mut hasher = Sha3_224::default();
+            hasher.update(key);
+            let hash_result = hasher.finalize();
+            let hash_result_slice = hash_result.as_slice();
+            let hash_result_slice_cut = &hash_result_slice[..16];
+            #[cfg(any(test, debug_assertions))]
+            {
+                use crate::utils::format_hex;
+                println!("Hash result: {}", format_hex(hash_result_slice).green());
+                println!(
+                    "Hash Cut result: {}",
+                    format_hex(hash_result_slice_cut).green()
+                );
+            }
+            hash_result_slice_cut.into()
+        })
+    }
 }
 
 pub trait GitCommand {

src/repo.rs

@@ -27,24 +27,16 @@ impl PathAppendExt for PathBuf {
         self.tap_mut(|p| p.as_mut_os_string().push(format!(".{ext}")))
     }
 }
-pub trait PathStripPrefix {
-    fn strip_prefix_better(&mut self, prefix: impl AsRef<Path>) -> &mut Self;
-}
-impl PathStripPrefix for PathBuf {
-    fn strip_prefix_better(&mut self, prefix: impl AsRef<Path>) -> &mut Self {
-        if let Ok(res) = self.strip_prefix(prefix) {
-            *self = res.to_path_buf();
-        }
-        self
-    }
-}
 #[allow(unused)]
 pub trait PathToUnixStyle {
     fn to_unix_style(&self) -> PathBuf;
 }
 impl<T: AsRef<Path>> PathToUnixStyle for T {
     fn to_unix_style(&self) -> PathBuf {
-        self.as_ref().to_string_lossy().replace('\\', "/").into()
+        #[cfg(not(unix))]
+        return self.as_ref().to_string_lossy().replace('\\', "/").into();
+        #[cfg(unix)]
+        return self.as_ref().to_path_buf();
     }
 }
 
@@ -55,12 +47,22 @@ pub trait Git2Patch {
 }
 impl<T: AsRef<Path>> Git2Patch for T {
     fn patch(&self) -> PathBuf {
-        self.as_ref().to_path_buf().tap_mut(|x| {
-            #[cfg(target_family = "unix")]
-            x.strip_prefix_better("./");
-            #[cfg(target_family = "windows")]
-            x.strip_prefix_better(".\\");
-        })
+        let path = self.as_ref().to_path_buf();
+
+        #[cfg(target_family = "unix")]
+        let mut prefix = "./".to_string();
+        #[cfg(target_family = "windows")]
+        let mut prefix = ".\\".to_string();
+
+        let res = path.strip_prefix(&prefix);
+        if let Ok(mut ok_stripped) = res {
+            prefix.remove(0);
+            while let Ok(stripped_again) = ok_stripped.strip_prefix(&prefix) {
+                ok_stripped = stripped_again;
+            }
+            return ok_stripped.to_path_buf();
+        }
+        path
     }
 }
 
@@ -77,3 +79,30 @@ impl<T: AsRef<Path>> PathToAbsolute for T {
         })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    #[cfg(target_family = "unix")]
+    fn test_patch() {
+        let patched = Path::new("src/utils/pathutils.rs");
+        assert_eq!(patched.patch(), patched);
+        let p = Path::new("./src/utils/pathutils.rs");
+        assert_eq!(p.patch(), patched);
+        let p = Path::new(".////src/utils/pathutils.rs");
+        assert_eq!(p.patch(), patched);
+    }
+
+    #[test]
+    #[cfg(target_family = "windows")]
+    fn test_patch() {
+        let patched = Path::new("src\\utils\\pathutils.rs");
+        assert_eq!(patched.patch(), patched);
+        let p = Path::new(".\\src\\utils\\pathutils.rs");
+        assert_eq!(p.patch(), patched);
+        let p = Path::new(".\\\\\\\\src\\utils\\pathutils.rs");
+        assert_eq!(p.patch(), patched);
+    }
+}

src/utils/pathutils.rs

@@ -1,11 +1,16 @@
 #![crate_type = "proc-macro"]
 extern crate proc_macro;
 
-use std::process::{Command, Output};
+use std::{
+    path::PathBuf,
+    process::{Command, Output},
+};
 
 use anyhow::Ok;
 use colored::Colorize;
 use git_simple_encrypt::*;
+use rand::seq::SliceRandom;
+use tap::Tap;
 use temp_testdir::TempDir;
 
 #[tokio::test]
@@ -170,3 +175,59 @@ async fn test_reencrypt() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+#[ignore = "This test takes too long to run, and it's not necessary to run it every time. You can run it manually if you want."]
+async fn test_many_files() -> anyhow::Result<()> {
+    let _lock = TempDir::default();
+    let temp_dir = &_lock;
+    let exec = |cmd: &str| -> std::io::Result<Output> {
+        let mut temp = cmd.split_whitespace();
+        let mut command = Command::new(temp.next().unwrap());
+        command.args(temp).current_dir(temp_dir).output()
+    };
+    macro_rules! run {
+        ($cmd:expr) => {
+            run(&Cli {
+                command: $cmd,
+                repo: temp_dir.to_path_buf(),
+            })
+            .await?;
+        };
+    }
+
+    exec("git init")?;
+    let dir = temp_dir.join("dir");
+    std::fs::create_dir(&dir)?;
+    let files = (1..2000)
+        .map(|i| {
+            dir.join(format!("file{}.txt", i))
+                .tap(|f| std::fs::write(f, "Hello").unwrap())
+        })
+        .collect::<Vec<PathBuf>>();
+
+    // Set key
+    run!(SubCommand::Set {
+        field: SetField::key,
+        value: "123".to_owned(),
+    });
+
+    // Add file
+    run!(SubCommand::Add {
+        path: vec!["dir".into()]
+    });
+
+    // Encrypt
+    run!(SubCommand::Encrypt);
+    // Decrypt
+    run!(SubCommand::Decrypt);
+
+    // Test
+    for _ in 1..10 {
+        let file_name = files.choose(&mut rand::thread_rng()).unwrap();
+        println!("Testing file: {}", file_name.display());
+        assert_eq!(std::fs::read_to_string(file_name)?, "Hello");
+    }
+
+    Ok(())
+}