From 061a552dc5c50a9c20caff633e0f60ceccd4b42f Mon Sep 17 00:00:00 2001 From: Diane Huxley Date: Mon, 13 May 2024 15:51:27 -0700 Subject: [PATCH] Refactor KeyManager --- .../key_store/in_memory_key_store.rs | 103 -------------- crates/keys/src/key_manager/key_store/mod.rs | 50 ------- .../keys/src/key_manager/local_key_manager.rs | 130 ++++++++++++++---- crates/keys/src/key_manager/mod.rs | 51 +++---- 4 files changed, 132 insertions(+), 202 deletions(-) delete mode 100644 crates/keys/src/key_manager/key_store/in_memory_key_store.rs delete mode 100644 crates/keys/src/key_manager/key_store/mod.rs diff --git a/crates/keys/src/key_manager/key_store/in_memory_key_store.rs b/crates/keys/src/key_manager/key_store/in_memory_key_store.rs deleted file mode 100644 index 059e4bf6..00000000 --- a/crates/keys/src/key_manager/key_store/in_memory_key_store.rs +++ /dev/null @@ -1,103 +0,0 @@ -use crate::key::{PrivateKey, PublicKey}; -use crate::key_manager::key_store::{KeyStore, KeyStoreError}; -use crypto::ed25519::Ed25519; -use crypto::secp256k1::Secp256k1; -use crypto::{Curve, CurveOperations}; -use std::collections::HashMap; -use std::sync::{Arc, RwLock}; - -pub struct InMemoryKeyStore { - map: RwLock>>, -} - -impl InMemoryKeyStore { - pub fn new() -> Self { - Self { - map: RwLock::new(HashMap::new()), - } - } -} - -impl KeyStore for InMemoryKeyStore { - fn generate_new( - &self, - curve: Curve, - key_alias: Option, - ) -> Result { - let private_key = Arc::new(match curve { - Curve::Ed25519 => Ed25519::generate(), - Curve::Secp256k1 => Secp256k1::generate(), - }?); - let key_alias = match key_alias { - Some(key_alias) => key_alias, - None => private_key.compute_thumbprint()?, - }; - let mut map_lock = self.map.write().map_err(|e| { - KeyStoreError::InternalKeyStoreError(format!("unable to acquire Mutex lock: {}", e)) - })?; - map_lock.insert(key_alias.clone(), private_key); - Ok(key_alias) - } - - fn get_all_aliases(&self) -> Result, KeyStoreError> { - let map_lock = self.map.read().map_err(|e| { - KeyStoreError::InternalKeyStoreError(format!("Unable to acquire Mutex lock: {}", e)) - })?; - let aliases = map_lock.keys().cloned().collect(); - Ok(aliases) - } - - fn sign(&self, key_alias: &str, payload: &[u8]) -> Result, KeyStoreError> { - let map_lock = self.map.read().map_err(|e| { - KeyStoreError::InternalKeyStoreError(format!("Unable to acquire Mutex lock: {}", e)) - })?; - let private_key = map_lock - .get(key_alias) - .ok_or(KeyStoreError::KeyNotFound(key_alias.to_string()))?; - - let signed_payload = private_key.sign(payload)?; - - Ok(signed_payload) - } - - fn get_public_key(&self, key_alias: &str) -> Result, KeyStoreError> { - let map_lock = self.map.read().map_err(|e| { - KeyStoreError::InternalKeyStoreError(format!("Unable to acquire Mutex lock: {}", e)) - })?; - let private_key = map_lock - .get(key_alias) - .ok_or(KeyStoreError::KeyNotFound(key_alias.to_string()))?; - let public_key = private_key.to_public()?; - Ok(public_key) - } - - fn export_private_keys(&self) -> Result>, KeyStoreError> { - let map_lock = self.map.read().map_err(|e| { - KeyStoreError::InternalKeyStoreError(format!("Unable to acquire Mutex lock: {}", e)) - })?; - let private_keys = map_lock.values().cloned().collect(); - Ok(private_keys) - } - - fn import_private_keys( - &self, - private_keys: Vec>, - ) -> Result<(), KeyStoreError> { - let mut map_lock = self.map.write().map_err(|e| { - KeyStoreError::InternalKeyStoreError(format!("Unable to acquire Mutex lock: {}", e)) - })?; - - for key in private_keys { - let key_alias = key.alias()?; - map_lock.insert(key_alias, key); - } - - Ok(()) - } -} - -impl Default for InMemoryKeyStore { - fn default() -> Self { - Self::new() - } -} diff --git a/crates/keys/src/key_manager/key_store/mod.rs b/crates/keys/src/key_manager/key_store/mod.rs deleted file mode 100644 index 93ab3b0c..00000000 --- a/crates/keys/src/key_manager/key_store/mod.rs +++ /dev/null @@ -1,50 +0,0 @@ -pub mod in_memory_key_store; - -use crate::key::{KeyError, PrivateKey, PublicKey}; -use crypto::{CryptoError, Curve}; -use jwk::JwkError; -use std::sync::Arc; - -#[derive(thiserror::Error, Debug, Clone, PartialEq)] -pub enum KeyStoreError { - #[error("{0}")] - InternalKeyStoreError(String), - #[error(transparent)] - KeyError(#[from] KeyError), - #[error("key not found {0}")] - KeyNotFound(String), - #[error(transparent)] - CryptoError(#[from] CryptoError), - #[error(transparent)] - JwkError(#[from] JwkError), - #[error("{0}")] - UnsupportedOperation(String), -} - -// Trait for storing and retrieving private keys. -// -// Implementations of this trait should be thread-safe and allow for concurrent access. -pub trait KeyStore: Send + Sync { - fn generate_new( - &self, - curve: Curve, - key_alias: Option, - ) -> Result; - fn get_all_aliases(&self) -> Result, KeyStoreError>; - fn sign(&self, key_alias: &str, payload: &[u8]) -> Result, KeyStoreError>; - fn get_public_key(&self, key_alias: &str) -> Result, KeyStoreError>; - - fn export_private_keys(&self) -> Result>, KeyStoreError> { - Err(KeyStoreError::UnsupportedOperation( - "exporting private keys is not supported".to_string(), - )) - } - fn import_private_keys( - &self, - _private_keys: Vec>, - ) -> Result<(), KeyStoreError> { - Err(KeyStoreError::UnsupportedOperation( - "importing private keys is not supported".to_string(), - )) - } -} diff --git a/crates/keys/src/key_manager/local_key_manager.rs b/crates/keys/src/key_manager/local_key_manager.rs index 9ab61f4d..e0b123cb 100644 --- a/crates/keys/src/key_manager/local_key_manager.rs +++ b/crates/keys/src/key_manager/local_key_manager.rs @@ -1,26 +1,25 @@ use crate::key::{PrivateKey, PublicKey}; -use crate::key_manager::key_store::in_memory_key_store::InMemoryKeyStore; -use crate::key_manager::key_store::KeyStore; use crate::key_manager::{KeyManager, KeyManagerError}; -use crypto::Curve; -use std::sync::Arc; +use crypto::ed25519::Ed25519; +use crypto::secp256k1::Secp256k1; +use crypto::{Curve, CurveOperations}; +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; + +use super::KeyImporter; /// Implementation of the [`KeyManager`] trait with key generation local to the device/platform it /// is being run. Key storage is provided by a [`KeyStore`] trait implementation, allowing the keys /// to be stored wherever is most appropriate for the application. pub struct LocalKeyManager { - key_store: Arc, + map: RwLock>>, } impl LocalKeyManager { - /// Constructs a new `LocalKeyManager` that stores keys in the provided `KeyStore`. - pub fn new(key_store: Arc) -> Self { - Self { key_store } - } - - pub fn new_in_memory() -> Self { + /// Constructs a new `LocalKeyManager` that stores keys in memory. + pub fn new() -> Self { Self { - key_store: Arc::new(InMemoryKeyStore::new()), + map: RwLock::new(HashMap::new()), } } } @@ -31,41 +30,80 @@ impl KeyManager for LocalKeyManager { curve: Curve, key_alias: Option, ) -> Result { - let key_alias = self.key_store.generate_new(curve, key_alias)?; + let private_key = Arc::new(match curve { + Curve::Ed25519 => Ed25519::generate(), + Curve::Secp256k1 => Secp256k1::generate(), + }?); + let key_alias = match key_alias { + Some(key_alias) => key_alias, + None => private_key.compute_thumbprint()?, + }; + let mut map_lock = self.map.write().map_err(|e| { + KeyManagerError::InternalKeyStoreError(format!("unable to acquire Mutex lock: {}", e)) + })?; + map_lock.insert(key_alias.clone(), private_key); Ok(key_alias) } fn get_public_key(&self, key_alias: &str) -> Result, KeyManagerError> { - let public_key = self.key_store.get_public_key(key_alias)?; + let map_lock = self.map.read().map_err(|e| { + KeyManagerError::InternalKeyStoreError(format!("Unable to acquire Mutex lock: {}", e)) + })?; + let private_key = map_lock + .get(key_alias) + .ok_or(KeyManagerError::KeyNotFound(key_alias.to_string()))?; + let public_key = private_key.to_public()?; Ok(public_key) } fn sign(&self, key_alias: &str, payload: &[u8]) -> Result, KeyManagerError> { - let signed_payload = self.key_store.sign(key_alias, payload)?; - Ok(signed_payload) - } + let map_lock: std::sync::RwLockReadGuard>> = self.map.read().map_err(|e| { + KeyManagerError::InternalKeyStoreError(format!("Unable to acquire Mutex lock: {}", e)) + })?; + let private_key = map_lock + .get(key_alias) + .ok_or(KeyManagerError::KeyNotFound(key_alias.to_string()))?; + + let signed_payload = private_key.sign(payload)?; - fn export_private_keys(&self) -> Result>, KeyManagerError> { - let private_keys = self.key_store.export_private_keys()?; - Ok(private_keys) + Ok(signed_payload) } +} - fn import_private_keys( +impl KeyImporter for LocalKeyManager { + fn import_with_alias( &self, - private_keys: Vec>, + private_key: Arc, + key_alias: &str, ) -> Result<(), KeyManagerError> { - self.key_store.import_private_keys(private_keys)?; + let mut map_lock = self.map.write().map_err(|e| { + KeyManagerError::InternalKeyStoreError(format!("Unable to acquire Mutex lock: {}", e)) + })?; + + map_lock.insert(key_alias.to_owned(), private_key); + Ok(()) } } + + +impl Default for LocalKeyManager { + fn default() -> Self { + Self::new() + } +} + + #[cfg(test)] mod tests { + use crate::key::Key; + use super::*; #[test] fn test_generate_private_key() { - let key_manager = LocalKeyManager::new_in_memory(); + let key_manager = LocalKeyManager::new(); key_manager .generate_private_key(Curve::Ed25519, None) @@ -84,7 +122,7 @@ mod tests { #[test] fn test_get_public_key() { - let key_manager = LocalKeyManager::new_in_memory(); + let key_manager = LocalKeyManager::new(); let key_alias = key_manager .generate_private_key(Curve::Ed25519, None) @@ -97,7 +135,7 @@ mod tests { #[test] fn test_sign() { - let key_manager = LocalKeyManager::new_in_memory(); + let key_manager: LocalKeyManager = LocalKeyManager::new(); let key_alias = key_manager .generate_private_key(Curve::Ed25519, None) .unwrap(); @@ -110,4 +148,44 @@ mod tests { let public_key = key_manager.get_public_key(&key_alias).unwrap(); assert!(!public_key.verify(payload, &signature).is_err()); } + + #[test] + fn test_import() { + let key_manager = LocalKeyManager::new(); + let private_key = Arc::new(Ed25519::generate().unwrap()); + + let key_alias = key_manager.import(private_key.clone()).expect("Failed to import private key"); + let default_alias = private_key.alias().expect("Failed to generate private key alias"); + assert_eq!(key_alias, default_alias); + + let key_manager_public_key = key_manager.get_public_key(&key_alias) + .expect("Failed to get public key") + .jwk() + .expect("Failed to get alias of public key"); + let public_key = private_key.to_public() + .expect("Failed to convert private key to public") + .jwk() + .expect("Failed to get alias of public key"); + assert_eq!(public_key, key_manager_public_key) + } + + #[test] + fn test_import_with_alias() { + let key_manager = LocalKeyManager::new(); + let private_key = Arc::new(Ed25519::generate().unwrap()); + + let key_alias = "1234".to_string(); + key_manager.import_with_alias(private_key.clone(), &key_alias) + .expect("Failed to import private key with alias"); + + let key_manager_public_key = key_manager.get_public_key(&key_alias) + .expect("Failed to get public key") + .jwk() + .expect("Failed to get alias of public key"); + let public_key = private_key.to_public() + .expect("Failed to convert private key to public") + .jwk() + .expect("Failed to get alias of public key"); + assert_eq!(public_key, key_manager_public_key) + } } diff --git a/crates/keys/src/key_manager/mod.rs b/crates/keys/src/key_manager/mod.rs index 200f9f5c..fd07e1d8 100644 --- a/crates/keys/src/key_manager/mod.rs +++ b/crates/keys/src/key_manager/mod.rs @@ -1,21 +1,24 @@ -pub mod key_store; pub mod local_key_manager; use crate::key::{KeyError, PrivateKey, PublicKey}; -use crate::key_manager::key_store::KeyStoreError; -use crypto::Curve; +use crypto::{CryptoError, Curve}; +use jwk::JwkError; use std::sync::Arc; #[derive(thiserror::Error, Debug, Clone, PartialEq)] pub enum KeyManagerError { + #[error(transparent)] + CryptoError(#[from] CryptoError), + #[error(transparent)] + JwkError(#[from] JwkError), #[error("Key generation failed")] KeyGenerationFailed, - #[error("Signing key not found in KeyManager")] - SigningKeyNotFound, #[error(transparent)] KeyError(#[from] KeyError), - #[error(transparent)] - KeyStoreError(#[from] KeyStoreError), + #[error("{0}")] + InternalKeyStoreError(String), + #[error("key not found {0}")] + KeyNotFound(String), } /// A key management trait for generating, storing, and utilizing keys private keys and their @@ -39,23 +42,25 @@ pub trait KeyManager: Send + Sync { /// Signs the provided payload using the private key identified by the provided `key_alias`. fn sign(&self, key_alias: &str, payload: &[u8]) -> Result, KeyManagerError>; +} - /// Exports all private keys managed by this key manager. - /// Default implementation returns an error indicating the feature is not supported. - fn export_private_keys(&self) -> Result>, KeyManagerError> { - Err(KeyStoreError::UnsupportedOperation( - "exporting private keys is not supported".to_string(), - ))? - } +pub trait KeyImporter: KeyManager { + /// Imports a private key with a custom key alias into the key manager. + /// Returns the key alias + fn import_with_alias( + &self, + private_key: Arc, + key_alias: &str, + ) -> Result<(), KeyManagerError>; - /// Imports a list of private keys into the key manager. - /// Default implementation returns an error indicating the feature is not supported. - fn import_private_keys( + /// Imports a private key into the key manager using private_key.alias() as the alias + /// Returns the key alias + fn import( &self, - _private_keys: Vec>, - ) -> Result<(), KeyManagerError> { - Err(KeyStoreError::UnsupportedOperation( - "importing private keys is not supported".to_string(), - ))? + private_key: Arc + ) -> Result { + let key_alias = private_key.alias()?; + self.import_with_alias(private_key, &key_alias)?; + Ok(key_alias) } -} +} \ No newline at end of file