change how key_manager uses nonce key types

tari-project · Aug 12, 2024 · b8318a6 · b8318a6
1 parent 84e7c0a
commit b8318a6
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 21 deletions.
diff --git a/base_layer/common_types/src/key_branches.rs b/base_layer/common_types/src/key_branches.rs
@@ -42,15 +42,15 @@ pub enum TransactionKeyManagerBranch {
     PreMine = Branch::PreMine as u8,
 }
 
-const DATA_ENCRYPTION: &str = "data encryption";
-const METADATA_EPHEMERAL_NONCE: &str = "metadata ephemeral nonce";
-const COMMITMENT_MASK: &str = "commitment mask";
-const NONCE: &str = "nonce";
-const KERNEL_NONCE: &str = "kernel nonce";
-const SENDER_OFFSET: &str = "sender offset";
-const ONE_SIDED_SENDER_OFFSET: &str = "one sided sender offset";
-const RANDOM_KEY: &str = "random key";
-const PRE_MINE: &str = "pre-mine";
+pub const DATA_ENCRYPTION: &str = "data encryption";
+pub const METADATA_EPHEMERAL_NONCE: &str = "metadata ephemeral nonce";
+pub const COMMITMENT_MASK: &str = "commitment mask";
+pub const NONCE: &str = "nonce";
+pub const KERNEL_NONCE: &str = "kernel nonce";
+pub const SENDER_OFFSET: &str = "sender offset";
+pub const ONE_SIDED_SENDER_OFFSET: &str = "one sided sender offset";
+pub const RANDOM_KEY: &str = "random key";
+pub const PRE_MINE: &str = "pre-mine";
 
 impl TransactionKeyManagerBranch {
     /// Warning: Changing these strings will affect the backwards compatibility of the wallet with older databases or

diff --git a/base_layer/core/src/transactions/key_manager/inner.rs b/base_layer/core/src/transactions/key_manager/inner.rs
@@ -35,12 +35,18 @@ use minotari_ledger_wallet_comms::accessor_methods::{
     ledger_get_script_signature,
     ScriptSignatureKey,
 };
-use rand::rngs::OsRng;
-#[cfg(feature = "ledger")]
-use rand::RngCore;
+use rand::{rngs::OsRng, RngCore};
 use strum::IntoEnumIterator;
 use tari_common_types::{
-    key_branches::TransactionKeyManagerBranch,
+    key_branches::{
+        TransactionKeyManagerBranch,
+        KERNEL_NONCE,
+        METADATA_EPHEMERAL_NONCE,
+        NONCE,
+        ONE_SIDED_SENDER_OFFSET,
+        RANDOM_KEY,
+        SENDER_OFFSET,
+    },
     tari_address::TariAddress,
     types::{ComAndPubSignature, Commitment, PrivateKey, PublicKey, RangeProof, Signature},
     wallet_types::WalletType,
@@ -172,14 +178,24 @@ where TBackend: KeyManagerBackend<PublicKey> + 'static
 
     pub async fn get_next_key(&self, branch: &str) -> Result<KeyAndId<PublicKey>, KeyManagerServiceError> {
         let index = {
-            let mut km = self
-                .key_managers
-                .get(branch)
-                .ok_or_else(|| self.unknown_key_branch_error("get_next_key", branch))?
-                .write()
-                .await;
-            self.db.increment_key_index(branch)?;
-            km.increment_key_index(1)
+            match branch {
+                METADATA_EPHEMERAL_NONCE |
+                NONCE |
+                KERNEL_NONCE |
+                SENDER_OFFSET |
+                ONE_SIDED_SENDER_OFFSET |
+                RANDOM_KEY => OsRng.next_u64(),
+                _ => {
+                    let mut km = self
+                        .key_managers
+                        .get(branch)
+                        .ok_or_else(|| self.unknown_key_branch_error("get_next_key", branch))?
+                        .write()
+                        .await;
+                    self.db.increment_key_index(branch)?;
+                    km.increment_key_index(1)
+                },
+            }
         };
         let key_id = KeyId::Managed {
             branch: branch.to_string(),