From 2e6622ff7efa370cd404cfff3a12c44223fdd14f Mon Sep 17 00:00:00 2001
From: Marcello <marcello@chainflip.io>
Date: Thu, 24 Oct 2024 12:55:12 +0200
Subject: [PATCH] bug: change_utxo not always present (#5340)

* fix bug: change_utxo not always present

Add logic to btc_utxos_balance rpc

* change comment

* fmt

* update logic in the rpc to match the callback

* address comments

* use log::info!

* address comments: check old key as well

* update rpc btc_utxos
---
 state-chain/runtime/src/chainflip.rs | 28 +++++++++++++++++++---------
 state-chain/runtime/src/lib.rs       | 28 ++++++++++++++++++++++++++--
 2 files changed, 45 insertions(+), 11 deletions(-)

diff --git a/state-chain/runtime/src/chainflip.rs b/state-chain/runtime/src/chainflip.rs
index 6226784623..28af980dfe 100644
--- a/state-chain/runtime/src/chainflip.rs
+++ b/state-chain/runtime/src/chainflip.rs
@@ -70,6 +70,7 @@ use cf_traits::{
 	ScheduledEgressDetails,
 };
 
+use cf_chains::{btc::ScriptPubkey, instances::BitcoinInstance};
 use codec::{Decode, Encode};
 use eth::Address as EvmAddress;
 use frame_support::{
@@ -815,15 +816,24 @@ impl OnBroadcastReady<Bitcoin> for BroadcastReadyProvider {
 		match api_call {
 			BitcoinApi::BatchTransfer(batch_transfer) => {
 				let tx_id = batch_transfer.bitcoin_transaction.txid();
-				let outputs = batch_transfer.bitcoin_transaction.outputs.clone();
-				let output_len = outputs.len();
-				let vout = output_len - 1;
-				let change_output = outputs.get(vout).unwrap();
-				Environment::add_bitcoin_change_utxo(
-					change_output.amount,
-					UtxoId { tx_id, vout: vout as u32 },
-					batch_transfer.change_utxo_key,
-				);
+				let outputs = &batch_transfer.bitcoin_transaction.outputs;
+				let btc_key = pallet_cf_threshold_signature::Pallet::<Runtime, BitcoinInstance>::keys(
+					pallet_cf_threshold_signature::Pallet::<Runtime, BitcoinInstance>::current_key_epoch()
+						.expect("We should always have an epoch set")).expect("We should always have a key set for the current epoch");
+				for (i, output) in outputs.iter().enumerate() {
+					if [
+						ScriptPubkey::Taproot(btc_key.previous.unwrap_or_default()),
+						ScriptPubkey::Taproot(btc_key.current),
+					]
+					.contains(&output.script_pubkey)
+					{
+						Environment::add_bitcoin_change_utxo(
+							output.amount,
+							UtxoId { tx_id, vout: i as u32 },
+							batch_transfer.change_utxo_key,
+						);
+					}
+				}
 			},
 			_ => unreachable!(),
 		}
diff --git a/state-chain/runtime/src/lib.rs b/state-chain/runtime/src/lib.rs
index 6c0cb87e9e..f982bbac9b 100644
--- a/state-chain/runtime/src/lib.rs
+++ b/state-chain/runtime/src/lib.rs
@@ -42,7 +42,7 @@ pub use cf_chains::instances::{
 use cf_chains::{
 	arb::api::ArbitrumApi,
 	assets::any::{AssetMap, ForeignChainAndAsset},
-	btc::{BitcoinCrypto, BitcoinRetryPolicy},
+	btc::{api::BitcoinApi, BitcoinCrypto, BitcoinRetryPolicy, ScriptPubkey},
 	dot::{self, PolkadotAccountId, PolkadotCrypto},
 	eth::{self, api::EthereumApi, Address as EthereumAddress, Ethereum},
 	evm::EvmCrypto,
@@ -2098,8 +2098,32 @@ impl_runtime_apis! {
 
 		fn cf_btc_utxos() -> BtcUtxos {
 			let utxos = pallet_cf_environment::BitcoinAvailableUtxos::<Runtime>::get();
+			let mut btc_balance = utxos.iter().fold(0, |acc, elem| acc + elem.amount);
+			//Sum the btc balance contained in the change utxos to the btc "free_balance"
+			let btc_ceremonies = pallet_cf_threshold_signature::PendingCeremonies::<Runtime,BitcoinInstance>::iter_values().map(|ceremony|{
+				ceremony.request_context.request_id
+			}).collect::<Vec<_>>();
+			let btc_key = pallet_cf_threshold_signature::Pallet::<Runtime, BitcoinInstance>::keys(
+				pallet_cf_threshold_signature::Pallet::<Runtime, BitcoinInstance>::current_key_epoch()
+				.expect("We should always have an epoch set")).expect("We should always have a key set for the current epoch");
+			for ceremony in btc_ceremonies {
+				if let RuntimeCall::BitcoinBroadcaster(pallet_cf_broadcast::pallet::Call::on_signature_ready{ api_call, ..}) = pallet_cf_threshold_signature::RequestCallback::<Runtime, BitcoinInstance>::get(ceremony).unwrap() {
+					if let BitcoinApi::BatchTransfer(batch_transfer) = *api_call {
+						for output in batch_transfer.bitcoin_transaction.outputs {
+							if [
+								ScriptPubkey::Taproot(btc_key.previous.unwrap_or_default()),
+								ScriptPubkey::Taproot(btc_key.current),
+							]
+							.contains(&output.script_pubkey)
+							{
+								btc_balance += output.amount;
+							}
+						}
+					}
+				}
+			}
 			BtcUtxos {
-				total_balance: utxos.iter().fold(0, |acc, elem| acc + elem.amount),
+				total_balance: btc_balance,
 				count: utxos.len() as u32,
 			}
 		}