fix: coinbase output recovery bug

base_layer/core/src/transactions/transaction/output_features.rs

@@ -209,6 +209,10 @@ impl OutputFeatures {
     pub fn is_non_fungible_burn(&self) -> bool {
         self.flags.contains(OutputFlags::BURN_NON_FUNGIBLE)
     }
+
+    pub fn is_coinbase(&self) -> bool {
+        self.flags.contains(OutputFlags::COINBASE_OUTPUT)
+    }
 }
 
 impl ConsensusEncoding for OutputFeatures {

base_layer/wallet/src/output_manager_service/error.rs

@@ -120,6 +120,8 @@ pub enum OutputManagerError {
     },
     #[error("Invalid message received:{0}")]
     InvalidMessageError(String),
+    #[error("Operation not support on this Key Manager branch")]
+    KeyManagerBranchNotSupported,
 }
 
 #[derive(Debug, Error)]

base_layer/wallet/src/output_manager_service/master_key_manager.rs

@@ -20,6 +20,8 @@
 // WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
 // USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+use std::fmt::{Display, Error, Formatter};
+
 use futures::lock::Mutex;
 use log::*;
 use tari_common_types::types::{PrivateKey, PublicKey};
@@ -41,14 +43,32 @@ use crate::{
 };
 
 const LOG_TARGET: &str = "wallet::output_manager_service::master_key_manager";
-
-const KEY_MANAGER_COINBASE_BRANCH_KEY: &str = "coinbase";
-const KEY_MANAGER_COINBASE_SCRIPT_BRANCH_KEY: &str = "coinbase_script";
-const KEY_MANAGER_SCRIPT_BRANCH_KEY: &str = "script";
-const KEY_MANAGER_RECOVERY_VIEWONLY_BRANCH_KEY: &str = "recovery_viewonly";
-const KEY_MANAGER_RECOVERY_BLINDING_BRANCH_KEY: &str = "recovery_blinding";
 const KEY_MANAGER_MAX_SEARCH_DEPTH: u64 = 1_000_000;
 
+#[derive(Clone, Copy)]
+pub enum KeyManagerBranch {
+    Spend,
+    SpendScript,
+    Coinbase,
+    CoinbaseScript,
+    RecoveryViewOnly,
+    RecoveryBlinding,
+}
+
+impl Display for KeyManagerBranch {
+    fn fmt(&self, fmt: &mut Formatter<'_>) -> Result<(), Error> {
+        let response = match self {
+            KeyManagerBranch::Spend => "",
+            KeyManagerBranch::SpendScript => "script",
+            KeyManagerBranch::Coinbase => "coinbase",
+            KeyManagerBranch::CoinbaseScript => "coinbase_script",
+            KeyManagerBranch::RecoveryViewOnly => "recovery_viewonly",
+            KeyManagerBranch::RecoveryBlinding => "recovery_blinding",
+        };
+        fmt.write_str(response)
+    }
+}
+
 pub(crate) struct MasterKeyManager<TBackend> {
     utxo_key_manager: Mutex<KeyManager<PrivateKey, KeyDigest>>,
     utxo_script_key_manager: Mutex<KeyManager<PrivateKey, KeyDigest>>,
@@ -67,7 +87,7 @@ where TBackend: OutputManagerBackend + 'static
             None => {
                 let starting_state = KeyManagerState {
                     seed: master_seed,
-                    branch_seed: "".to_string(),
+                    branch_seed: KeyManagerBranch::Spend.to_string(),
                     primary_key_index: 0,
                 };
                 db.set_key_manager_state(starting_state.clone()).await?;
@@ -89,32 +109,32 @@ where TBackend: OutputManagerBackend + 'static
 
         let utxo_script_key_manager = KeyManager::<PrivateKey, KeyDigest>::from(
             key_manager_state.seed.clone(),
-            KEY_MANAGER_SCRIPT_BRANCH_KEY.to_string(),
+            KeyManagerBranch::SpendScript.to_string(),
             key_manager_state.primary_key_index,
         );
 
         let coinbase_key_manager = KeyManager::<PrivateKey, KeyDigest>::from(
             key_manager_state.seed.clone(),
-            KEY_MANAGER_COINBASE_BRANCH_KEY.to_string(),
+            KeyManagerBranch::Coinbase.to_string(),
             0,
         );
 
         let coinbase_script_key_manager = KeyManager::<PrivateKey, KeyDigest>::from(
             key_manager_state.seed.clone(),
-            KEY_MANAGER_COINBASE_SCRIPT_BRANCH_KEY.to_string(),
+            KeyManagerBranch::CoinbaseScript.to_string(),
             0,
         );
 
         let rewind_key_manager = KeyManager::<PrivateKey, KeyDigest>::from(
             key_manager_state.seed.clone(),
-            KEY_MANAGER_RECOVERY_VIEWONLY_BRANCH_KEY.to_string(),
+            KeyManagerBranch::RecoveryViewOnly.to_string(),
             0,
         );
         let rewind_key = rewind_key_manager.derive_key(0)?.k;
 
         let rewind_blinding_key_manager = KeyManager::<PrivateKey, KeyDigest>::from(
             key_manager_state.seed,
-            KEY_MANAGER_RECOVERY_BLINDING_BRANCH_KEY.to_string(),
+            KeyManagerBranch::RecoveryBlinding.to_string(),
             0,
         );
         let rewind_blinding_key = rewind_blinding_key_manager.derive_key(0)?.k;
@@ -158,6 +178,12 @@ where TBackend: OutputManagerBackend + 'static
         Ok(script_key.k)
     }
 
+    pub async fn get_coinbase_script_key_at_index(&self, index: u64) -> Result<PrivateKey, OutputManagerError> {
+        let skm = self.coinbase_script_key_manager.lock().await;
+        let script_key = skm.derive_key(index)?;
+        Ok(script_key.k)
+    }
+
     pub async fn get_coinbase_spend_and_script_key_for_height(
         &self,
         height: u64,
@@ -185,14 +211,19 @@ where TBackend: OutputManagerBackend + 'static
         }
     }
 
