Add ChannelContext::sort_{inbounds, outbounds}_by_state

This function sorts included from excluded htlcs based on their state,
and runs the two passed closures on the resulting sets.

To avoid mutating the outbound preimages vector in both the included and
the excluded closures, this commit stops including outbound preimages
for HTLCs which are included in the commitment transaction, and only
includes outbound preimages for HTLCs which are not included in the
commitment transaction. This seems ok from the available docs on
`sign_counterparty_commitment`, but test coverage is missing.

Author: tankyleo

lightning/src/ln/channel.rs

@@ -974,10 +974,10 @@ struct HTLCStats {
 }
 
 /// A struct gathering data on a commitment, either local or remote.
-struct CommitmentData<'a> {
+struct CommitmentData {
 	tx: CommitmentTransaction,
 	stats: CommitmentStats,
-	htlcs_included: Vec<(HTLCOutputInCommitment, Option<&'a HTLCSource>)>, // the list of HTLCs (dust HTLCs *included*) which were not ignored when building the transaction
+	htlcs_included: Vec<(HTLCOutputInCommitment, Option<Box<HTLCSource>>)>, // the list of HTLCs (dust HTLCs *included*) which were not ignored when building the transaction
 	outbound_htlc_preimages: Vec<PaymentPreimage>, // preimages for successful offered HTLCs since last commitment
 	inbound_htlc_preimages: Vec<PaymentPreimage>, // preimages for successful received HTLCs since last commitment
 }
@@ -3604,7 +3604,6 @@ impl<SP: Deref> ChannelContext<SP> where SP::Target: SignerProvider {
 			}
 			bitcoin_tx.txid
 		};
-		let mut htlcs_cloned: Vec<_> = commitment_data.htlcs_included.iter().map(|htlc| (htlc.0.clone(), htlc.1.map(|h| h.clone()))).collect();
 
 		// If our counterparty updated the channel fee in this commitment transaction, check that
 		// they can actually afford the new fee now.
@@ -3655,10 +3654,10 @@ impl<SP: Deref> ChannelContext<SP> where SP::Target: SignerProvider {
 			separate_nondust_htlc_sources = rand_val % 2 == 0;
 		}
 
-		let mut nondust_htlc_sources = Vec::with_capacity(htlcs_cloned.len());
-		let mut htlcs_and_sigs = Vec::with_capacity(htlcs_cloned.len());
+		let mut nondust_htlc_sources = Vec::with_capacity(commitment_data.htlcs_included.len());
+		let mut htlcs_and_sigs = Vec::with_capacity(commitment_data.htlcs_included.len());
 		let holder_keys = commitment_data.tx.trust().keys();
-		for (idx, (htlc, mut source_opt)) in htlcs_cloned.drain(..).enumerate() {
+		for (idx, (htlc, mut source_opt)) in commitment_data.htlcs_included.into_iter().enumerate() {
 			if let Some(_) = htlc.transaction_output_index {
 				let htlc_tx = chan_utils::build_htlc_transaction(&commitment_txid, commitment_data.tx.feerate_per_kw(),
 					funding.get_counterparty_selected_contest_delay().unwrap(), &htlc, &self.channel_type,
@@ -3674,14 +3673,14 @@ impl<SP: Deref> ChannelContext<SP> where SP::Target: SignerProvider {
 					return Err(ChannelError::close("Invalid HTLC tx signature from peer".to_owned()));
 				}
 				if !separate_nondust_htlc_sources {
-					htlcs_and_sigs.push((htlc, Some(msg.htlc_signatures[idx]), source_opt.take()));
+					htlcs_and_sigs.push((htlc, Some(msg.htlc_signatures[idx]), source_opt.take().map(|source| *source)));
 				}
 			} else {
-				htlcs_and_sigs.push((htlc, None, source_opt.take()));
+				htlcs_and_sigs.push((htlc, None, source_opt.take().map(|source| *source)));
 			}
 			if separate_nondust_htlc_sources {
 				if let Some(source) = source_opt.take() {
-					nondust_htlc_sources.push(source);
+					nondust_htlc_sources.push(*source);
 				}
 			}
 			debug_assert!(source_opt.is_none(), "HTLCSource should have been put somewhere");
@@ -3705,6 +3704,30 @@ impl<SP: Deref> ChannelContext<SP> where SP::Target: SignerProvider {
 		})
 	}
 
+	fn sort_inbounds_by_state<I, E>(&self, generated_by_local: bool, mut for_each_included: I, mut for_each_excluded: E)
+		where I: FnMut(&InboundHTLCOutput), E: FnMut(&InboundHTLCOutput)
+	{
+		self.pending_inbound_htlcs.iter().for_each(|htlc| {
+			if htlc.state.included_in_commitment(generated_by_local) {
+				for_each_included(htlc);
+			} else {
+				for_each_excluded(htlc);
+			}
+		})
+	}
+
+	fn sort_outbounds_by_state<I, E>(&self, generated_by_local: bool, mut for_each_included: I, mut for_each_excluded: E)
+		where I: FnMut(&OutboundHTLCOutput), E: FnMut(&OutboundHTLCOutput)
+	{
+		self.pending_outbound_htlcs.iter().for_each(|htlc| {
+			if htlc.state.included_in_commitment(generated_by_local) {
+				for_each_included(htlc);
+			} else {
+				for_each_excluded(htlc);
+			}
+		})
+	}
+
 	/// Generates stats on a potential commitment transaction build, without actually building the
 	/// commitment transaction. See `build_commitment_transaction` for further docs.
 	#[inline]
@@ -3744,38 +3767,33 @@ impl<SP: Deref> ChannelContext<SP> where SP::Target: SignerProvider {
 				}
 			}
 		}
