From aec87254743d32007c0d89816eba65d392827824 Mon Sep 17 00:00:00 2001 From: Pablu23 <43807157+Pablu23@users.noreply.github.com> Date: Thu, 9 Mar 2023 19:45:56 +0100 Subject: [PATCH] Added Parallelism, also cleaned up code a bit --- src/main.rs | 153 +++++++++++++++++++++++++++++++++++----------------- 1 file changed, 105 insertions(+), 48 deletions(-) diff --git a/src/main.rs b/src/main.rs index 4717826..c533c17 100644 --- a/src/main.rs +++ b/src/main.rs @@ -9,7 +9,9 @@ use std::{ fs::File, io, io::{Read, Write}, - // time::SystemTime, + sync::Arc, // time::SystemTime, + thread, + thread::{available_parallelism, JoinHandle}, }; use uuid::Uuid; @@ -22,10 +24,10 @@ const BUFFER_LEN: usize = 50 * 1024 * 1024; // 50 MiB fn encrypt_file( source_path: String, dest_path: String, - nonce: &[u8; 24], - key: &Vec, + nonce: [u8; 24], + key: &[u8], ) -> io::Result<()> { - let mut cipher = XChaCha20::new(key[..32].as_ref().into(), nonce.into()); + let mut cipher = XChaCha20::new(key[..32].as_ref().into(), &nonce.into()); let source_file_path = Path::new(&source_path); @@ -48,10 +50,12 @@ fn encrypt_file( // Stack allocated buffer // let mut buffer = [0u8; BUFFER_LEN]; + println!("Start encrypting File: {}", file_name.to_str().unwrap()); + // Heap allocated buffer (Allows larger sized buffer, up to 50 % max ram) let mut buffer = vec![0u8; BUFFER_LEN].into_boxed_slice(); - dest_file.write(nonce)?; + dest_file.write(&nonce)?; let mut f_name_bytes = file_name.to_str().unwrap_or_default().as_bytes().to_owned(); @@ -83,29 +87,23 @@ fn encrypt_file( fs::remove_file(source_file_path)?; - return Ok(()); + println!("Finished encrypting File: {}", file_name.to_str().unwrap()); + + Ok(()) } -fn decrypt_file(source_path: String, pwd: &String, config: &argon2::Config) -> io::Result<()> { +fn decrypt_file(source_path: &Path, pwd: &String, config: &argon2::Config) -> io::Result<()> { let mut nonce = [0u8; 24]; - let path_info = Path::new(&source_path); - - if !path_info.try_exists()? { + if !source_path.try_exists()? { return Err(io::Error::from(io::ErrorKind::NotFound)); } - //let path = &source_path[..&source_path.len() - 4]; - - // if Path::new(&path).exists() { - // fs::remove_file(&path)?; - // } - let mut source_file = File::open(&source_path)?; source_file.read(&mut nonce)?; - let key = argon2::hash_raw(pwd.as_bytes(), &nonce, &config).unwrap(); + let key = argon2::hash_raw(pwd.as_bytes(), &nonce, config).unwrap(); let mut cipher = XChaCha20::new(key[..32].as_ref().into(), &nonce.into()); @@ -116,14 +114,12 @@ fn decrypt_file(source_path: String, pwd: &String, config: &argon2::Config) -> i source_file.read(&mut file_name_size_buffer)?; let file_name_size = u16::from_le_bytes(file_name_size_buffer); - println!("File Name Size: {file_name_size}"); - let mut file_name_bytes = vec![0u8; file_name_size.into()]; source_file.read_exact(&mut file_name_bytes)?; cipher.apply_keystream(&mut file_name_bytes); - let private_dir_path = match path_info.parent() { + let private_dir_path = match source_path.parent() { Some(p) => Ok(p), None => Err(io::Error::from(io::ErrorKind::AddrNotAvailable)), }?; @@ -134,22 +130,18 @@ fn decrypt_file(source_path: String, pwd: &String, config: &argon2::Config) -> i }?; let file_name = String::from_utf8(file_name_bytes).unwrap_or_default(); - let path = root_dir_path.join(&file_name); //to_str().unwrap_or_default().to_owned() + &file_name; + let path = root_dir_path.join(&file_name); - println!("File Name: {file_name}\n Path: {path:?}"); + println!("Start decrypting File: {file_name}"); let mut dest_file = File::create(path)?; - println!("Worked!"); - - // Heap allocated buffer (Allows larger sized buffer, up to 50 % max ram) + // Heap allocated buffer (Allows larger sized buffer, up to 50 % max ram, technically more) let mut buffer = vec![0u8; BUFFER_LEN].into_boxed_slice(); loop { let read_count = source_file.read(&mut buffer)?; - println!("Read Bytes {read_count}"); - if read_count == BUFFER_LEN { cipher.apply_keystream(&mut buffer); dest_file.write(&buffer)?; @@ -160,43 +152,77 @@ fn decrypt_file(source_path: String, pwd: &String, config: &argon2::Config) -> i } } - println!("Fully written"); - return Ok(()); + println!("Finished decrypting File: {file_name}"); + Ok(()) } fn main() -> io::Result<()> { - let config = argon2::Config { + let config = Arc::new(argon2::Config { variant: argon2::Variant::Argon2id, hash_length: 32, lanes: 8, mem_cost: 16 * 1024, time_cost: 8, ..Default::default() - }; + }); let exe = env::current_exe().unwrap(); let cwd = env::current_dir().unwrap(); let private = cwd.join("private"); + let mut max_threads: usize = available_parallelism().unwrap().into(); + + let percent = max_threads as f32 * 0.7; + max_threads = percent.floor() as usize; + + if max_threads < 1 { + max_threads = 1; + } + + println!("Parallelism: {max_threads}"); + if private.exists() { let paths = fs::read_dir(&private).unwrap(); print!("Type password for encrypted files: "); std::io::stdout().flush().unwrap(); - let pwd = read_password().unwrap(); + let pwd = Arc::new(read_password().unwrap()); + + let config = Arc::new(config); + + let mut handles: Vec> = Vec::with_capacity(max_threads); + + let mut current_threads = 0; for path_result in paths { let path = path_result.unwrap().path(); + let pwd = pwd.clone(); + let config = config.clone(); - if path.is_file() { - decrypt_file(String::from(path.to_str().unwrap()), &pwd, &config)?; - fs::remove_file(String::from(path.to_str().unwrap()))?; + handles.push(thread::spawn(move || { + if path.is_file() { + decrypt_file(path.as_path(), &pwd, &config).unwrap(); + fs::remove_file(String::from(path.to_str().unwrap())).unwrap(); + } + })); + current_threads += 1; + + if current_threads >= max_threads { + while let Some(handle) = handles.pop() { + handle.join().unwrap(); + current_threads -= 1; + } + } + } + + if current_threads > 0 { + while let Some(handle) = handles.pop() { + handle.join().unwrap(); + current_threads -= 1; } } fs::remove_dir(private)?; - - //fs::remove_dir_all(private).unwrap(); } else { let paths = fs::read_dir(&cwd).unwrap(); @@ -222,26 +248,57 @@ fn main() -> io::Result<()> { print!("Type password to encrypt files: "); std::io::stdout().flush().unwrap(); - let pwd = read_password().unwrap(); + let pwd = Arc::new(read_password().unwrap()); let mut nonce = [0u8; 24]; OsRng.fill_bytes(&mut nonce); - let key = argon2::hash_raw(pwd.as_bytes(), &nonce, &config).unwrap(); + // let key = Arc::new(argon2::hash_raw(pwd.as_bytes(), &nonce, &config).unwrap()); let paths = fs::read_dir(cwd).unwrap(); + let exe = Arc::new(exe); + let private = Arc::new(private); + + let mut handles: Vec> = Vec::with_capacity(max_threads); + let mut current_threads = 0; for path_result in paths { - let path = path_result.unwrap().path(); + let pwd = pwd.clone(); + let exe = exe.clone(); + let private = private.clone(); + let config = config.clone(); - if path.is_file() && path != exe { - encrypt_file( - String::from(path.to_str().unwrap()), - String::from(private.to_str().unwrap()), - &nonce, - &key, - ) - .unwrap(); + handles.push(thread::spawn(move || { + let mut nonce = [0u8; 24]; + OsRng.fill_bytes(&mut nonce); + let key = argon2::hash_raw(pwd.as_bytes(), &nonce, &config).unwrap(); + + let path = path_result.unwrap().path(); + + if path.is_file() && path.as_os_str() != exe.as_os_str() { + encrypt_file( + String::from(path.to_str().unwrap()), + String::from(private.to_str().unwrap()), + nonce, + &key, + ) + .unwrap(); + } + })); + current_threads += 1; + + if current_threads >= max_threads { + while let Some(handle) = handles.pop() { + handle.join().unwrap(); + current_threads -= 1; + } + } + } + + if current_threads > 0 { + while let Some(handle) = handles.pop() { + handle.join().unwrap(); + current_threads -= 1; } } }