Add ChannelMonitor::get_justice_txs for simplified watchtower integration

Adds a single-call API that returns signed justice transactions for all
revoked counterparty commitments, eliminating the need for Persist
implementors to maintain external state between commitment receipt and
revocation.

- New JusticeTransaction struct and get_justice_txs() method
- Stores recent counterparty CommitmentTransactions in ChannelMonitorImpl
- TLV field 39 (optional, backwards-compatible)
- Simplified test WatchtowerPersister (~80 lines of queue management removed)

Implements the approach suggested by TheBlueMatt in ldk-node#813 and
in review of #2552.

Author: FreeOnlineUser

lightning/src/chain/channelmonitor.rs

@@ -262,6 +262,19 @@ impl_writeable_tlv_based!(HTLCUpdate, {
 	(4, payment_preimage, option),
 });
 
+/// A signed justice transaction ready for broadcast or watchtower submission.
+///
+/// Returned by [`ChannelMonitor::get_justice_txs`].
+#[derive(Clone, Debug)]
+pub struct JusticeTransaction {
+	/// The fully signed justice transaction.
+	pub tx: Transaction,
+	/// The txid of the revoked counterparty commitment transaction.
+	pub revoked_commitment_txid: Txid,
+	/// The commitment number of the revoked commitment transaction.
+	pub commitment_number: u64,
+}
+
 /// If an output goes from claimable only by us to claimable by us or our counterparty within this
 /// many blocks, we consider it pinnable for the purposes of aggregating claims in a single
 /// transaction.
@@ -1372,6 +1385,13 @@ pub(crate) struct ChannelMonitorImpl<Signer: EcdsaChannelSigner> {
 	/// we now provide the transaction outright.
 	initial_counterparty_commitment_tx: Option<CommitmentTransaction>,
 
+	/// The latest counterparty commitment transaction(s), stored so that justice
+	/// transactions can be built and signed in a single call via [`ChannelMonitor::get_justice_txs`].
+	/// Contains the current and previous counterparty commitment(s). With splicing,
+	/// there may be multiple entries per commitment number (one per funding scope).
+	/// Pruned to remove entries more than one revocation old.
+	latest_counterparty_commitment_txs: Vec<CommitmentTransaction>,
+
 	/// The first block height at which we had no remaining claimable balances.
 	balances_empty_height: Option<u32>,
 
@@ -1755,6 +1775,7 @@ pub(crate) fn write_chanmon_internal<Signer: EcdsaChannelSigner, W: Writer>(
 		(34, channel_monitor.alternative_funding_confirmed, option),
 		(35, channel_monitor.is_manual_broadcast, required),
 		(37, channel_monitor.funding_seen_onchain, required),
+		(39, channel_monitor.latest_counterparty_commitment_txs, optional_vec),
 	});
 
 	Ok(())
@@ -1960,6 +1981,7 @@ impl<Signer: EcdsaChannelSigner> ChannelMonitor<Signer> {
 			counterparty_node_id: counterparty_node_id,
 			initial_counterparty_commitment_info: None,
 			initial_counterparty_commitment_tx: None,
+			latest_counterparty_commitment_txs: Vec::new(),
 			balances_empty_height: None,
 
 			failed_back_htlc_ids: new_hash_set(),
@@ -2271,6 +2293,25 @@ impl<Signer: EcdsaChannelSigner> ChannelMonitor<Signer> {
 		self.inner.lock().unwrap().sign_to_local_justice_tx(justice_tx, input_idx, value, commitment_number)
 	}
 
+	/// Returns signed justice transactions for all revoked counterparty commitment
+	/// transactions that this monitor knows about.
+	///
+	/// This is a convenience method that combines the functionality of
+	/// [`Self::counterparty_commitment_txs_from_update`], building justice transactions,
+	/// and [`Self::sign_to_local_justice_tx`] into a single call, eliminating the need
+	/// for callers to track intermediate state.
+	///
+	/// `feerate_per_kw` is used for the justice transaction fee calculation.
+	/// `destination_script` is the script where swept funds will be sent.
+	///
+	/// Returns a list of [`JusticeTransaction`]s, each containing a fully signed
+	/// transaction and metadata about the revoked commitment it punishes.
+	pub fn get_justice_txs(
+		&self, feerate_per_kw: u64, destination_script: ScriptBuf,
+	) -> Vec<JusticeTransaction> {
+		self.inner.lock().unwrap().get_justice_txs(feerate_per_kw, destination_script)
+	}
+
 	pub(crate) fn get_min_seen_secret(&self) -> u64 {
 		self.inner.lock().unwrap().get_min_seen_secret()
 	}
@@ -3483,6 +3524,7 @@ impl<Signer: EcdsaChannelSigner> ChannelMonitorImpl<Signer> {
 		self.provide_latest_counterparty_commitment_tx(commitment_tx.trust().txid(), Vec::new(), commitment_tx.commitment_number(),
 				commitment_tx.per_commitment_point());
 		// Soon, we will only populate this field
+		self.latest_counterparty_commitment_txs = vec![commitment_tx.clone()];
 		self.initial_counterparty_commitment_tx = Some(commitment_tx);
 	}
 
@@ -4284,8 +4326,21 @@ impl<Signer: EcdsaChannelSigner> ChannelMonitorImpl<Signer> {
 			}
 		}
 
