From 290bf4b9696b9641a4d8f892534a604e62585d63 Mon Sep 17 00:00:00 2001
From: Hansie Odendaal <pluto@tari.com>
Date: Wed, 31 Jul 2024 13:32:07 +0200
Subject: [PATCH] wip

---
 .../src/output_manager_service/service.rs     | 74 ++++++++++++-------
 1 file changed, 48 insertions(+), 26 deletions(-)

diff --git a/base_layer/wallet/src/output_manager_service/service.rs b/base_layer/wallet/src/output_manager_service/service.rs
index b1ba319885..81d2c56eec 100644
--- a/base_layer/wallet/src/output_manager_service/service.rs
+++ b/base_layer/wallet/src/output_manager_service/service.rs
@@ -68,7 +68,7 @@ use tari_core::{
     },
 };
 use tari_crypto::{keys::SecretKey, ristretto::pedersen::PedersenCommitment};
-use tari_key_manager::key_manager_service::KeyAndId;
+use tari_key_manager::key_manager_service::{KeyAndId, KeyId};
 use tari_script::{
     inputs,
     push_pubkey_script,
@@ -1168,6 +1168,32 @@ where
         Ok((tx_id, stp.into_transaction()?))
     }
 
+    async fn pre_mine_script_key_from_payment_id(
+        &self,
+        payment_id: PaymentId,
+        tx_id: TxId,
+    ) -> Result<KeyAndId<PublicKey>, OutputManagerError> {
+        if let PaymentId::U64(index) = payment_id {
+            let script_key_id = KeyId::Managed {
+                branch: TransactionKeyManagerBranch::PreMine.get_branch_key(),
+                index,
+            };
+            Ok(KeyAndId::<PublicKey> {
+                pub_key: self
+                    .resources
+                    .key_manager
+                    .get_public_key_at_key_id(&script_key_id)
+                    .await?,
+                key_id: script_key_id,
+            })
+        } else {
+            return Err(OutputManagerError::ServiceError(format!(
+                "Invalid payment id (TxId: {}): expected 'PaymentId::U64(_)', received {:?}",
+                tx_id, payment_id
+            )));
+        }
+    }
+
     /// Create a partial transaction in order to prepare output
     #[allow(clippy::too_many_lines)]
     #[allow(clippy::mutable_key_type)]
@@ -1187,7 +1213,7 @@ where
         maturity: u64,
         range_proof_type: RangeProofType,
         minimum_value_promise: MicroMinotari,
-        script_key_id: TariKeyId,
+        _script_key_id: TariKeyId,
     ) -> Result<
         (
             Transaction,
@@ -1200,14 +1226,6 @@ where
         OutputManagerError,
     > {
         trace!(target: LOG_TARGET, "encumber_aggregate_utxo: start");
-        let script_key = KeyAndId::<PublicKey> {
-            pub_key: self
-                .resources
-                .key_manager
-                .get_public_key_at_key_id(&script_key_id)
-                .await?,
-            key_id: script_key_id,
-        };
         // Fetch the output from the blockchain
         let output = self
             .fetch_unspent_outputs_from_node(vec![output_hash])
@@ -1227,20 +1245,7 @@ where
         }
         trace!(target: LOG_TARGET, "encumber_aggregate_utxo: fetched outputs");
         // Retrieve the list of n public keys from the script
-        let (multi_sig_public_keys, threshold) = if let Some(Opcode::CheckMultiSigVerifyAggregatePubKey(
-            m,
-            _n,
-            keys,
-            _msg,
-        )) = output.script.as_slice().get(3)
-        {
-            (keys.clone(), m)
-        } else {
-            return Err(OutputManagerError::ServiceError(format!(
-                "Invalid script (TxId: {})",
-                tx_id
-            )));
-        };
+        let (multi_sig_public_keys, threshold) = get_multi_sig_script_components(&output.script, tx_id)?;
         trace!(target: LOG_TARGET, "encumber_aggregate_utxo: retrieved public keys from script");
         // Create a deterministic encryption key from the sum of the public keys
         let sum_public_keys = multi_sig_public_keys
@@ -1254,6 +1259,9 @@ where
             EncryptedData::decrypt_data(&encryption_private_key, &output.commitment, &output.encrypted_data)
         {
             if output.verify_mask(&self.resources.factories.range_proof, &commitment_mask, amount.as_u64())? {
+                let script_key = self
+                    .pre_mine_script_key_from_payment_id(payment_id.clone(), tx_id)
+                    .await?;
                 let mut script_signatures = Vec::new();
                 // lets add our own signature to the list
                 let self_signature = self
@@ -1274,7 +1282,7 @@ where
                         }
                     }
                 }
-                if script_signatures.len() != usize::from(*threshold) {
+                if script_signatures.len() != usize::from(threshold) {
                     return Err(OutputManagerError::ServiceError(format!(
                         "Invalid number of signatures (TxId: {}), expected {}, received {}",
                         tx_id,
@@ -1290,7 +1298,7 @@ where
                     output.features,
                     output.script,
                     ExecutionStack::new(script_signatures),
-                    script_key.key_id, // Only of the master wallet
+                    script_key.key_id.clone(), // Only of the master wallet
                     output.sender_offset_public_key,
                     output.metadata_signature,
                     0,
@@ -2881,6 +2889,20 @@ where
     }
 }
 
+fn get_multi_sig_script_components(
+    script: &TariScript,
+    tx_id: TxId,
+) -> Result<(Vec<PublicKey>, u8), OutputManagerError> {
+    if let Some(Opcode::CheckMultiSigVerifyAggregatePubKey(m, _n, keys, _msg)) = script.as_slice().get(3) {
+        Ok((keys.clone(), *m))
+    } else {
+        return Err(OutputManagerError::ServiceError(format!(
+            "Invalid script (TxId: {})",
+            tx_id
+        )));
+    }
+}
+
 fn service_error_with_id(tx_id: TxId, err: String, log_error: bool) -> OutputManagerError {
     let err_str = format!("TxId: {} ({})", tx_id, err);
     if log_error {