-    /// Search the current key manager key chain to find the index of the specified key.
-    pub async fn find_utxo_key_index(&self, key: PrivateKey) -> Result<u64, OutputManagerError> {
-        let utxo_key_manager = self.utxo_key_manager.lock().await;
-        let current_index = (*utxo_key_manager).key_index();
+    /// Search the specified branch key manager key chain to find the index of the specified key.
+    pub async fn find_key_index(&self, key: PrivateKey, branch: KeyManagerBranch) -> Result<u64, OutputManagerError> {
+        let key_manager = match branch {
+            KeyManagerBranch::Spend => self.utxo_key_manager.lock().await,
+            KeyManagerBranch::Coinbase => self.coinbase_key_manager.lock().await,
+            _ => return Err(OutputManagerError::KeyManagerBranchNotSupported),
+        };
+
+        let current_index = (*key_manager).key_index();
 
         for i in 0u64..current_index + KEY_MANAGER_MAX_SEARCH_DEPTH {
-            if (*utxo_key_manager).derive_key(i)?.k == key {
-                trace!(target: LOG_TARGET, "Key found in Key Chain at index {}", i);
+            if (*key_manager).derive_key(i)?.k == key {
+                trace!(target: LOG_TARGET, "Key found in {} Key Chain at index {}", branch, i);
                 return Ok(i);
             }
         }
@@ -201,7 +232,7 @@ where TBackend: OutputManagerBackend + 'static
     }
 
     /// If the supplied index is higher than the current UTXO key chain indices then they will be updated.
-    pub async fn update_current_index_if_higher(&self, index: u64) -> Result<(), OutputManagerError> {
+    pub async fn update_current_spend_key_index_if_higher(&self, index: u64) -> Result<(), OutputManagerError> {
         let mut utxo_key_manager = self.utxo_key_manager.lock().await;
         let mut utxo_script_key_manager = self.utxo_script_key_manager.lock().await;
         let current_index = (*utxo_key_manager).key_index();

base_layer/wallet/src/output_manager_service/recovery/standard_outputs_recoverer.rs

@@ -37,6 +37,7 @@ use tari_crypto::{
 
 use crate::output_manager_service::{
     error::{OutputManagerError, OutputManagerStorageError},
+    master_key_manager::KeyManagerBranch,
     storage::{
         database::{OutputManagerBackend, OutputManagerDatabase},
         models::DbUnblindedOutput,
@@ -164,18 +165,30 @@ where TBackend: OutputManagerBackend + 'static
         &mut self,
         output: &mut UnblindedOutput,
     ) -> Result<(), OutputManagerError> {
-        let found_index = self
-            .master_key_manager
-            .find_utxo_key_index(output.spending_key.clone())
-            .await?;
-
-        self.master_key_manager
-            .update_current_index_if_higher(found_index)
-            .await?;
-
-        let script_private_key = self.master_key_manager.get_script_key_at_index(found_index).await?;
-        output.input_data = inputs!(PublicKey::from_secret_key(&script_private_key));
-        output.script_private_key = script_private_key;
+        let script_key = if output.features.is_coinbase() {
+            let found_index = self
+                .master_key_manager
+                .find_key_index(output.spending_key.clone(), KeyManagerBranch::Coinbase)
+                .await?;
+
+            self.master_key_manager
+                .get_coinbase_script_key_at_index(found_index)
+                .await?
+        } else {
+            let found_index = self
+                .master_key_manager
+                .find_key_index(output.spending_key.clone(), KeyManagerBranch::Spend)
+                .await?;
+
+            self.master_key_manager
+                .update_current_spend_key_index_if_higher(found_index)
+                .await?;
+
+            self.master_key_manager.get_script_key_at_index(found_index).await?
+        };
+
+        output.input_data = inputs!(PublicKey::from_secret_key(&script_key));
+        output.script_private_key = script_key;
         Ok(())
     }
 }

integration_tests/features/WalletRecovery.feature

@@ -13,41 +13,39 @@ Feature: Wallet Recovery
         When I wait for wallet WALLET_A to have at least 55000000000 uT
         Then all nodes are at height 15
         And I send 200000 uT from wallet WALLET_A to wallet WALLET_B at fee 100
+        And I have mining node MINER_B connected to base node NODE and wallet WALLET_B
+        When mining node MINER_B mines 2 blocks
         When I mine 5 blocks on NODE
-        Then all nodes are at height 20
+        Then all nodes are at height 22
         Then I stop wallet WALLET_B
         When I recover wallet WALLET_B into wallet WALLET_C connected to all seed nodes
-        When I wait for wallet WALLET_C to have at least 200000 uT
+        When I wait for wallet WALLET_C to have at least 10000200000 uT
         And I have wallet WALLET_D connected to all seed nodes
         And I send 100000 uT from wallet WALLET_C to wallet WALLET_D at fee 100
         When I mine 5 blocks on NODE
-        Then all nodes are at height 25
+        Then all nodes are at height 27
         Then I wait for wallet WALLET_D to have at least 100000 uT
 
-    @broken
     Scenario Outline: Multiple Wallet recovery from seed node
         Given I have a seed node NODE
-        And I have wallet WALLET_A connected to all seed nodes
-        And I have mining node MINER connected to base node NODE and wallet WALLET_A
-        When mining node MINER mines 15 blocks
-        When I wait for wallet WALLET_A to have at least 55000000000 uT
-        Then all nodes are at height 15
-        Then I stop wallet WALLET_A
-        When I recover wallet WALLET_A into <NumWallets> wallets connected to all seed nodes
-        When I wait for <NumWallets> wallets to have at least 55000000000 uT
-        # TODO: having multiple wallet with the same network id is problematic, use N separate wallets or ensure that both are not trying to connect to the same base node
-        # Then Wallet WALLET_A and <NumWallets> wallets have the same balance
+        And I have <NumWallets> non-default wallets connected to all seed nodes using DirectAndStoreAndForward
+        And I have individual mining nodes connected to each wallet and base node NODE
+        Then I have each mining node mine 3 blocks
+        Then all nodes are at height 3*<NumWallets>
+        Then I stop all wallets
+        When I recover all wallets connected to all seed nodes
+        Then I wait for recovered wallets to have at least 15000000000 uT
         @critical
         Examples:
             | NumWallets |
-            | 1        |
+            | 4        |
 
         @long-running
         Examples:
             | NumWallets |
-            | 2        |
             | 5        |
             | 10        |
+            | 20        |
 
 
     @critical

integration_tests/features/support/mining_node_steps.js

@@ -103,6 +103,36 @@ Given(
   }
 );
 
+Given(
+  /I have individual mining nodes connected to each wallet and (?:base|seed) node (.*)/,
+  async function (node) {
+    let walletNames = Object.keys(this.wallets);
+    const promises = [];
+    for (let i = 0; i < walletNames.length; i++) {
+      let name = "Miner_" + String(i).padStart(2, "0");
+      promises.push(this.createMiningNode(name, node, walletNames[i]));
+    }
+    await Promise.all(promises);
+  }
+);
+
+Given(
+  /I have each mining node mine (\d+) blocks?$/,
+  { timeout: 1200 * 1000 }, // Must allow many blocks to be mined; dynamic time out below limits actual time
+  async function (numBlocks) {
+    let miningNodes = Object.keys(this.miners);
+    for (let i = 0; i < miningNodes.length; i++) {
+      console.log("MN", miningNodes[i]);
+      const miningNode = this.getMiningNode(miningNodes[i]);
+      await miningNode.init(numBlocks, null, 1, i + 2, false, null);
+      await withTimeout(
+        (10 + parseInt(numBlocks * miningNodes.length) * 1) * 1000,
+        await miningNode.startNew()
+      );
+    }
+  }
+);
+
 Given(
   /I have mine-before-tip mining node (.*) connected to base node (.*) and wallet (.*)/,
   function (miner, node, wallet) {

integration_tests/features/support/node_steps.js

@@ -354,28 +354,15 @@ Then(
   "all nodes are at height {int}",
   { timeout: 800 * 1000 },
   async function (height) {
-    await waitFor(
-      async () => {
-        let result = true;
-        await this.forEachClientAsync(async (client, name) => {
-          await waitFor(
-            async () => await client.getTipHeight(),
-            height,
-            5 * height * 1000 /* 5 seconds per block */
-          );
-          const currTip = await client.getTipHeight();
-          console.log(
-            `Node ${name} is at tip: ${currTip} (should be ${height})`
-          );
-          result = result && currTip == height;
-        });
-        return result;
-      },
-      true,
-      600 * 1000,
-      5 * 1000,
-      5
-    );
+    await this.all_nodes_are_at_height(height);
+  }
+);
+
+Then(
+  "all nodes are at height {int}*{int}",
+  { timeout: 800 * 1000 },
+  async function (a, b) {
+    await this.all_nodes_are_at_height(a * b);
   }
 );