tari-project · aviator-app · Apr 29, 2022 · Apr 25, 2022 · Apr 25, 2022 · Apr 25, 2022
@@ -26,7 +26,7 @@ use tari_script::ScriptError;
 use tari_utilities::{hex::HexError, ByteArrayError};
 
 use crate::error::WalletStorageError;
-
+/// Error enum for the [KeyManagerService]
 #[derive(Debug, thiserror::Error)]
 pub enum KeyManagerServiceError {
     #[error("Branch does not exist")]
@@ -42,7 +42,7 @@ pub enum KeyManagerServiceError {
     #[error("Tari Key Manager error: `{0}`")]
     TariKeyManagerError(#[from] KMError),
 }
-
+/// Error enum for the [KeyManagerStorage]
 #[derive(Debug, thiserror::Error)]
 pub enum KeyManagerStorageError {
     #[error("Value not found")]

@@ -35,7 +35,11 @@ use crate::key_manager_service::{
     KeyManagerInner,
     KeyManagerInterface,
 };
-
+/// The key manager provides a hierarchical key derivation function (KDF) that derives uniformly random secret keys from
+/// a single seed key for arbitrary branches, using an implementation of `KeyManagerBackend` to store the current index
+/// for each branch.
+///
+/// This handle can be cloned cheaply and safely shared across multiple threads.
 #[derive(Clone)]
 pub struct KeyManagerHandle<TBackend> {
     key_manager_inner: Arc<RwLock<KeyManagerInner<TBackend>>>,
@@ -44,6 +48,9 @@ pub struct KeyManagerHandle<TBackend> {
 impl<TBackend> KeyManagerHandle<TBackend>
 where TBackend: KeyManagerBackend + 'static
 {
+    /// Creates a new key manager.
+    /// * `master_seed` is the primary seed that will be used to derive all unique branch keys with their indexes
+    /// * `db` implements `KeyManagerBackend` and is used for persistent storage of branches and indices.
     pub fn new(master_seed: CipherSeed, db: KeyManagerDatabase<TBackend>) -> Self {
         KeyManagerHandle {
             key_manager_inner: Arc::new(RwLock::new(KeyManagerInner::new(master_seed, db))),

@@ -36,6 +36,7 @@ use crate::key_manager_service::{
     KeyManagerHandle,
 };
 
+/// Initializes the key manager service by implementing the [ServiceInitializer] trait.
 pub struct KeyManagerInitializer<T>
 where T: KeyManagerBackend
 {
@@ -46,6 +47,7 @@ where T: KeyManagerBackend
 impl<T> KeyManagerInitializer<T>
 where T: KeyManagerBackend + 'static
 {
+    /// Creates a new [KeyManagerInitializer] from the provided [KeyManagerBackend] and [CipherSeed]
     pub fn new(backend: T, master_seed: CipherSeed) -> Self {
         Self {
             backend: Some(backend),

@@ -25,6 +25,8 @@ use tari_common_types::types::PrivateKey;
 
 use crate::key_manager_service::error::KeyManagerServiceError;
 
+/// The value returned from [add_new_branch]. `AlreadyExists` is returned if the branch was previously created,
+/// otherwise `NewEntry` is returned.
 #[derive(Debug, PartialEq)]
 pub enum AddResult {
     NewEntry,
@@ -36,28 +38,41 @@ pub struct NextKeyResult {
     pub index: u64,
 }
 
+/// Behaviour required for the Key manager service
 #[async_trait::async_trait]
 pub trait KeyManagerInterface: Clone + Send + Sync + 'static {
+    /// Creates a new branch for the key manager service to track
+    /// If this is an existing branch, that is not yet tracked in memory, the key manager service will load the key
+    /// manager from the backend to track in memory, will return `Ok(AddResult::NewEntry)`. If the branch is already
+    /// tracked in memory the result will be `Ok(AddResult::AlreadyExists)`. If the branch does not exist in memory
+    /// or in the backend, a new branch will be created and tracked the backend, `Ok(AddResult::NewEntry)`.
     async fn add_new_branch<T: Into<String> + Send>(&self, branch: T) -> Result<AddResult, KeyManagerServiceError>;
 
+    /// Encrypts the key manager state using the provided cipher. An error is returned if the state is already
+    /// encrypted.
     async fn apply_encryption(&self, cipher: Aes256Gcm) -> Result<(), KeyManagerServiceError>;
 
+    /// Decrypts the key manager state using the provided cipher. An error is returned if the state is not encrypted.
     async fn remove_encryption(&self) -> Result<(), KeyManagerServiceError>;
 
+    /// Gets the next key from the branch. This will auto-increment the branch key index by 1
     async fn get_next_key<T: Into<String> + Send>(&self, branch: T) -> Result<NextKeyResult, KeyManagerServiceError>;
 
+    /// Gets the key at the specified index
     async fn get_key_at_index<T: Into<String> + Send>(
         &self,
         branch: T,
         index: u64,
     ) -> Result<PrivateKey, KeyManagerServiceError>;
 
+    /// Searches the branch to find the index used to generated the key, O(N) where N = index used.
     async fn find_key_index<T: Into<String> + Send>(
         &self,
         branch: T,
         key: &PrivateKey,
     ) -> Result<u64, KeyManagerServiceError>;
 
+    /// Will update the index of the branch if the index given is higher than the current saved index
     async fn update_current_key_index_if_higher<T: Into<String> + Send>(
         &self,
         branch: T,

@@ -37,13 +37,16 @@ use std::{collections::HashMap, sync::Arc};
 
 use crate::key_manager_service::{error::KeyManagerServiceError, storage::database::KeyManagerState};
 
+/// Testing Mock for the key manager service
+/// Contains all functionality of the normal key manager service except persistent storage
 #[derive(Clone)]
 pub struct KeyManagerMock {
     key_managers: Arc<RwLock<HashMap<String, KeyManager<PrivateKey, KeyDigest>>>>,
     master_seed: CipherSeed,
 }
 
 impl KeyManagerMock {
+    /// Creates a new testing mock key manager service
     pub fn new(master_seed: CipherSeed) -> Self {
         KeyManagerMock {
             key_managers: Arc::new(RwLock::new(HashMap::new())),
@@ -53,6 +56,7 @@ impl KeyManagerMock {
 }
 
 impl KeyManagerMock {
+    /// Adds a new branch for the key manager mock to track
     pub async fn add_key_manager_mock(&self, branch: String) -> Result<AddResult, KeyManagerServiceError> {
         let result = if self.key_managers.read().await.contains_key(&branch) {
             AddResult::AlreadyExists
@@ -75,6 +79,7 @@ impl KeyManagerMock {
         Ok(result)
     }
 
+    /// Gets the next key in the branch and increments the index
     pub async fn get_next_key_mock(&self, branch: String) -> Result<NextKeyResult, KeyManagerServiceError> {
         let mut lock = self.key_managers.write().await;
         let km = lock.get_mut(&branch).ok_or(KeyManagerServiceError::UnknownKeyBranch)?;
@@ -85,6 +90,7 @@ impl KeyManagerMock {
         })
     }
 
+    /// get the key at the request index for the branch
     pub async fn get_key_at_index_mock(
         &self,
         branch: String,

@@ -24,14 +24,15 @@ use aes_gcm::Aes256Gcm;
 use crate::key_manager_service::{error::KeyManagerStorageError, storage::database::KeyManagerState};
 
 /// This trait defines the required behaviour that a storage backend must provide for the Key Manager service.
-/// Data is passed to and from the backend via the [DbKey], [DbValue], and [DbValueKey] enums. If new data types are
-/// required to be supported by the backends then these enums can be updated to reflect this requirement and the trait
-/// will remain the same
 pub trait KeyManagerBackend: Send + Sync + Clone {
+    /// This will retrieve the key manager specified by the branch string, None is returned if the key manager is not
+    /// found for the branch.
     fn get_key_manager(&self, branch: String) -> Result<Option<KeyManagerState>, KeyManagerStorageError>;
+    /// This will add an additional branch for the key manager to track.
     fn add_key_manager(&self, key_manager: KeyManagerState) -> Result<(), KeyManagerStorageError>;
+    /// This will increase the key index of the specified branch, and returns an error if the branch does not exist.
     fn increment_key_index(&self, branch: String) -> Result<(), KeyManagerStorageError>;
-    /// This method will set the currently stored key index for the key manager
+    /// This method will set the currently stored key index for the key manager.
     fn set_key_index(&self, branch: String, index: u64) -> Result<(), KeyManagerStorageError>;
     /// Apply encryption to the backend.
     fn apply_encryption(&self, cipher: Aes256Gcm) -> Result<(), KeyManagerStorageError>;

@@ -45,10 +45,13 @@ pub struct KeyManagerDatabase<T> {
 impl<T> KeyManagerDatabase<T>
 where T: KeyManagerBackend + 'static
 {
+    /// Creates a new [KeyManagerDatabase] linked to the provided KeyManagerBackend
     pub fn new(db: T) -> Self {
         Self { db: Arc::new(db) }
     }
 
+    /// Retrieves the key manager state of the provided branch
+    /// Returns None if the request branch does not exist.
     pub async fn get_key_manager_state(
         &self,
         branch: String,
@@ -60,6 +63,7 @@ where T: KeyManagerBackend + 'static
             .and_then(|inner_result| inner_result)
     }
 
+    /// Saves the specified key manager state to the backend database.
     pub async fn set_key_manager_state(&self, state: KeyManagerState) -> Result<(), KeyManagerStorageError> {
         let db_clone = self.db.clone();
         tokio::task::spawn_blocking(move || db_clone.add_key_manager(state))
@@ -69,6 +73,8 @@ where T: KeyManagerBackend + 'static
         Ok(())
     }
 
+    /// Increment the key index of the provided branch of the key manager.
+    /// Will error if the branch does not exist.
     pub async fn increment_key_index(&self, branch: String) -> Result<(), KeyManagerStorageError> {
         let db_clone = self.db.clone();
         tokio::task::spawn_blocking(move || db_clone.increment_key_index(branch))
@@ -77,6 +83,8 @@ where T: KeyManagerBackend + 'static
         Ok(())
     }
 
+    /// Sets the key index of the provided branch of the key manager.
+    /// Will error if the branch does not exist.
     pub async fn set_key_index(&self, branch: String, index: u64) -> Result<(), KeyManagerStorageError> {
         let db_clone = self.db.clone();
         tokio::task::spawn_blocking(move || db_clone.set_key_index(branch, index))
@@ -85,6 +93,8 @@ where T: KeyManagerBackend + 'static
         Ok(())
     }
 
+    /// Encrypts the entire key manager with all branches.
+    /// This will only encrypt the index used, as the master seed phrase is not directly stored with the key manager.
     pub async fn apply_encryption(&self, cipher: Aes256Gcm) -> Result<(), KeyManagerStorageError> {
         let db_clone = self.db.clone();
         tokio::task::spawn_blocking(move || db_clone.apply_encryption(cipher))
@@ -93,6 +103,7 @@ where T: KeyManagerBackend + 'static
             .and_then(|inner_result| inner_result)
     }
 
+    /// Decrypts the entire key manager.
     pub async fn remove_encryption(&self) -> Result<(), KeyManagerStorageError> {
         let db_clone = self.db.clone();
         tokio::task::spawn_blocking(move || db_clone.remove_encryption())

@@ -35,6 +35,7 @@ use crate::{
     },
 };
 
+/// Represents a row in the key_manager_states table.
 #[derive(Clone, Debug, Queryable, Identifiable)]
 #[table_name = "key_manager_states"]
 #[primary_key(id)]
@@ -45,6 +46,7 @@ pub struct KeyManagerStateSql {
     pub timestamp: NaiveDateTime,
 }
 
+/// Struct used to create a new Key manager in the database
 #[derive(Clone, Debug, Insertable)]
 #[table_name = "key_manager_states"]
 pub struct NewKeyManagerStateSql {
@@ -76,6 +78,7 @@ impl TryFrom<KeyManagerStateSql> for KeyManagerState {
 }
 
 impl NewKeyManagerStateSql {
+    /// Commits a new key manager into the database
     pub fn commit(&self, conn: &SqliteConnection) -> Result<(), KeyManagerStorageError> {
         diesel::insert_into(key_manager_states::table)
             .values(self.clone())
@@ -85,17 +88,22 @@ impl NewKeyManagerStateSql {
 }
 
 impl KeyManagerStateSql {
+    /// Retrieve every key manager branch currently in the database.
+    /// Returns a `Vec` of [KeyManagerStateSql], if none are found, it will return an empty `Vec`.
     pub fn index(conn: &SqliteConnection) -> Result<Vec<KeyManagerStateSql>, KeyManagerStorageError> {
         Ok(key_manager_states::table.load::<KeyManagerStateSql>(conn)?)
     }
 
+    /// Retrieve the key manager for the provided branch
+    /// Will return Err if the branch does not exist in the database
     pub fn get_state(branch: &str, conn: &SqliteConnection) -> Result<KeyManagerStateSql, KeyManagerStorageError> {
         key_manager_states::table
             .filter(key_manager_states::branch_seed.eq(branch.to_string()))
             .first::<KeyManagerStateSql>(conn)
             .map_err(|_| KeyManagerStorageError::KeyManagerNotInitialized)
     }
 
+    /// Creates or updates the database with the key manager state in this instance.
     pub fn set_state(&self, conn: &SqliteConnection) -> Result<(), KeyManagerStorageError> {
         match KeyManagerStateSql::get_state(&self.branch_seed, conn) {
             Ok(km) => {
@@ -121,6 +129,7 @@ impl KeyManagerStateSql {
         Ok(())
     }
 
+    /// Updates the key index of the of the provided key manager indicated by the id.
     pub fn set_index(id: i32, index: Vec<u8>, conn: &SqliteConnection) -> Result<(), KeyManagerStorageError> {
         let update = KeyManagerStateUpdateSql {
             branch_seed: None,

@@ -25,8 +25,8 @@ use diesel::{prelude::*, SqliteConnection};
 
 use crate::{key_manager_service::error::KeyManagerStorageError, schema::key_manager_states_old};
 
-// This is a temporary migration file to convert existing indexes to new ones.
-// Todo remove at next testnet reset (currently on Dibbler) #testnet_reset
+/// This is a temporary migration file to convert existing indexes to new ones.
+/// Todo remove at next testnet reset (currently on Dibbler) #testnet_reset
 #[derive(Clone, Debug, Queryable, Identifiable)]
 #[table_name = "key_manager_states_old"]
 #[primary_key(id)]
@@ -39,10 +39,12 @@ pub struct KeyManagerStateSqlOld {
 }
 
 impl KeyManagerStateSqlOld {
+    /// Retrieve all key manager states stored in the database.
     pub fn index(conn: &SqliteConnection) -> Result<Vec<KeyManagerStateSqlOld>, KeyManagerStorageError> {
         Ok(key_manager_states_old::table.load::<KeyManagerStateSqlOld>(conn)?)
     }
 
+    /// Deletes all the stored key manager states in the database.
     pub fn delete(conn: &SqliteConnection) -> Result<(), KeyManagerStorageError> {
         diesel::delete(key_manager_states_old::dsl::key_manager_states_old).execute(conn)?;
         Ok(())

@@ -56,6 +56,9 @@ pub struct KeyManagerSqliteDatabase {
 }
 
 impl KeyManagerSqliteDatabase {
+    /// Creates a new sql backend from provided wallet db connection
+    /// * `cipher` is used to encrypt the sensitive fields in the database, if no cipher is provided, the database will
+    ///   not encrypt sensitive fields
     pub fn new(
         database_connection: WalletDbConnection,
         cipher: Option<Aes256Gcm>,