feat: add utility for benchmarking kdf functions

Author: louib

Cargo.toml

@@ -130,5 +130,10 @@ required-features = ["utilities", "save_kdbx4", "challenge_response"]
 name = "kp-yk-recover"
 required-features = ["utilities", "save_kdbx4", "challenge_response"]
 
+[[bin]]
+name = "kp-benchmark-kdf"
+required-features = ["utilities"]
+
 [lints.rust]
 unexpected_cfgs = { level = "warn", check-cfg = ['cfg(tarpaulin_include)'] }
+

src/bin/kp-benchmark-kdf.rs

@@ -0,0 +1,57 @@
+use std::time::Duration;
+
+use clap::{Parser, ValueEnum};
+use keepass::crypt::kdf::{AesKdf, Argon2Kdf, Kdf};
+
+#[derive(Parser, Debug)]
+#[command(version, about, long_about = None)]
+struct Args {
+    /// KDF to benchmark
+    #[arg(value_enum)]
+    kdf: KdfChoice,
+
+    /// Duration for each KDF in milliseconds
+    #[arg(short, long, default_value_t = 1000)]
+    msecs: u64,
+}
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
+enum KdfChoice {
+    Aes,
+    Argon2,
+}
+
+fn main() {
+    let args = Args::parse();
+    let duration = Duration::from_millis(args.msecs);
+
+    match args.kdf {
+        KdfChoice::Aes => {
+            let kdf = AesKdf {
+                seed: vec![0; 32],
+                rounds: 100_000,
+            };
+            println!("Benchmarking AES KDF for {} ms...", args.msecs);
+            let rounds = kdf.benchmark(duration);
+            println!("AES KDF: {} rounds in {} ms", rounds, args.msecs);
+        }
+        KdfChoice::Argon2 => {
+            let kdf = Argon2Kdf {
+                salt: vec![0; 32],
+                parallelism: 1,
+                memory: 1024 * 1024, // 1 MiB
+                iterations: 1,
+                variant: argon2::Variant::Argon2id,
+                version: argon2::Version::Version13,
+            };
+            println!(
+                "Benchmarking Argon2id KDF with {} KiB memory and parallelism {} for {} ms...",
+                kdf.memory / 1024,
+                kdf.parallelism,
+                args.msecs
+            );
+            let iterations = kdf.benchmark(duration);
+            println!("Argon2id KDF: {} iterations in {} ms", iterations, args.msecs);
+        }
+    }
+}

src/crypt/kdf.rs

@@ -1,3 +1,5 @@
+use std::time::Duration;
+
 use aes::Aes256;
 use cipher::{
     generic_array::{typenum::U32, GenericArray},
@@ -7,11 +9,13 @@ use sha2::{Digest, Sha256};
 
 use super::CryptographyError;
 
-pub(crate) trait Kdf {
+pub trait Kdf {
     fn transform_key(
         &self,
         composite_key: &GenericArray<u8, U32>,
     ) -> Result<GenericArray<u8, U32>, CryptographyError>;
+
+    fn benchmark(&self, duration: Duration) -> usize;
 }
 
 pub struct AesKdf {
@@ -39,6 +43,23 @@ impl Kdf for AesKdf {
 
         Ok(digest.finalize())
     }
+
+    fn benchmark(&self, duration: Duration) -> usize {
+        let composite_key: GenericArray<u8, U32> = GenericArray::clone_from_slice(&[0; 32]);
+        let mut rounds = 0;
+
+        let cipher = Aes256::new(&GenericArray::clone_from_slice(&self.seed));
+        let mut block1 = GenericArray::clone_from_slice(&composite_key[..16]);
+        let mut block2 = GenericArray::clone_from_slice(&composite_key[16..]);
+
+        let start_time = std::time::Instant::now();
+        while start_time.elapsed() < duration {
+            cipher.encrypt_block(&mut block1);
+            cipher.encrypt_block(&mut block2);
+            rounds = rounds + 1;
+        }
+        rounds
+    }
 }
 
 pub struct Argon2Kdf {
@@ -71,16 +92,31 @@ impl Kdf for Argon2Kdf {
 
         Ok(*GenericArray::from_slice(&key))
     }
-}
 
-/*
-pub(crate) fn transform_key_argon2(
-    composite_key: &GenericArray<u8, U32>,
-) -> Result<GenericArray<u8, U32>> {
-    let version = match version {
-        0x10 => argon2::Version::Version10,
-        0x13 => argon2::Version::Version13,
-        _ => return Err(DatabaseIntegrityError::InvalidKDFVersion { version: version }.into()),
-    };
+    fn benchmark(&self, duration: Duration) -> usize {
+        let composite_key: GenericArray<u8, U32> = GenericArray::clone_from_slice(&[0; 32]);
+
+        let config = argon2::Config {
+            thread_mode: argon2::ThreadMode::Parallel,
+            ad: &[],
+            hash_length: 32,
+            lanes: self.parallelism,
+            mem_cost: (self.memory / 1024) as u32,
+            secret: &[],
+            time_cost: 1, // benchmark for one iteration
+            variant: self.variant,
+            version: self.version,
+        };
+
+        let start_time = std::time::Instant::now();
+        let _ = argon2::hash_raw(&composite_key, &self.salt, &config);
+        let elapsed = start_time.elapsed();
+
+        if elapsed.is_zero() {
+            // Should not happen, but to be safe
+            return 0;
+        }
+
+        (duration.as_nanos() / elapsed.as_nanos()) as usize
+    }
 }
-*/

src/crypt/mod.rs

@@ -13,7 +13,7 @@ use sha2::{Digest, Sha256, Sha512};
 use crate::error::CryptographyError;
 
 pub(crate) mod ciphers;
-pub(crate) mod kdf;
+pub mod kdf;
 
 pub(crate) fn calculate_hmac(
     elements: &[&[u8]],

src/lib.rs

@@ -3,7 +3,7 @@
 
 mod compression;
 pub mod config;
-pub(crate) mod crypt;
+pub mod crypt;
 pub mod db;
 pub mod error;
 pub(crate) mod format;