Refine KeyStore logic

Author: DraftedDev

src/lib.rs

@@ -5,10 +5,7 @@ use std::path::PathBuf;
 
 use serde::{Deserialize, Serialize};
 
-use crate::{
-    error::Error,
-    secure::{KeyStore, PlatformKeyStore},
-};
+use crate::{error::Error, secure::get_key_store};
 
 /// Contains the [error::Error] type.
 pub mod error;
@@ -91,7 +88,7 @@ pub trait GlobalByteBox<const SECURE: bool>: Serialize + for<'de> Deserialize<'d
             let name = format!("{}", Self::path().display());
 
             let bytes = bitcode::serialize(self)?;
-            let key = PlatformKeyStore::get_key_or_generate(&name)?;
+            let key = get_key_store().get_key_or_generate(&name)?;
 
             secure::encrypt(&bytes, key)
         } else {
@@ -105,7 +102,7 @@ pub trait GlobalByteBox<const SECURE: bool>: Serialize + for<'de> Deserialize<'d
     fn decode(bytes: &[u8]) -> Result<Self, Error> {
         if SECURE {
             let name = Self::path().to_string_lossy().to_string();
-            let key = PlatformKeyStore::get_key_or_generate(&name)?;
+            let key = get_key_store().get_key_or_generate(&name)?;
             let decrypted = secure::decrypt(bytes.to_vec(), key)?;
 
             Ok(bitcode::deserialize(&decrypted)?)
@@ -189,7 +186,7 @@ pub trait ByteBox<const SECURE: bool>: Serialize + for<'de> Deserialize<'de> {
         if SECURE {
             let name = format!("{}", self.path().display());
             let bytes = bitcode::serialize(self)?;
-            let key = PlatformKeyStore::get_key_or_generate(&name)?;
+            let key = get_key_store().get_key_or_generate(&name)?;
 
             secure::encrypt(&bytes, key)
         } else {
@@ -204,7 +201,7 @@ pub trait ByteBox<const SECURE: bool>: Serialize + for<'de> Deserialize<'de> {
         if SECURE {
             let name = format!("{}", self.path().display());
 
-            let key = PlatformKeyStore::get_key_or_generate(&name)?;
+            let key = get_key_store().get_key_or_generate(&name)?;
             let decrypted = secure::decrypt(bytes.to_vec(), key)?;
 
             Ok(bitcode::deserialize(&decrypted)?)

src/secure/keyring.rs

@@ -6,7 +6,7 @@ use crate::secure::namespace;
 pub struct KeyStore;
 
 impl super::KeyStore for KeyStore {
-    fn get_key(name: &str) -> Result<Option<Vec<u8>>, crate::error::Error> {
+    fn get_key(&self, name: &str) -> Result<Option<Vec<u8>>, crate::error::Error> {
         let entry = Entry::new(namespace()?.as_str(), name)?;
 
         match entry.get_secret() {
@@ -16,7 +16,7 @@ impl super::KeyStore for KeyStore {
         }
     }
 
-    fn set_key(name: &str, key: Vec<u8>) -> Result<(), crate::error::Error> {
+    fn set_key(&self, name: &str, key: Vec<u8>) -> Result<(), crate::error::Error> {
         let entry = Entry::new(namespace()?.as_str(), name)?;
 
         entry.set_secret(&key)?;

src/secure/mod.rs

@@ -1,3 +1,5 @@
+use std::sync::OnceLock;
+
 use chacha20poly1305::{
     aead::{AeadMut, Key, Nonce, OsRng},
     AeadCore, KeyInit, XChaCha20Poly1305,
@@ -15,6 +17,28 @@ pub type PlatformKeyStore = keyring::KeyStore;
 
 const NONCE_SIZE: usize = 24;
 
+/// The global [KeyStore] instance.
+///
+/// Normally initialized with the default implementation on first access.
+pub static KEY_STORE: OnceLock<Box<dyn KeyStore>> = OnceLock::new();
+
+/// Returns the set [KeyStore] implementation, initializing it with the default if necessary.
+pub fn get_key_store<'a>() -> &'a Box<dyn KeyStore> {
+    KEY_STORE.get_or_init(default_key_store)
+}
+
+/// Returns the default [KeyStore] implementation for this platform.
+pub fn default_key_store() -> Box<dyn KeyStore> {
+    #[cfg(use_keyring)]
+    {
+        return Box::new(keyring::KeyStore);
+    }
+    #[cfg(not(use_keyring))]
+    {
+        panic!("Platform not supported!");
+    }
+}
+
 /// Encrypts the given data using the provided key.
 pub fn encrypt(data: &[u8], key: Vec<u8>) -> Result<Vec<u8>, Error> {
     let key = Key::<XChaCha20Poly1305>::from_iter(key);
@@ -64,23 +88,23 @@ pub fn namespace() -> Result<String, Error> {
 }
 
 /// A trait for storing and retrieving keys.
-pub trait KeyStore {
+pub trait KeyStore: Send + Sync + 'static {
     /// Returns the key with the given name, generating a new one if none exists.
-    fn get_key_or_generate(name: &str) -> Result<Vec<u8>, Error> {
-        let key = Self::get_key(name)?;
+    fn get_key_or_generate(&self, name: &str) -> Result<Vec<u8>, Error> {
+        let key = self.get_key(name)?;
 
         if let Some(key) = key {
             Ok(key)
         } else {
             let key = generate_key();
-            Self::set_key(name, key.to_vec())?;
+            self.set_key(name, key.to_vec())?;
             Ok(key.to_vec())
         }
     }
 
     /// Returns the key with the given name, if one exists.
-    fn get_key(name: &str) -> Result<Option<Vec<u8>>, Error>;
+    fn get_key(&self, name: &str) -> Result<Option<Vec<u8>>, Error>;
 
     /// Sets the key with the given name.
-    fn set_key(name: &str, key: Vec<u8>) -> Result<(), Error>;
+    fn set_key(&self, name: &str, key: Vec<u8>) -> Result<(), Error>;
 }