fixup! simplify CommitmentTransaction constructor

devrandom · devrandom · commit 49a2fc6bee29 · 2020-12-19T17:29:15.000-08:00
diff --git a/lightning/src/ln/chan_utils.rs b/lightning/src/ln/chan_utils.rs
@@ -736,8 +736,8 @@ impl HolderCommitmentTransaction {
 			counterparty_parameters: Some(CounterpartyChannelTransactionParameters { pubkeys: channel_pubkeys.clone(), selected_contest_delay: 0 }),
 			funding_outpoint: Some(chain::transaction::OutPoint { txid: Default::default(), index: 0 })
 		};
-		let aux: Vec<()> = Vec::new();
-		let inner = CommitmentTransaction::new_with_auxiliary_htlc_data(0, 0, 0, keys, 0, Vec::new(), aux, &channel_parameters.as_counterparty_broadcastable()).0;
+		let mut htlcs_with_aux: Vec<(_, ())> = Vec::new();
+		let inner = CommitmentTransaction::new_with_auxiliary_htlc_data(0, 0, 0, keys, 0, &mut htlcs_with_aux, &channel_parameters.as_counterparty_broadcastable());
 		HolderCommitmentTransaction {
 			inner,
 			counterparty_sig: dummy_sig,
@@ -887,32 +887,20 @@ impl_writeable!(CommitmentTransaction, 0, {
 impl CommitmentTransaction {
 	/// Construct an object of the class while assigning transaction output indices to HTLCs.
 	///
-	/// Also keeps track of auxiliary HTLC data and returns it along with the mutated and sorted HTLCs.
-	/// This allows the caller to match the HTLC output index with the auxiliary data.
+	/// Populates HTLCOutputInCommitment.transaction_output_index in htlcs_with_aux.
+	///
+	/// The generic T allows the caller to match the HTLC output index with auxiliary data.
 	/// This auxiliary data is not stored in this object.
 	///
 	/// Only include HTLCs that are above the dust limit for the channel.
-	///
-	/// Panics if the length of htlcs and aux are different.
-	pub fn new_with_auxiliary_htlc_data<T>(commitment_number: u64, to_broadcaster_value_sat: u64, to_countersignatory_value_sat: u64, keys: TxCreationKeys, feerate_per_kw: u32, htlcs: Vec<HTLCOutputInCommitment>, aux: Vec<T>, channel_parameters: &DirectedChannelTransactionParameters) -> (CommitmentTransaction, Vec<(HTLCOutputInCommitment, T)>) {
+	pub fn new_with_auxiliary_htlc_data<T>(commitment_number: u64, to_broadcaster_value_sat: u64, to_countersignatory_value_sat: u64, keys: TxCreationKeys, feerate_per_kw: u32, htlcs_with_aux: &mut Vec<(HTLCOutputInCommitment, T)>, channel_parameters: &DirectedChannelTransactionParameters) -> CommitmentTransaction {
 		// Sort outputs and populate output indices while keeping track of the auxiliary data
-		let mut txouts = Self::internal_build_outputs(&keys, to_broadcaster_value_sat, to_countersignatory_value_sat, &htlcs, aux, channel_parameters).unwrap();
-		let mut result_htlcs_with_aux = Vec::new();
-		let mut htlcs = Vec::new();
-		let mut outputs = Vec::new();
-		for (idx, mut out) in txouts.drain(..).enumerate() {
-			if let Some(mut htlc) = (out.1).1.take() {
-				htlc.0.transaction_output_index = Some(idx as u32);
-				result_htlcs_with_aux.push((htlc.0.clone(), htlc.1));
-				htlcs.push(htlc.0);
-			}
-			outputs.push(out.0);
-		}
+		let (outputs, htlcs, _) = Self::internal_build_outputs(&keys, to_broadcaster_value_sat, to_countersignatory_value_sat, htlcs_with_aux, channel_parameters).unwrap();
 
 		let (obscured_commitment_transaction_number, txins) = Self::internal_build_inputs(commitment_number, channel_parameters);
 		let transaction = Self::make_transaction(obscured_commitment_transaction_number, txins, outputs);
 		let txid = transaction.txid();
-		let info = CommitmentTransaction {
+		CommitmentTransaction {
 			commitment_number,
 			to_broadcaster_value_sat,
 			to_countersignatory_value_sat,
@@ -923,44 +911,32 @@ impl CommitmentTransaction {
 				transaction,
 				txid
 			},
-		};
-		(info, result_htlcs_with_aux)
+		}
 	}
 
 	#[cfg(feature = "_test_utils")]
-	/// Re-build the Bitcoin transaction.
-	///
-	/// This function ignores the pre-built transaction and the pre-derived TxCreationKeys
-	/// for use in an environment, such as an external signer, where we don't trust these to be
-	/// derived by the Lightning node.
+	/// Get the transaction output scripts.
 	///
-	/// The second return value are the scripts that were built for debugging purposes, since
-	/// only the hash is present in the actual transaction.
-	pub fn rebuild_transaction(&self, keys: &TxCreationKeys, channel_parameters: &DirectedChannelTransactionParameters) -> Result<(BuiltCommitmentTransaction, Vec<Script>), ()> {
-		self.internal_rebuild_transaction(keys, channel_parameters)
+	/// For debugging purposes, since only the hash is present in the actual transaction.
+	pub fn get_output_scripts(&self, keys: &TxCreationKeys, channel_parameters: &DirectedChannelTransactionParameters) -> Result<Vec<Script>, ()> {
+		let mut htlcs_with_aux = self.htlcs.iter().map(|h| (h.clone(), ())).collect();
+		let (_, _, scripts) = Self::internal_build_outputs(keys, self.to_broadcaster_value_sat, self.to_countersignatory_value_sat, &mut htlcs_with_aux, channel_parameters)?;
+		Ok(scripts)
 	}
 
-	fn internal_rebuild_transaction(&self, keys: &TxCreationKeys, channel_parameters: &DirectedChannelTransactionParameters) -> Result<(BuiltCommitmentTransaction, Vec<Script>), ()> {
+	fn internal_rebuild_transaction(&self, keys: &TxCreationKeys, channel_parameters: &DirectedChannelTransactionParameters) -> Result<BuiltCommitmentTransaction, ()> {
 		let (obscured_commitment_transaction_number, txins) = Self::internal_build_inputs(self.commitment_number, channel_parameters);
 
-		let mut aux = Vec::with_capacity(self.htlcs.len());
-		aux.resize(self.htlcs.len(), ());
-		let mut txouts = Self::internal_build_outputs(keys, self.to_broadcaster_value_sat, self.to_countersignatory_value_sat, &self.htlcs, aux, channel_parameters)?;
-
-		let mut outputs = Vec::with_capacity(txouts.len());
-		let mut scripts = Vec::with_capacity(txouts.len());
-		for (txout, (script, _)) in txouts.drain(..) {
-			outputs.push(txout);
-			scripts.push(script);
-		}
+		let mut htlcs_with_aux = self.htlcs.iter().map(|h| (h.clone(), ())).collect();
+		let (outputs, _, _) = Self::internal_build_outputs(keys, self.to_broadcaster_value_sat, self.to_countersignatory_value_sat, &mut htlcs_with_aux, channel_parameters)?;
 
 		let transaction = Self::make_transaction(obscured_commitment_transaction_number, txins, outputs);
 		let txid = transaction.txid();
 		let built_transaction = BuiltCommitmentTransaction {
 			transaction,
 			txid
 		};
-		Ok((built_transaction, scripts))
+		Ok(built_transaction)
 	}
 
 	fn make_transaction(obscured_commitment_transaction_number: u64, txins: Vec<TxIn>, outputs: Vec<TxOut>) -> Transaction {
@@ -976,12 +952,11 @@ impl CommitmentTransaction {
 	// - initial sorting of outputs / HTLCs in the constructor, in which case T is auxiliary data the
 	//   caller needs to have sorted together with the HTLCs so it can keep track of the output index
 	// - building of a bitcoin transaction (build -> build_outputs), in which case T is just ()
-	fn internal_build_outputs<T>(keys: &TxCreationKeys, to_broadcaster_value_sat: u64, to_countersignatory_value_sat: u64, htlcs: &Vec<HTLCOutputInCommitment>, aux: Vec<T>, channel_parameters: &DirectedChannelTransactionParameters) -> Result<Vec<(TxOut, (Script, Option<(HTLCOutputInCommitment, T)>))>, ()> {
-		assert_eq!(htlcs.len(), aux.len(), "htlcs vs aux length mismatch");
+	fn internal_build_outputs<T>(keys: &TxCreationKeys, to_broadcaster_value_sat: u64, to_countersignatory_value_sat: u64, htlcs_with_aux: &mut Vec<(HTLCOutputInCommitment, T)>, channel_parameters: &DirectedChannelTransactionParameters) -> Result<(Vec<TxOut>, Vec<HTLCOutputInCommitment>, Vec<Script>), ()> {
 		let countersignatory_pubkeys = channel_parameters.countersignatory_pubkeys();
 		let contest_delay = channel_parameters.contest_delay();
 
-		let mut txouts: Vec<(TxOut, (Script, Option<(HTLCOutputInCommitment, T)>))> = Vec::new();
+		let mut txouts: Vec<(TxOut, (Script, Option<&mut HTLCOutputInCommitment>))> = Vec::new();
 
 		if to_countersignatory_value_sat > 0 {
 			let script = script_for_p2wpkh(&countersignatory_pubkeys.payment_point);
@@ -1009,29 +984,43 @@ impl CommitmentTransaction {
 			));
 		}
 
-		for (htlc, t) in htlcs.iter().zip(aux) {
+		let mut htlcs = Vec::with_capacity(htlcs_with_aux.len());
+		for (htlc, _) in htlcs_with_aux {
 			let script = chan_utils::get_htlc_redeemscript(&htlc, &keys);
 			let txout = TxOut {
 				script_pubkey: script.to_v0_p2wsh(),
 				value: htlc.amount_msat / 1000,
 			};
-			txouts.push((txout, (script, Some((htlc.clone(), t)))));
+			txouts.push((txout, (script, Some(htlc))));
 		}
 
 		// Sort output in BIP-69 order (amount, scriptPubkey).  Tie-breaks based on HTLC
 		// CLTV expiration height.
 		sort_outputs(&mut txouts, |a, b| {
 			if let &(_, Some(ref a_htlcout)) = a {
 				if let &(_, Some(ref b_htlcout)) = b {
-					a_htlcout.0.cltv_expiry.cmp(&b_htlcout.0.cltv_expiry)
+					a_htlcout.cltv_expiry.cmp(&b_htlcout.cltv_expiry)
 				} else {
 					cmp::Ordering::Equal
 				}
 			} else {
 				cmp::Ordering::Equal
 			}
 		});
-		Ok(txouts)
+
+		let mut outputs = Vec::with_capacity(txouts.len());
+		let mut scripts = Vec::new();
+		scripts.clear(); // Prevent mut warning
+		for (idx, out) in txouts.drain(..).enumerate() {
+			if let Some(htlc) = (out.1).1 {
+				htlc.transaction_output_index = Some(idx as u32);
+				htlcs.push(htlc.clone());
+			}
+			outputs.push(out.0);
+			#[cfg(feature = "_test_utils")]
+			scripts.push((out.1).0);
+		}
+		Ok((outputs, htlcs, scripts))
 	}
 
 	fn internal_build_inputs(commitment_number: u64, channel_parameters: &DirectedChannelTransactionParameters) -> (u64, Vec<TxIn>) {
@@ -1116,7 +1105,7 @@ impl CommitmentTransaction {
 		if keys != self.keys {
 			return Err(());
 		}
-		let (tx, _) = self.internal_rebuild_transaction(&keys, channel_parameters)?;
+		let tx = self.internal_rebuild_transaction(&keys, channel_parameters)?;
 		if self.built.transaction != tx.transaction || self.built.txid != tx.txid {
 			return Err(());
 		}
diff --git a/lightning/src/ln/channel.rs b/lightning/src/ln/channel.rs
@@ -839,8 +839,7 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 	fn build_commitment_transaction<L: Deref>(&self, commitment_number: u64, keys: &TxCreationKeys, local: bool, generated_by_local: bool, feerate_per_kw: u32, logger: &L) -> (CommitmentTransaction, usize, Vec<(HTLCOutputInCommitment, Option<&HTLCSource>)>) where L::Target: Logger {
 		let mut included_dust_htlcs: Vec<(HTLCOutputInCommitment, Option<&HTLCSource>)> = Vec::new();
 		let num_htlcs = self.pending_inbound_htlcs.len() + self.pending_outbound_htlcs.len();
-		let mut included_non_dust_htlcs: Vec<HTLCOutputInCommitment> = Vec::with_capacity(num_htlcs);
-		let mut included_non_dust_sources: Vec<Option<&HTLCSource>> = Vec::with_capacity(num_htlcs);
+		let mut included_non_dust_htlcs: Vec<(HTLCOutputInCommitment, Option<&HTLCSource>)> = Vec::with_capacity(num_htlcs);
 
 		let broadcaster_dust_limit_satoshis = if local { self.holder_dust_limit_satoshis } else { self.counterparty_dust_limit_satoshis };
 		let mut remote_htlc_total_msat = 0;
@@ -867,8 +866,7 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 					let htlc_in_tx = get_htlc_in_commitment!($htlc, true);
 					if $htlc.amount_msat / 1000 >= broadcaster_dust_limit_satoshis + (feerate_per_kw as u64 * HTLC_TIMEOUT_TX_WEIGHT / 1000) {
 						log_trace!(logger, "   ...including {} {} HTLC {} (hash {}) with value {}", if $outbound { "outbound" } else { "inbound" }, $state_name, $htlc.htlc_id, log_bytes!($htlc.payment_hash.0), $htlc.amount_msat);
-						included_non_dust_htlcs.push(htlc_in_tx);
-						included_non_dust_sources.push($source);
+						included_non_dust_htlcs.push((htlc_in_tx, $source));
 					} else {
 						log_trace!(logger, "   ...including {} {} dust HTLC {} (hash {}) with value {} due to dust limit", if $outbound { "outbound" } else { "inbound" }, $state_name, $htlc.htlc_id, log_bytes!($htlc.payment_hash.0), $htlc.amount_msat);
 						included_dust_htlcs.push((htlc_in_tx, $source));
@@ -877,8 +875,7 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 					let htlc_in_tx = get_htlc_in_commitment!($htlc, false);
 					if $htlc.amount_msat / 1000 >= broadcaster_dust_limit_satoshis + (feerate_per_kw as u64 * HTLC_SUCCESS_TX_WEIGHT / 1000) {
 						log_trace!(logger, "   ...including {} {} HTLC {} (hash {}) with value {}", if $outbound { "outbound" } else { "inbound" }, $state_name, $htlc.htlc_id, log_bytes!($htlc.payment_hash.0), $htlc.amount_msat);
-						included_non_dust_htlcs.push(htlc_in_tx);
-						included_non_dust_sources.push($source);
+						included_non_dust_htlcs.push((htlc_in_tx, $source));
 					} else {
 						log_trace!(logger, "   ...including {} {} dust HTLC {} (hash {}) with value {}", if $outbound { "outbound" } else { "inbound" }, $state_name, $htlc.htlc_id, log_bytes!($htlc.payment_hash.0), $htlc.amount_msat);
 						included_dust_htlcs.push((htlc_in_tx, $source));
@@ -993,19 +990,20 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 		let channel_parameters =
 			if local { self.channel_transaction_parameters.as_holder_broadcastable() }
 			else { self.channel_transaction_parameters.as_counterparty_broadcastable() };
-		let (info, mut htlcs_included) = CommitmentTransaction::new_with_auxiliary_htlc_data(commitment_number,
-		                                                                                     value_to_a as u64,
-		                                                                                     value_to_b as u64,
-		                                                                                     keys.clone(),
-		                                                                                     feerate_per_kw,
-		                                                                                     included_non_dust_htlcs,
-		                                                                                     included_non_dust_sources,
-		                                                                                     &channel_parameters
+		let tx = CommitmentTransaction::new_with_auxiliary_htlc_data(commitment_number,
+		                                                             value_to_a as u64,
+		                                                             value_to_b as u64,
+		                                                             keys.clone(),
+		                                                             feerate_per_kw,
+		                                                             &mut included_non_dust_htlcs,
+		                                                             &channel_parameters
 		);
-
+		let mut htlcs_included = included_non_dust_htlcs;
+		// The unwrap is safe, because all non-dust HTLCs have been assigned an output index
+		htlcs_included.sort_unstable_by_key(|h| h.0.transaction_output_index.unwrap());
 		htlcs_included.append(&mut included_dust_htlcs);
 
-		(info, num_nondust_htlcs, htlcs_included)
+		(tx, num_nondust_htlcs, htlcs_included)
 	}
 
 	#[inline]
diff --git a/lightning/src/ln/functional_tests.rs b/lightning/src/ln/functional_tests.rs
@@ -1663,11 +1663,10 @@ fn test_fee_spike_violation_fails_htlc() {
 			local_chan_balance,
 			commit_tx_keys.clone(),
 			feerate_per_kw,
-			vec![accepted_htlc_info],
-			vec![()],
+			&mut vec![(accepted_htlc_info, ())],
 			&local_chan.channel_transaction_parameters.as_counterparty_broadcastable()
 		);
-		local_chan_keys.sign_counterparty_commitment(&commitment_tx.0, &secp_ctx).unwrap()
+		local_chan_keys.sign_counterparty_commitment(&commitment_tx, &secp_ctx).unwrap()
 	};
 
 	let commit_signed_msg = msgs::CommitmentSigned {
