Implement UTXO Return Address RPC command

cli/src/modules/rpc.rs

@@ -269,6 +269,21 @@ impl Rpc {
                 let result = rpc.get_current_block_color_call(None, GetCurrentBlockColorRequest { hash }).await?;
                 self.println(&ctx, result);
             }
+            RpcApiOps::GetUtxoReturnAddress => {
+                if argv.is_empty() || argv.len() != 2 {
+                    return Err(Error::custom("Please specify a txid and a accepting_block_daa_score"));
+                }
+
+                let txid = argv.remove(0);
+                let txid = RpcHash::from_hex(txid.as_str())?;
+
+                let accepting_block_daa_score = argv.remove(0).parse::<u64>()?;
+
+                let result =
+                    rpc.get_utxo_return_address_call(None, GetUtxoReturnAddressRequest { txid, accepting_block_daa_score }).await?;
+
+                self.println(&ctx, result);
+            }
             _ => {
                 tprintln!(ctx, "rpc method exists but is not supported by the cli: '{op_str}'\r\n");
                 return Ok(());

components/consensusmanager/src/session.rs

@@ -4,7 +4,7 @@
 
 use kaspa_consensus_core::{
     acceptance_data::AcceptanceData,
-    api::{BlockCount, BlockValidationFutures, ConsensusApi, ConsensusStats, DynConsensus},
+    api::{BlockCount, BlockValidationFutures, ConsensusApi, ConsensusStats, DynConsensus, ReturnAddress},
     block::Block,
     blockstatus::BlockStatus,
     daa_score_timestamp::DaaScoreTimestamp,
@@ -313,6 +313,10 @@ impl ConsensusSessionOwned {
         self.clone().spawn_blocking(|c| c.get_chain_block_samples()).await
     }
 
+    pub async fn async_get_utxo_return_script_public_key(&self, txid: Hash, accepting_block_daa_score: u64) -> ReturnAddress {
+        self.clone().spawn_blocking(move |c| c.get_utxo_return_address(txid, accepting_block_daa_score)).await
+    }
+
     /// Returns the antipast of block `hash` from the POV of `context`, i.e. `antipast(hash) ∩ past(context)`.
     /// Since this might be an expensive operation for deep blocks, we allow the caller to specify a limit
     /// `max_traversal_allowed` on the maximum amount of blocks to traverse for obtaining the answer

consensus/core/src/api/mod.rs

@@ -1,6 +1,10 @@
 use futures_util::future::BoxFuture;
+use kaspa_addresses::Address;
 use kaspa_muhash::MuHash;
-use std::sync::Arc;
+use std::{
+    fmt::{Display, Formatter},
+    sync::Arc,
+};
 
 use crate::{
     acceptance_data::AcceptanceData,
@@ -43,6 +47,31 @@ pub struct BlockValidationFutures {
     pub virtual_state_task: BlockValidationFuture,
 }
 
+#[derive(Debug, Clone)]
+pub enum ReturnAddress {
+    Found(Address),
+    AlreadyPruned,
+    TxFromCoinbase,
+    NoTxAtScore,
+    NonStandard,
+    NotFound(String),
+}
+
+impl Display for ReturnAddress {
+    #[inline]
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        let s = match self {
+            ReturnAddress::AlreadyPruned => "Transaction is already pruned".to_string(),
+            ReturnAddress::NoTxAtScore => "Transaction not found at given accepting daa score".to_string(),
+            ReturnAddress::NonStandard => "Transaction was found but not standard".to_string(),
+            ReturnAddress::TxFromCoinbase => "Transaction return address is coinbase".to_string(),
+            ReturnAddress::NotFound(reason) => format!("Transaction return address not found: {}", reason),
+            ReturnAddress::Found(address) => address.to_string(),
+        };
+        f.write_str(&s)
+    }
+}
+
 /// Abstracts the consensus external API
 #[allow(unused_variables)]
 pub trait ConsensusApi: Send + Sync {
@@ -170,6 +199,10 @@ pub trait ConsensusApi: Send + Sync {
         unimplemented!()
     }
 
+    fn get_utxo_return_address(&self, txid: Hash, daa_score: u64) -> ReturnAddress {
+        unimplemented!()
+    }
+
     fn get_virtual_parents(&self) -> BlockHashSet {
         unimplemented!()
     }

consensus/src/consensus/mod.rs

@@ -46,7 +46,7 @@ use kaspa_consensus_core::{
     api::{
         args::{TransactionValidationArgs, TransactionValidationBatchArgs},
         stats::BlockCount,
-        BlockValidationFutures, ConsensusApi, ConsensusStats,
+        BlockValidationFutures, ConsensusApi, ConsensusStats, ReturnAddress,
     },
     block::{Block, BlockTemplate, TemplateBuildMode, TemplateTransactionSelector, VirtualStateApproxId},
     blockhash::BlockHashExtensions,
@@ -77,10 +77,11 @@ use crossbeam_channel::{
 use itertools::Itertools;
 use kaspa_consensusmanager::{SessionLock, SessionReadGuard};
 
+use kaspa_core::{trace, warn};
 use kaspa_database::prelude::StoreResultExtensions;
 use kaspa_hashes::Hash;
 use kaspa_muhash::MuHash;
-use kaspa_txscript::caches::TxScriptCacheCounters;
+use kaspa_txscript::{caches::TxScriptCacheCounters, extract_script_pub_key_address};
 
 use std::{
     cmp::Reverse,
@@ -102,6 +103,9 @@ use crate::model::stores::selected_chain::SelectedChainStoreReader;
 
 use std::cmp;
 
+use crate::model::stores::utxo_diffs::UtxoDiffsStoreReader;
+use kaspa_consensus_core::utxo::utxo_diff::ImmutableUtxoDiff;
+
 pub struct Consensus {
     // DB
     db: Arc<DB>,
@@ -696,6 +700,120 @@ impl ConsensusApi for Consensus {
         sample_headers
     }
 
+    fn get_utxo_return_address(&self, txid: Hash, target_daa_score: u64) -> ReturnAddress {
+        // We need consistency between the past pruning points, selected chain and header store reads
+        let _guard = self.pruning_lock.blocking_read();
+
+        let sc_read = self.selected_chain_store.read();
+
+        let source_hash = self.get_source();
+
+        // Pruning Point hash is always expected to be in get_compact_header_data so unwrap should never fail
+        if target_daa_score < self.headers_store.get_compact_header_data(source_hash).unwrap().daa_score {
+            // Early exit if target daa score is lower than that of pruning point's daa score:
+            return ReturnAddress::AlreadyPruned;
+        }
+
+        let source_index = sc_read.get_by_hash(source_hash).unwrap();
+        let (tip_index, tip_hash) = sc_read.get_tip().unwrap();
+        let tip_daa_score = self.headers_store.get_compact_header_data(tip_hash).unwrap().daa_score;
+
+        let mut low_index = tip_index.saturating_sub(tip_daa_score.saturating_sub(target_daa_score)).max(source_index);
+        let mut high_index = tip_index;
+
+        let matching_chain_block_hash = loop {
+            // Binary search for the chain block that matches the target_daa_score
+            // 0. Get the mid point index
+            let mid = low_index + (high_index - low_index) / 2;
+
+            // 1. Get the chain block hash at that index. Error if we don't find a hash at an index
+            let hash = match sc_read.get_by_index(mid) {
+                Ok(hash) => hash,
+                Err(_) => {
+                    trace!("Did not find a hash at index {}", mid);
+                    return ReturnAddress::NotFound(format!("Did not find a hash at index {}", mid));
+                }
+            };
+
+            // 2. Get the compact header so we have access to the daa_score. Error if we
+            let compact_header = match self.headers_store.get_compact_header_data(hash) {
+                Ok(compact_header) => compact_header,
+                Err(_) => {
+                    trace!("Did not find a compact header with hash {}", hash);
+                    return ReturnAddress::NotFound(format!("Did not find a compact header with hash {}", hash));
+                }
+            };
+
+            // 3. Compare block daa score to our target
+            match compact_header.daa_score.cmp(&target_daa_score) {
+                cmp::Ordering::Equal => {
+                    // We found the chain block we need
+                    break hash;
+                }
+                cmp::Ordering::Greater => {
+                    high_index = mid - 1;
+                }
+                cmp::Ordering::Less => {
+                    low_index = mid + 1;
+                }
+            }
+
+            if low_index > high_index {
+                return ReturnAddress::NoTxAtScore;
+            }
+        };
+
+        let acceptance_data = match self.acceptance_data_store.get(matching_chain_block_hash) {
+            Ok(acceptance_data) => acceptance_data,
+            Err(_) => {
+                return ReturnAddress::NotFound("Did not find acceptance data".to_string());
+            }
+        };
+        let (index, containing_acceptance) = match acceptance_data.iter().find_map(|mbad| {
+            let tx_arr_index =
+                mbad.accepted_transactions.iter().find_map(|tx| (tx.transaction_id == txid).then_some(tx.index_within_block as usize));
+            tx_arr_index.map(|index| (index, mbad.clone()))
+        }) {
+            Some((index, containing_acceptance)) => (index, containing_acceptance),
+            None => {
+                return ReturnAddress::NotFound("Did not find containing_acceptance".to_string());
+            }
+        };
+
+        // Found Merged block containing the TXID
+        let tx = &self.block_transactions_store.get(containing_acceptance.block_hash).unwrap()[index];
+
+        if tx.id() != txid {
+            // Should never happen, but do a sanity check. This would mean something went wrong with storing block transactions
+            // Sanity check is necessary to guarantee that this function will never give back a wrong address (err on the side of NotFound)
+            warn!("Expected {} to match {} when checking block_transaction_store using array index of transaction", tx.id(), txid);
+            return ReturnAddress::NotFound(format!(
+                "Expected {} to match {} when checking block_transaction_store using array index of transaction",
+                tx.id(),
+                txid
+            ));
+        }
+
+        if tx.inputs.is_empty() {
+            // A transaction may have no inputs (like a coinbase transaction)
+            return ReturnAddress::TxFromCoinbase;
+        }
+
+        let first_input_prev_outpoint = &tx.inputs[0].previous_outpoint;
+        // Expected to never fail, since we found the acceptance data and therefore there must be matching diff
+        let utxo_diff = self.utxo_diffs_store.get(matching_chain_block_hash).unwrap();
+        let removed_diffs = utxo_diff.removed();
+
+        if let Ok(address) = extract_script_pub_key_address(
+            &removed_diffs.get(first_input_prev_outpoint).unwrap().script_public_key,
+            self.config.prefix(),
+        ) {
+            ReturnAddress::Found(address)
+        } else {
+            ReturnAddress::NonStandard
+        }
+    }
+
     fn get_virtual_parents(&self) -> BlockHashSet {
         self.lkg_virtual_state.load().parents.iter().copied().collect()
     }

rpc/core/src/api/ops.rs

@@ -136,6 +136,8 @@ pub enum RpcApiOps {
     GetFeeEstimateExperimental = 148,
     /// Block color determination by iterating DAG.
     GetCurrentBlockColor = 149,
+    /// Get UTXO Return Addresses
+    GetUtxoReturnAddress = 150,
 }
 
 impl RpcApiOps {

rpc/core/src/api/rpc.rs

@@ -438,6 +438,18 @@ pub trait RpcApi: Sync + Send + AnySync {
         request: GetDaaScoreTimestampEstimateRequest,
     ) -> RpcResult<GetDaaScoreTimestampEstimateResponse>;
 
+    async fn get_utxo_return_address(&self, txid: RpcHash, accepting_block_daa_score: u64) -> RpcResult<Option<RpcAddress>> {
+        Ok(self
+            .get_utxo_return_address_call(None, GetUtxoReturnAddressRequest { txid, accepting_block_daa_score })
+            .await?
+            .return_address)
+    }
+    async fn get_utxo_return_address_call(
+        &self,
+        _connection: Option<&DynRpcConnection>,
+        request: GetUtxoReturnAddressRequest,
+    ) -> RpcResult<GetUtxoReturnAddressResponse>;
+
     // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
     // Fee estimation API
 

rpc/core/src/error.rs

@@ -2,7 +2,7 @@
 //! [`RpcError`] enum used by RPC primitives.
 //!
 
-use kaspa_consensus_core::{subnets::SubnetworkConversionError, tx::TransactionId};
+use kaspa_consensus_core::{api::ReturnAddress, subnets::SubnetworkConversionError, tx::TransactionId};
 use kaspa_utils::networking::IpAddress;
 use std::{net::AddrParseError, num::TryFromIntError};
 use thiserror::Error;
@@ -134,6 +134,9 @@ pub enum RpcError {
 
     #[error(transparent)]
     ConsensusClient(#[from] kaspa_consensus_client::error::Error),
+
+    #[error("utxo return address could not be found -> {0}")]
+    UtxoReturnAddressNotFound(ReturnAddress),
 }
 
 impl From<String> for RpcError {

rpc/core/src/model/message.rs

@@ -2666,6 +2666,69 @@ impl Deserializer for GetCurrentBlockColorResponse {
     }
 }
 
+#[derive(Clone, Debug, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct GetUtxoReturnAddressRequest {
+    pub txid: RpcHash,
+    pub accepting_block_daa_score: u64,
+}
+
+impl GetUtxoReturnAddressRequest {
+    pub fn new(txid: RpcHash, accepting_block_daa_score: u64) -> Self {
+        Self { txid, accepting_block_daa_score }
+    }
+}
+
+impl Serializer for GetUtxoReturnAddressRequest {
+    fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
+        store!(u16, &1, writer)?;
+        store!(RpcHash, &self.txid, writer)?;
+        store!(u64, &self.accepting_block_daa_score, writer)?;
+
+        Ok(())
+    }
+}
+
+impl Deserializer for GetUtxoReturnAddressRequest {
+    fn deserialize<R: std::io::Read>(reader: &mut R) -> std::io::Result<Self> {
+        let _version = load!(u16, reader)?;
+        let txid = load!(RpcHash, reader)?;
+        let accepting_block_daa_score = load!(u64, reader)?;
+
+        Ok(Self { txid, accepting_block_daa_score })
+    }
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct GetUtxoReturnAddressResponse {
+    pub return_address: Option<RpcAddress>,
+}
+
+impl GetUtxoReturnAddressResponse {
+    pub fn new(return_address: Option<RpcAddress>) -> Self {
+        Self { return_address }
+    }
+}
+
+impl Serializer for GetUtxoReturnAddressResponse {
+    fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
+        store!(u16, &1, writer)?;
+        store!(Option<RpcAddress>, &self.return_address, writer)?;
+
+        Ok(())
+    }
+}
+
+impl Deserializer for GetUtxoReturnAddressResponse {
+    fn deserialize<R: std::io::Read>(reader: &mut R) -> std::io::Result<Self> {
+        let _version = load!(u16, reader)?;
+        let return_address = load!(Option<RpcAddress>, reader)?;
+
+        Ok(Self { return_address })
+    }
+}
+
 // ----------------------------------------------------------------------------
 // Subscriptions & notifications
 // ----------------------------------------------------------------------------

rpc/grpc/client/src/lib.rs

@@ -277,6 +277,7 @@ impl RpcApi for GrpcClient {
     route!(get_fee_estimate_call, GetFeeEstimate);
     route!(get_fee_estimate_experimental_call, GetFeeEstimateExperimental);
     route!(get_current_block_color_call, GetCurrentBlockColor);
+    route!(get_utxo_return_address_call, GetUtxoReturnAddress);
 
     // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
     // Notification API

rpc/grpc/core/proto/messages.proto

@@ -65,6 +65,7 @@ message KaspadRequest {
     GetFeeEstimateRequestMessage getFeeEstimateRequest = 1106;
     GetFeeEstimateExperimentalRequestMessage getFeeEstimateExperimentalRequest = 1108;
     GetCurrentBlockColorRequestMessage getCurrentBlockColorRequest = 1110;
+    GetUtxoReturnAddressRequestMessage GetUtxoReturnAddressRequest = 1112;
   }
 }
 
@@ -130,6 +131,7 @@ message KaspadResponse {
     GetFeeEstimateResponseMessage getFeeEstimateResponse = 1107;
     GetFeeEstimateExperimentalResponseMessage getFeeEstimateExperimentalResponse = 1109;
     GetCurrentBlockColorResponseMessage getCurrentBlockColorResponse = 1111;
+    GetUtxoReturnAddressResponseMessage GetUtxoReturnAddressResponse = 1113;
   }
 }