-		#[cfg(debug_assertions)] {
-			self.counterparty_commitment_txs_from_update(updates);
+		let new_commitment_txs = self.counterparty_commitment_txs_from_update(updates);
+		if !new_commitment_txs.is_empty() {
+			self.latest_counterparty_commitment_txs.extend(new_commitment_txs);
+		}
+		// Prune commitment txs that are two or more revocations old. We keep one
+		// revocation depth so that get_justice_txs can sign the just-revoked
+		// commitment during this update_persisted_channel call.
+		if self.latest_counterparty_commitment_txs.len() > 1 {
+			let current = self.current_counterparty_commitment_number;
+			self.latest_counterparty_commitment_txs.retain(|tx| {
+				// Commitment numbers count down. Keep entries within 1 of current
+				// (current and the one just prior, which may have just been revoked).
+				// Also keep anything with a matching number (splicing).
+				tx.commitment_number() <= current + 1
+			});
 		}
 
 		self.latest_update_id = updates.update_id;
@@ -4563,6 +4618,52 @@ impl<Signer: EcdsaChannelSigner> ChannelMonitorImpl<Signer> {
 		Ok(justice_tx)
 	}
 
+	fn get_justice_txs(
+		&self, feerate_per_kw: u64, destination_script: ScriptBuf,
+	) -> Vec<JusticeTransaction> {
+		let mut result = Vec::new();
+
+		for commitment_tx in &self.latest_counterparty_commitment_txs {
+			let commitment_number = commitment_tx.commitment_number();
+
+			// Check if we have the revocation secret (i.e., this commitment is revoked)
+			let _secret = match self.get_secret(commitment_number) {
+				Some(s) => s,
+				None => continue,
+			};
+
+			let trusted = commitment_tx.trust();
+			let output_idx = match trusted.revokeable_output_index() {
+				Some(idx) => idx,
+				None => continue,
+			};
+
+			let built = trusted.built_transaction();
+			let value = built.transaction.output[output_idx].value;
+			let txid = built.txid;
+
+			let justice_tx = match trusted
+				.build_to_local_justice_tx(feerate_per_kw, destination_script.clone())
+			{
+				Ok(tx) => tx,
+				Err(_) => continue,
+			};
+
+			match self.sign_to_local_justice_tx(justice_tx, 0, value.to_sat(), commitment_number) {
+				Ok(signed_tx) => {
+					result.push(JusticeTransaction {
+						tx: signed_tx,
+						revoked_commitment_txid: txid,
+						commitment_number,
+					});
+				},
+				Err(_) => continue,
+			}
+		}
+
+		result
+	}
+
 	/// Can only fail if idx is < get_min_seen_secret
 	fn get_secret(&self, idx: u64) -> Option<[u8; 32]> {
 		self.commitment_secrets.get_secret(idx)
@@ -6521,6 +6622,7 @@ impl<'a, 'b, ES: EntropySource, SP: SignerProvider> ReadableArgs<(&'a ES, &'b SP
 		let mut alternative_funding_confirmed = None;
 		let mut is_manual_broadcast = RequiredWrapper(None);
 		let mut funding_seen_onchain = RequiredWrapper(None);
+		let mut latest_counterparty_commitment_txs: Option<Vec<CommitmentTransaction>> = Some(Vec::new());
 		read_tlv_fields!(reader, {
 			(1, funding_spend_confirmed, option),
 			(3, htlcs_resolved_on_chain, optional_vec),
@@ -6543,6 +6645,7 @@ impl<'a, 'b, ES: EntropySource, SP: SignerProvider> ReadableArgs<(&'a ES, &'b SP
 			(34, alternative_funding_confirmed, option),
 			(35, is_manual_broadcast, (default_value, false)),
 			(37, funding_seen_onchain, (default_value, true)),
+			(39, latest_counterparty_commitment_txs, optional_vec),
 		});
 		// Note that `payment_preimages_with_info` was added (and is always written) in LDK 0.1, so
 		// we can use it to determine if this monitor was last written by LDK 0.1 or later.
@@ -6714,6 +6817,7 @@ impl<'a, 'b, ES: EntropySource, SP: SignerProvider> ReadableArgs<(&'a ES, &'b SP
 			counterparty_node_id: counterparty_node_id.unwrap_or(dummy_node_id),
 			initial_counterparty_commitment_info,
 			initial_counterparty_commitment_tx,
+			latest_counterparty_commitment_txs: latest_counterparty_commitment_txs.unwrap_or_default(),
 			balances_empty_height,
 			failed_back_htlc_ids: new_hash_set(),
 

lightning/src/util/test_utils.rs

@@ -21,8 +21,6 @@ use crate::chain::channelmonitor::{
 };
 use crate::chain::transaction::OutPoint;
 use crate::chain::WatchedOutput;
-#[cfg(any(test, feature = "_externalize_tests"))]
-use crate::ln::chan_utils::CommitmentTransaction;
 use crate::ln::channel_state::ChannelDetails;
 use crate::ln::channelmanager;
 use crate::ln::inbound_payment::ExpandedKey;
@@ -679,21 +677,9 @@ impl<'a> chain::Watch<TestChannelSigner> for TestChainMonitor<'a> {
 	}
 }
 
-#[cfg(any(test, feature = "_externalize_tests"))]
-struct JusticeTxData {
-	justice_tx: Transaction,
-	value: Amount,
-	commitment_number: u64,
-}
-
 #[cfg(any(test, feature = "_externalize_tests"))]
 pub(crate) struct WatchtowerPersister {
 	persister: TestPersister,
-	/// Upon a new commitment_signed, we'll get a
-	/// ChannelMonitorUpdateStep::LatestCounterpartyCommitmentTxInfo. We'll store the justice tx
-	/// amount, and commitment number so we can build the justice tx after our counterparty
-	/// revokes it.
-	unsigned_justice_tx_data: Mutex<HashMap<ChannelId, VecDeque<JusticeTxData>>>,
 	/// After receiving a revoke_and_ack for a commitment number, we'll form and store the justice
 	/// tx which would be used to provide a watchtower with the data it needs.
 	watchtower_state: Mutex<HashMap<ChannelId, HashMap<Txid, Transaction>>>,
@@ -703,12 +689,9 @@ pub(crate) struct WatchtowerPersister {
 #[cfg(any(test, feature = "_externalize_tests"))]
 impl WatchtowerPersister {
 	pub(crate) fn new(destination_script: ScriptBuf) -> Self {
-		let unsigned_justice_tx_data = Mutex::new(new_hash_map());
-		let watchtower_state = Mutex::new(new_hash_map());
 		WatchtowerPersister {
 			persister: TestPersister::new(),
-			unsigned_justice_tx_data,
-			watchtower_state,
+			watchtower_state: Mutex::new(new_hash_map()),
 			destination_script,
 		}
 	}
@@ -724,23 +707,6 @@ impl WatchtowerPersister {
 			.get(commitment_txid)
 			.cloned()
 	}
-
-	fn form_justice_data_from_commitment(
-		&self, counterparty_commitment_tx: &CommitmentTransaction,
-	) -> Option<JusticeTxData> {
-		let trusted_tx = counterparty_commitment_tx.trust();
-		let output_idx = trusted_tx.revokeable_output_index()?;
-		let built_tx = trusted_tx.built_transaction();
-		let value = built_tx.transaction.output[output_idx as usize].value;
-		let justice_tx = trusted_tx
-			.build_to_local_justice_tx(
-				FEERATE_FLOOR_SATS_PER_KW as u64,
-				self.destination_script.clone(),
-			)
-			.ok()?;
-		let commitment_number = counterparty_commitment_tx.commitment_number();
-		Some(JusticeTxData { justice_tx, value, commitment_number })
-	}
 }
 
 #[cfg(any(test, feature = "_externalize_tests"))]
@@ -750,31 +716,25 @@ impl<Signer: sign::ecdsa::EcdsaChannelSigner> Persist<Signer> for WatchtowerPers
 	) -> chain::ChannelMonitorUpdateStatus {
 		let res = self.persister.persist_new_channel(monitor_name, data);
 
-		assert!(self
-			.unsigned_justice_tx_data
-			.lock()
-			.unwrap()
-			.insert(data.channel_id(), VecDeque::new())
-			.is_none());
 		assert!(self
 			.watchtower_state
 			.lock()
 			.unwrap()
 			.insert(data.channel_id(), new_hash_map())
 			.is_none());
 
-		let initial_counterparty_commitment_tx =
-			data.initial_counterparty_commitment_tx().expect("First and only call expects Some");
-		if let Some(justice_data) =
-			self.form_justice_data_from_commitment(&initial_counterparty_commitment_tx)
-		{
-			self.unsigned_justice_tx_data
+		// Use the simplified get_justice_txs API
+		let justice_txs =
+			data.get_justice_txs(FEERATE_FLOOR_SATS_PER_KW as u64, self.destination_script.clone());
+		for jtx in justice_txs {
+			self.watchtower_state
 				.lock()
 				.unwrap()
 				.get_mut(&data.channel_id())
 				.unwrap()
-				.push_back(justice_data);
+				.insert(jtx.revoked_commitment_txid, jtx.tx);
 		}
+
 		res
 	}
 
@@ -784,41 +744,18 @@ impl<Signer: sign::ecdsa::EcdsaChannelSigner> Persist<Signer> for WatchtowerPers
 	) -> chain::ChannelMonitorUpdateStatus {
 		let res = self.persister.update_persisted_channel(monitor_name, update, data);
 
-		if let Some(update) = update {
-			let commitment_txs = data.counterparty_commitment_txs_from_update(update);
-			let justice_datas = commitment_txs
-				.into_iter()
-				.filter_map(|commitment_tx| self.form_justice_data_from_commitment(&commitment_tx));
-			let mut channels_justice_txs = self.unsigned_justice_tx_data.lock().unwrap();
-			let channel_state = channels_justice_txs.get_mut(&data.channel_id()).unwrap();
-			channel_state.extend(justice_datas);
-
-			while let Some(JusticeTxData { justice_tx, value, commitment_number }) =
-				channel_state.front()
-			{
-				let input_idx = 0;
-				let commitment_txid = justice_tx.input[input_idx].previous_output.txid;
-				match data.sign_to_local_justice_tx(
-					justice_tx.clone(),
-					input_idx,
-					value.to_sat(),
-					*commitment_number,
-				) {
-					Ok(signed_justice_tx) => {
-						let dup = self
-							.watchtower_state
-							.lock()
-							.unwrap()
-							.get_mut(&data.channel_id())
-							.unwrap()
-							.insert(commitment_txid, signed_justice_tx);
-						assert!(dup.is_none());
-						channel_state.pop_front();
-					},
-					Err(_) => break,
-				}
-			}
+		// Use the simplified get_justice_txs API
+		let justice_txs =
+			data.get_justice_txs(FEERATE_FLOOR_SATS_PER_KW as u64, self.destination_script.clone());
+		for jtx in justice_txs {
+			self.watchtower_state
+				.lock()
+				.unwrap()
+				.get_mut(&data.channel_id())
+				.unwrap()
+				.insert(jtx.revoked_commitment_txid, jtx.tx);
 		}
+
 		res
 	}