From e77433e0b0049cb77f7da544928de367abce953b Mon Sep 17 00:00:00 2001
From: Lucas Nogueira <lucas@tauri.studio>
Date: Tue, 8 Aug 2023 07:47:00 -0300
Subject: [PATCH] simplify impl

---
 .changes/stronghold-argon2.md |  5 ++++
 plugins/stronghold/src/kdf.rs | 23 +++-------------
 plugins/stronghold/src/lib.rs | 51 ++++++++++++++++++++++++-----------
 3 files changed, 44 insertions(+), 35 deletions(-)
 create mode 100644 .changes/stronghold-argon2.md

diff --git a/.changes/stronghold-argon2.md b/.changes/stronghold-argon2.md
new file mode 100644
index 00000000..598e4f80
--- /dev/null
+++ b/.changes/stronghold-argon2.md
@@ -0,0 +1,5 @@
+---
+"stronghold": patch
+---
+
+Added `Builder::with_argon2`.
diff --git a/plugins/stronghold/src/kdf.rs b/plugins/stronghold/src/kdf.rs
index f8ddf237..9d8dd120 100644
--- a/plugins/stronghold/src/kdf.rs
+++ b/plugins/stronghold/src/kdf.rs
@@ -1,13 +1,11 @@
 use argon2::Argon2;
 use rand_chacha::ChaCha20Rng;
 use rand_core::{RngCore, SeedableRng};
-use std::path::PathBuf;
-use tauri::Config;
+use std::path::Path;
 
 /// NOTE: Hash supplied to Stronghold must be 32bits long.
 /// This is a current limitation of Stronghold.
 const HASH_LENGTH: usize = 32;
-const SALT_FILENAME: &str = "stronghold_salt.txt";
 
 pub struct KeyDerivation {}
 
@@ -15,7 +13,7 @@ impl KeyDerivation {
     /// Will create a key from [`password`] and a generated salt.
     /// Salt will be generated to file [`salt_path`] or taken from it
     /// if file already exists
-    pub fn argon2(password: &str, salt_path: &PathBuf) -> Vec<u8> {
+    pub fn argon2(password: &str, salt_path: &Path) -> Vec<u8> {
         let mut salt = [0u8; HASH_LENGTH];
         create_or_get_salt(&mut salt, salt_path);
 
@@ -25,24 +23,9 @@ impl KeyDerivation {
             .expect("Failed to generate hash for password");
         encoded.to_vec()
     }
-
-    /// Will create a key from [`password`] and a generated salt.
-    /// Salt will be generated/taken from a default file in the Tauri local
-    /// directory
-    pub fn argon2_with_config(password: &str, tauri_config: &Config) -> Vec<u8> {
-        let salt_dir = tauri::api::path::app_local_data_dir(tauri_config)
-            .expect("Application local directory not found");
-        let mut salt_path = PathBuf::new();
-        salt_path.push(salt_dir);
-        salt_path.push(SALT_FILENAME);
-
-        KeyDerivation::argon2(password, &salt_path)
-    }
 }
 
-// NOTE: this is not ideal as we produce a single salt per application
-// rather than having different salt for each Stronghold snapshot/password
-fn create_or_get_salt(salt: &mut [u8], salt_path: &PathBuf) {
+fn create_or_get_salt(salt: &mut [u8], salt_path: &Path) {
     if salt_path.is_file() {
         // Get existing salt
         let tmp = std::fs::read(salt_path).unwrap();
diff --git a/plugins/stronghold/src/lib.rs b/plugins/stronghold/src/lib.rs
index 0695d73d..6e83f215 100644
--- a/plugins/stronghold/src/lib.rs
+++ b/plugins/stronghold/src/lib.rs
@@ -1,7 +1,7 @@
 use std::{
     collections::HashMap,
     fmt,
-    path::PathBuf,
+    path::{Path, PathBuf},
     sync::{Arc, Mutex},
     time::Duration,
 };
@@ -397,28 +397,44 @@ fn get_client(
     }
 }
 
+enum PasswordHashFunctionKind {
+    Argon2(PathBuf),
+    Custom(Box<PasswordHashFn>),
+}
+
 pub struct Builder {
-    password_hash_function: Box<PasswordHashFn>,
+    password_hash_function: PasswordHashFunctionKind,
 }
 
 impl Builder {
     pub fn new<F: Fn(&str) -> Vec<u8> + Send + Sync + 'static>(password_hash_function: F) -> Self {
         Self {
-            password_hash_function: Box::new(password_hash_function),
+            password_hash_function: PasswordHashFunctionKind::Custom(Box::new(
+                password_hash_function,
+            )),
         }
     }
 
-    /// Initializes a stronghold plugin with argon2 as a default kdf
-    pub fn init_and_build_with_argon2<R: Runtime>() -> TauriPlugin<R> {
-        let plugin_builder = PluginBuilder::new("stronghold").setup(move |app| {
-            let app2 = app.clone();
-            app.manage(StrongholdCollection::default());
-            app.manage(PasswordHashFunction(Box::new(move |pwd: &str| {
-                kdf::KeyDerivation::argon2_with_config(pwd, &app2.config())
-            })));
-            Ok(())
-        });
-        Builder::invoke_stronghold_handlers_and_build(plugin_builder)
+    /// Initializes [`Self`] with argon2 as password hash function.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// tauri::Builder::default()
+    ///     .setup(|app| {
+    ///         let salt_path = app
+    ///             .path_resolver()
+    ///             .app_local_data_dir()
+    ///             .expect("could not resolve app local data path")
+    ///             .join("salt.txt");
+    ///         app.handle().plugin(tauri_plugin_stronghold::Builder::with_argon2(&salt_path).build())?;
+    ///         Ok(())
+    ///     });
+    /// ```
+    pub fn with_argon2(salt_path: &Path) -> Self {
+        Self {
+            password_hash_function: PasswordHashFunctionKind::Argon2(salt_path.to_owned()),
+        }
     }
 
     pub fn build<R: Runtime>(self) -> TauriPlugin<R> {
@@ -426,7 +442,12 @@ impl Builder {
 
         let plugin_builder = PluginBuilder::new("stronghold").setup(move |app| {
             app.manage(StrongholdCollection::default());
-            app.manage(PasswordHashFunction(password_hash_function));
+            app.manage(PasswordHashFunction(match password_hash_function {
+                PasswordHashFunctionKind::Argon2(path) => {
+                    Box::new(move |p| kdf::KeyDerivation::argon2(p, &path))
+                }
+                PasswordHashFunctionKind::Custom(f) => f,
+            }));
             Ok(())
         });