fix: coinbase output recovery bug (#3789)

Description --- A bug was discovered when recovering Coinbase outputs that was revealed when the Key Manager was fixed to actually use branch seeds. Coinbase output keys are derived on a separate key manager branch so when the UTXO scanner tried to update the key manager index it would not find the coinabse key in the main spending key branch which caused an error. This PR updates the UTXO scanner to check if a recovered output has the coinbase flag or not and then searches on the Coinbase branch when looking for that outputs key index and script key. How Has This Been Tested? --- The PR updates the `Wallet recovery with connected base node staying online` cucumber test to include recovering some coinbase outputs. The PR also updates the `Multiple Wallet recovery from seed node` cucumber test to work now that the coinbase issue is fixed. The test is also updated so that it recovered N distinct wallets to the same seed node where as before it would create N wallet with the same seed words and thus network identity.
tari-project · Feb 3, 2022 · beb299e · beb299e
1 parent 3495e85
commit beb299e
Show file tree

Hide file tree

Showing 10 changed files with 288 additions and 114 deletions.
diff --git a/base_layer/core/src/transactions/transaction/output_features.rs b/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 {

diff --git a/base_layer/wallet/src/output_manager_service/error.rs b/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)]

diff --git a/base_layer/wallet/src/output_manager_service/master_key_manager.rs b/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();

diff --git a/base_layer/wallet/src/output_manager_service/recovery/standard_outputs_recoverer.rs b/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(())
     }
 }
diff --git a/integration_tests/features/WalletRecovery.feature b/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

diff --git a/integration_tests/features/support/mining_node_steps.js b/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) {

diff --git a/integration_tests/features/support/node_steps.js b/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);
   }
 );