Implement Confirm/Listen interfaces

src/sweep.rs

@@ -5,7 +5,7 @@ use crate::wallet::{num_blocks_from_conf_target, Wallet};
 use crate::{Error, KeysManager};
 
 use lightning::chain::chaininterface::{BroadcasterInterface, ConfirmationTarget, FeeEstimator};
-use lightning::chain::BestBlock;
+use lightning::chain::{self, BestBlock, Confirm, Filter, Listen};
 use lightning::impl_writeable_tlv_based;
 use lightning::sign::{EntropySource, SpendableOutputDescriptor};
 use lightning::util::ser::Writeable;
@@ -16,6 +16,8 @@ use bitcoin::{BlockHash, BlockHeader, LockTime, PackedLockTime, Script, Transact
 use std::ops::Deref;
 use std::sync::{Arc, Mutex};
 
+const CONSIDERED_SPENT_THRESHOLD_CONF: u32 = 6;
+
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub(crate) struct SpendableOutputInfo {
 	id: [u8; 32],
@@ -33,29 +35,33 @@ impl_writeable_tlv_based!(SpendableOutputInfo, {
 	(8, confirmed_in_block, option),
 });
 
-pub(crate) struct OutputSweeper<K: KVStore + Sync + Send, L: Deref>
+pub(crate) struct OutputSweeper<K: KVStore + Sync + Send, F: Deref, L: Deref>
 where
+	F::Target: Filter,
 	L::Target: Logger,
 {
 	outputs: Mutex<Vec<SpendableOutputInfo>>,
 	wallet: Arc<Wallet<bdk::database::SqliteDatabase, L>>,
 	keys_manager: Arc<KeysManager>,
 	kv_store: Arc<K>,
 	best_block: Mutex<BestBlock>,
+	chain_source: Option<F>,
 	logger: L,
 }
 
-impl<K: KVStore + Sync + Send, L: Deref> OutputSweeper<K, L>
+impl<K: KVStore + Sync + Send, F: Deref, L: Deref> OutputSweeper<K, F, L>
 where
+	F::Target: Filter,
 	L::Target: Logger,
 {
 	pub(crate) fn new(
 		outputs: Vec<SpendableOutputInfo>, wallet: Arc<Wallet<bdk::database::SqliteDatabase, L>>,
-		keys_manager: Arc<KeysManager>, kv_store: Arc<K>, best_block: BestBlock, logger: L,
+		keys_manager: Arc<KeysManager>, kv_store: Arc<K>, best_block: BestBlock,
+		chain_source: Option<F>, logger: L,
 	) -> Self {
 		let outputs = Mutex::new(outputs);
 		let best_block = Mutex::new(best_block);
-		Self { outputs, wallet, keys_manager, kv_store, best_block, logger }
+		Self { outputs, wallet, keys_manager, kv_store, best_block, chain_source, logger }
 	}
 
 	pub(crate) fn add_outputs(&self, output_descriptors: Vec<SpendableOutputDescriptor>) {
@@ -64,6 +70,9 @@ where
 		let spending_tx = match self.get_spending_tx(&output_descriptors) {
 			Ok(Some(spending_tx)) => {
 				self.wallet.broadcast_transactions(&[&spending_tx]);
+				if let Some(filter) = self.chain_source.as_ref() {
+					filter.register_tx(&spending_tx.txid(), &Script::new())
+				}
 				spending_tx
 			}
 			Ok(None) => {
@@ -138,3 +147,145 @@ where
 		})
 	}
 }
+
+impl<K: KVStore + Sync + Send, F: Deref, L: Deref> Listen for OutputSweeper<K, F, L>
+where
+	F::Target: Filter,
+	L::Target: Logger,
+{
+	fn filtered_block_connected(
+		&self, header: &BlockHeader, txdata: &chain::transaction::TransactionData, height: u32,
+	) {
+		{
+			let best_block = self.best_block.lock().unwrap();
+			assert_eq!(best_block.block_hash(), header.prev_blockhash,
+			"Blocks must be connected in chain-order - the connected header must build on the last connected header");
+			assert_eq!(best_block.height(), height - 1,
+			"Blocks must be connected in chain-order - the connected block height must be one greater than the previous height");
+		}
+
+		self.transactions_confirmed(header, txdata, height);
+		self.best_block_updated(header, height);
+	}
+
+	fn block_disconnected(&self, header: &BlockHeader, height: u32) {
+		let new_height = height - 1;
+		{
+			let mut best_block = self.best_block.lock().unwrap();
+			assert_eq!(best_block.block_hash(), header.block_hash(),
+				"Blocks must be disconnected in chain-order - the disconnected header must be the last connected header");
+			assert_eq!(best_block.height(), height,
+				"Blocks must be disconnected in chain-order - the disconnected block must have the correct height");
+			*best_block = BestBlock::new(header.prev_blockhash, new_height)
+		}
+
+		let mut locked_outputs = self.outputs.lock().unwrap();
+		for output_info in locked_outputs.iter_mut() {
+			if output_info.confirmed_in_block == Some((height, header.block_hash())) {
+				output_info.confirmed_in_block = None;
+			}
+		}
+	}
+}
+
+impl<K: KVStore + Sync + Send, F: Deref, L: Deref> Confirm for OutputSweeper<K, F, L>
+where
+	F::Target: Filter,
+	L::Target: Logger,
+{
+	fn transactions_confirmed(
+		&self, header: &BlockHeader, txdata: &chain::transaction::TransactionData, height: u32,
+	) {
+		let mut locked_outputs = self.outputs.lock().unwrap();
+		for (_, tx) in txdata {
+			locked_outputs
+				.iter_mut()
+				.filter(|o| o.spending_tx.txid() == tx.txid())
+				.for_each(|o| o.confirmed_in_block = Some((height, header.block_hash())));
+		}
+	}
+
+	fn transaction_unconfirmed(&self, txid: &Txid) {
+		let mut locked_outputs = self.outputs.lock().unwrap();
+
+		// Get what height was unconfirmed.
+		let unconf_height = locked_outputs
+			.iter()
+			.find(|o| o.spending_tx.txid() == *txid)
+			.and_then(|o| o.confirmed_in_block)
+			.map(|t| t.0);
+
+		// Unconfirm all >= this height.
+		locked_outputs
+			.iter_mut()
+			.filter(|o| o.confirmed_in_block.map(|t| t.0) >= unconf_height)
+			.for_each(|o| o.confirmed_in_block = None);
+	}
+
+	fn best_block_updated(&self, header: &BlockHeader, height: u32) {
+		*self.best_block.lock().unwrap() = BestBlock::new(header.block_hash(), height);
+
+		let mut locked_outputs = self.outputs.lock().unwrap();
+
+		// Rebroadcast all outputs that didn't get confirmed by now.
+		for output_info in locked_outputs.iter_mut().filter(|o| o.confirmed_in_block.is_none()) {
+			if height
+				>= output_info.broadcast_height
+					+ num_blocks_from_conf_target(ConfirmationTarget::Background)
+			{
+				let output_descriptors = vec![output_info.descriptor.clone()];
+				match self.get_spending_tx(&output_descriptors) {
+					Ok(Some(spending_tx)) => {
+						self.wallet.broadcast_transactions(&[&spending_tx]);
+						if let Some(filter) = self.chain_source.as_ref() {
+							filter.register_tx(&spending_tx.txid(), &Script::new())
+						}
+						output_info.spending_tx = spending_tx;
+						output_info.broadcast_height = height;
+					}
+					Ok(None) => {
+						log_debug!(
+							self.logger,
+							"Omitted spending static outputs: {:?}",
+							output_descriptors
+						);
+					}
+					Err(err) => {
+						log_error!(self.logger, "Error spending outputs: {:?}", err);
+					}
+				};
+			}
+		}
+
+		// Prune all outputs that have sufficient depth by now.
+		locked_outputs.retain(|o| {
+			if let Some((conf_height, _)) = o.confirmed_in_block {
+				if height >= conf_height + CONSIDERED_SPENT_THRESHOLD_CONF {
+					let key = hex_utils::to_string(&o.id);
+					match self.kv_store.remove(SPENDABLE_OUTPUT_INFO_PERSISTENCE_NAMESPACE, &key) {
+						Ok(_) => return false,
+						Err(e) => {
+							log_error!(
+								self.logger,
+								"Removal of key {}/{} failed due to: {}",
+								SPENDABLE_OUTPUT_INFO_PERSISTENCE_NAMESPACE,
+								key,
+								e
+							);
+							return true;
+						}
+					}
+				}
+			}
+			true
+		});
+	}
+
+	fn get_relevant_txids(&self) -> Vec<(Txid, Option<BlockHash>)> {
+		let locked_outputs = self.outputs.lock().unwrap();
+		locked_outputs
+			.iter()
+			.map(|o| (o.spending_tx.txid(), o.confirmed_in_block.map(|c| c.1)))
+			.collect::<Vec<_>>()
+	}
+}