-
-		for ref htlc in self.pending_inbound_htlcs.iter() {
-			if htlc.state.included_in_commitment(generated_by_local) {
-				count_nondust_htlc!(htlc, false);
-				remote_htlc_total_msat += htlc.amount_msat;
-			} else {
-				log_trace!(logger, "   ...not counting inbound HTLC {} (hash {}) with value {} due to state ({})", htlc.htlc_id, htlc.payment_hash, htlc.amount_msat, htlc.state.as_str());
-				match htlc.state {
-					InboundHTLCState::LocalRemoved(InboundHTLCRemovalReason::Fulfill(_preimage)) => {
-						value_to_self_msat_offset += htlc.amount_msat as i64;
-					},
-					_ => {},
-				}
+		
+		let for_each_included = |htlc: &InboundHTLCOutput| {
+			count_nondust_htlc!(htlc, false);
+			remote_htlc_total_msat += htlc.amount_msat;
+		};
+		let for_each_excluded = |htlc: &InboundHTLCOutput| {
+			log_trace!(logger, "   ...not counting inbound HTLC {} (hash {}) with value {} due to state ({})", htlc.htlc_id, htlc.payment_hash, htlc.amount_msat, htlc.state.as_str());
+			if let InboundHTLCState::LocalRemoved(InboundHTLCRemovalReason::Fulfill(_preimage)) = htlc.state {
+				value_to_self_msat_offset += htlc.amount_msat as i64;
 			}
-		}
+		};
+		self.sort_inbounds_by_state(generated_by_local, for_each_included, for_each_excluded);
 
-		for ref htlc in self.pending_outbound_htlcs.iter() {
-			if htlc.state.included_in_commitment(generated_by_local) {
-				count_nondust_htlc!(htlc, true);
-				local_htlc_total_msat += htlc.amount_msat;
-			} else {
-				log_trace!(logger, "   ...not counting outbound HTLC {} (hash {}) with value {} due to state ({})", htlc.htlc_id, htlc.payment_hash, htlc.amount_msat, htlc.state.as_str());
-				match htlc.state {
-					OutboundHTLCState::AwaitingRemoteRevokeToRemove(OutboundHTLCOutcome::Success(_)) |
-					OutboundHTLCState::AwaitingRemovedRemoteRevoke(OutboundHTLCOutcome::Success(_)) |
-					OutboundHTLCState::RemoteRemoved(OutboundHTLCOutcome::Success(_)) => {
-						value_to_self_msat_offset -= htlc.amount_msat as i64;
-					},
-					_ => {},
-				}
+		let for_each_included = |htlc: &OutboundHTLCOutput| {
+			count_nondust_htlc!(htlc, true);
+			local_htlc_total_msat += htlc.amount_msat;
+		};
+		let for_each_excluded = |htlc: &OutboundHTLCOutput| {
+			log_trace!(logger, "   ...not counting outbound HTLC {} (hash {}) with value {} due to state ({})", htlc.htlc_id, htlc.payment_hash, htlc.amount_msat, htlc.state.as_str());
+			if let OutboundHTLCState::AwaitingRemoteRevokeToRemove(OutboundHTLCOutcome::Success(_)) |
+				   OutboundHTLCState::AwaitingRemovedRemoteRevoke(OutboundHTLCOutcome::Success(_)) |
+				   OutboundHTLCState::RemoteRemoved(OutboundHTLCOutcome::Success(_))
+			= htlc.state {
+				value_to_self_msat_offset -= htlc.amount_msat as i64;
 			}
-		}
+		};
+		self.sort_outbounds_by_state(generated_by_local, for_each_included, for_each_excluded);
 
 		// # Panics
 		//
@@ -3844,9 +3862,9 @@ impl<SP: Deref> ChannelContext<SP> where SP::Target: SignerProvider {
 			remote_balance_before_fee_anchors_msat
 		} = stats;
 
-		let mut included_dust_htlcs: Vec<(HTLCOutputInCommitment, Option<&HTLCSource>)> = Vec::new();
+		let mut included_dust_htlcs: Vec<(HTLCOutputInCommitment, Option<Box<HTLCSource>>)> = Vec::new();
 		let num_htlcs = self.pending_inbound_htlcs.len() + self.pending_outbound_htlcs.len();
-		let mut included_non_dust_htlcs: Vec<(HTLCOutputInCommitment, Option<&HTLCSource>)> = Vec::with_capacity(num_htlcs);
+		let mut included_non_dust_htlcs: Vec<(HTLCOutputInCommitment, Option<Box<HTLCSource>>)> = Vec::with_capacity(num_htlcs);
 
 		log_trace!(logger, "Building commitment transaction number {} (really {} xor {}) for channel {} for {}, generated by {} with fee {}...",
 			commitment_number, (INITIAL_COMMITMENT_NUMBER - commitment_number),
@@ -3880,30 +3898,32 @@ impl<SP: Deref> ChannelContext<SP> where SP::Target: SignerProvider {
 		}
 
 		let mut inbound_htlc_preimages: Vec<PaymentPreimage> = Vec::new();
+		let mut outbound_htlc_preimages: Vec<PaymentPreimage> = Vec::new();
 
-		for ref htlc in self.pending_inbound_htlcs.iter() {
-			if htlc.state.included_in_commitment(generated_by_local) {
-				add_htlc_output!(htlc, false, None);
-			} else {
-				log_trace!(logger, "   ...not including inbound HTLC {} (hash {}) with value {} due to state ({})", htlc.htlc_id, htlc.payment_hash, htlc.amount_msat, htlc.state.as_str());
-				if let Some(preimage) = htlc.state.preimage() {
-					inbound_htlc_preimages.push(preimage);
-				}
+		let for_each_included = |htlc: &InboundHTLCOutput| {
+			add_htlc_output!(htlc, false, None);
+		};
+		let for_each_excluded = |htlc: &InboundHTLCOutput| {
+			log_trace!(logger, "   ...not including inbound HTLC {} (hash {}) with value {} due to state ({})", htlc.htlc_id, htlc.payment_hash, htlc.amount_msat, htlc.state.as_str());
+			if let Some(preimage) = htlc.state.preimage() {
+				inbound_htlc_preimages.push(preimage);
 			}
-		}
-
-		let mut outbound_htlc_preimages: Vec<PaymentPreimage> = Vec::new();
+		};
+		self.sort_inbounds_by_state(generated_by_local, for_each_included, for_each_excluded);
 
-		for ref htlc in self.pending_outbound_htlcs.iter() {
+		let for_each_included = |htlc: &OutboundHTLCOutput| {
+			// We box the source here because borrowed data cannot escape a closure
+			add_htlc_output!(htlc, true, Some(Box::new(htlc.source.clone())));
+		};
+		let for_each_excluded = |htlc: &OutboundHTLCOutput| {
+			log_trace!(logger, "   ...not including outbound HTLC {} (hash {}) with value {} due to state ({})", htlc.htlc_id, htlc.payment_hash, htlc.amount_msat, htlc.state.as_str());
+			// We previously included outbound preimages for both included and excluded HTLCs,
+			// but now only include preimages for HTLCs excluded from the commitment transaction.
 			if let Some(preimage) = htlc.state.preimage() {
 				outbound_htlc_preimages.push(preimage);
 			}
-			if htlc.state.included_in_commitment(generated_by_local) {
-				add_htlc_output!(htlc, true, Some(&htlc.source));
-			} else {
-				log_trace!(logger, "   ...not including outbound HTLC {} (hash {}) with value {} due to state ({})", htlc.htlc_id, htlc.payment_hash, htlc.amount_msat, htlc.state.as_str());
-			}
-		}
+		};
+		self.sort_outbounds_by_state(generated_by_local, for_each_included, for_each_excluded);
 
 		// We MUST use saturating subs here, as the funder's balance is not guaranteed to be greater
 		// than or equal to the sum of `total_fee_sat` and `total_anchors_sat`.
@@ -8814,11 +8834,9 @@ impl<SP: Deref> FundedChannel<SP> where
 		}
 		self.context.resend_order = RAACommitmentOrder::RevokeAndACKFirst;
 
-		let (mut htlcs_ref, counterparty_commitment_tx) =
+		let (htlcs, counterparty_commitment_tx) =
 			self.build_commitment_no_state_update(logger);
 		let counterparty_commitment_txid = counterparty_commitment_tx.trust().txid();
-		let htlcs: Vec<(HTLCOutputInCommitment, Option<Box<HTLCSource>>)> =
-			htlcs_ref.drain(..).map(|(htlc, htlc_source)| (htlc, htlc_source.map(|source_ref| Box::new(source_ref.clone())))).collect();
 
 		if self.context.announcement_sigs_state == AnnouncementSigsState::MessageSent {
 			self.context.announcement_sigs_state = AnnouncementSigsState::Committed;
@@ -8845,7 +8863,7 @@ impl<SP: Deref> FundedChannel<SP> where
 	}
 
 	fn build_commitment_no_state_update<L: Deref>(&self, logger: &L)
-	-> (Vec<(HTLCOutputInCommitment, Option<&HTLCSource>)>, CommitmentTransaction)
+	-> (Vec<(HTLCOutputInCommitment, Option<Box<HTLCSource>>)>, CommitmentTransaction)
 	where L::Target: Logger
 	{
 		let commitment_data = self.context.build_commitment_transaction(&self.funding,