Ensure build_commitment_tx and next_local/remote_tx_fee agree on the fee

... when possible. Sometimes they'll inevitably count different HTLCs, so we skip those cases.

Author: valentinewallace

lightning/src/ln/channel.rs

@@ -42,6 +42,8 @@ use std;
 use std::default::Default;
 use std::{cmp,mem,fmt};
 use std::ops::Deref;
+#[cfg(any(test, feature = "fuzztarget"))]
+use std::sync::Mutex;
 use bitcoin::hashes::hex::ToHex;
 
 #[cfg(test)]
@@ -386,6 +388,15 @@ pub(super) struct Channel<ChanSigner: ChannelKeys> {
 	commitment_secrets: CounterpartyCommitmentSecrets,
 
 	network_sync: UpdateStatus,
+
+	// We save these values so we can make sure `next_local_commit_tx_fee_msat` and
+	// `next_remote_commit_tx_fee_msat` properly predict what the next commitment transaction fee will
+	// be, by comparing the cached values to the fee of the tranaction generated by
+	// `build_commitment_transaction`.
+	#[cfg(any(test, feature = "fuzztarget"))]
+	next_local_commitment_tx_fee_cached: Mutex<Option<u64>>,
+	#[cfg(any(test, feature = "fuzztarget"))]
+	next_remote_commitment_tx_fee_cached: Mutex<Option<u64>>,
 }
 
 pub const OUR_MAX_HTLCS: u16 = 50; //TODO
@@ -557,6 +568,11 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 			commitment_secrets: CounterpartyCommitmentSecrets::new(),
 
 			network_sync: UpdateStatus::Fresh,
+
+			#[cfg(any(test, feature = "fuzztarget"))]
+			next_local_commitment_tx_fee_cached: Mutex::new(None),
+			#[cfg(any(test, feature = "fuzztarget"))]
+			next_remote_commitment_tx_fee_cached: Mutex::new(None),
 		})
 	}
 
@@ -790,6 +806,11 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 			commitment_secrets: CounterpartyCommitmentSecrets::new(),
 
 			network_sync: UpdateStatus::Fresh,
+
+			#[cfg(any(test, feature = "fuzztarget"))]
+			next_local_commitment_tx_fee_cached: Mutex::new(None),
+			#[cfg(any(test, feature = "fuzztarget"))]
+			next_remote_commitment_tx_fee_cached: Mutex::new(None),
 		};
 
 		Ok(chan)
@@ -867,7 +888,21 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 				InboundHTLCState::AwaitingRemoteRevokeToAnnounce(_) => (!generated_by_local, "AwaitingRemoteRevokeToAnnounce"),
 				InboundHTLCState::AwaitingAnnouncedRemoteRevoke(_) => (true, "AwaitingAnnouncedRemoteRevoke"),
 				InboundHTLCState::Committed => (true, "Committed"),
-				InboundHTLCState::LocalRemoved(_) => (!generated_by_local, "LocalRemoved"),
+				InboundHTLCState::LocalRemoved(_) => {
+					#[cfg(any(test, feature = "fuzztarget"))]
+					{
+						// Remote commitment transactions generated by us won't include our locally removed
+						// HTLCs, even though we include them when deciding whether to accept/send new HTLCs and
+						// setting these cached values. Therefore, set these values to None here so we don't
+						// assert that the transaction generated in this function has the same fee as the cached
+						// fee. Similar reasoning applies to LocalAnnounced and RemoveRemoved HTLCs below.
+						if !local && generated_by_local {
+							*self.next_local_commitment_tx_fee_cached.lock().unwrap() = None;
+							*self.next_remote_commitment_tx_fee_cached.lock().unwrap() = None;
+						}
+					}
+					(!generated_by_local, "LocalRemoved")
+				},
 			};
 
 			if include {
@@ -890,9 +925,27 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 
 		for ref htlc in self.pending_outbound_htlcs.iter() {
 			let (include, state_name) = match htlc.state {
-				OutboundHTLCState::LocalAnnounced(_) => (generated_by_local, "LocalAnnounced"),
+				OutboundHTLCState::LocalAnnounced(_) => {
+					#[cfg(any(test, feature = "fuzztarget"))]
+					{
+						if local && !generated_by_local {
+							*self.next_local_commitment_tx_fee_cached.lock().unwrap() = None;
+							*self.next_remote_commitment_tx_fee_cached.lock().unwrap() = None;
+						}
+					}
+					(generated_by_local, "LocalAnnounced")
+				},
 				OutboundHTLCState::Committed => (true, "Committed"),
-				OutboundHTLCState::RemoteRemoved(_) => (generated_by_local, "RemoteRemoved"),
+				OutboundHTLCState::RemoteRemoved(_) => {
+					#[cfg(any(test, feature = "fuzztarget"))]
+					{
+						if local && !generated_by_local {
+							*self.next_local_commitment_tx_fee_cached.lock().unwrap() = None;
+							*self.next_remote_commitment_tx_fee_cached.lock().unwrap() = None;
+						}
+					}
+					(generated_by_local, "RemoteRemoved")
+				},
 				OutboundHTLCState::AwaitingRemoteRevokeToRemove(_) => (generated_by_local, "AwaitingRemoteRevokeToRemove"),
 				OutboundHTLCState::AwaitingRemovedRemoteRevoke(_) => (false, "AwaitingRemovedRemoteRevoke"),
 			};
@@ -1740,7 +1793,18 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 			}
 		}
 
-		self.commit_tx_fee_msat(included_htlcs + addl_htlcs)
+		let num_htlcs = included_htlcs + addl_htlcs;
+		let res = self.commit_tx_fee_msat(num_htlcs);
+		#[cfg(any(test, feature = "fuzztarget"))]
+		{
+			if fee_spike_buffer_htlc.is_none() {
+				*self.next_local_commitment_tx_fee_cached.lock().unwrap() = Some(res);
+			} else {
+				let fee = self.commit_tx_fee_msat(num_htlcs - 1);
+				*self.next_local_commitment_tx_fee_cached.lock().unwrap() = Some(fee);
+			}
+		}
+		res
 	}
 
 	// Get the commitment tx fee for the remote's next commitment transaction based on the number of
@@ -1785,10 +1849,38 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 			}
 		}
 
-		self.commit_tx_fee_msat(included_htlcs + addl_htlcs)
+		let num_htlcs = included_htlcs + addl_htlcs;
+		let res = self.commit_tx_fee_msat(num_htlcs);
+		#[cfg(any(test, feature = "fuzztarget"))]
+		{
+			if fee_spike_buffer_htlc.is_none() {
+				*self.next_remote_commitment_tx_fee_cached.lock().unwrap() = Some(res);
+			} else {
+				let fee = self.commit_tx_fee_msat(num_htlcs - 1);
+				*self.next_remote_commitment_tx_fee_cached.lock().unwrap() = Some(fee);
+			}
+		}
+		res
 	}
 
-	pub fn update_add_htlc<F, L: Deref>(&mut self, msg: &msgs::UpdateAddHTLC, mut pending_forward_status: PendingHTLCStatus, create_pending_htlc_status: F, logger: &L) -> Result<(), ChannelError>
+
+	pub fn update_add_htlc<F, L: Deref>(&mut self, msg: &msgs::UpdateAddHTLC, pending_forward_status: PendingHTLCStatus, create_pending_htlc_status: F, logger: &L) -> Result<(), ChannelError>
+	where F: for<'a> Fn(&'a Self, PendingHTLCStatus, u16) -> PendingHTLCStatus, L::Target: Logger {
+		match self.update_add_htlc_internal(msg, pending_forward_status, create_pending_htlc_status, logger) {
+			Ok(()) => return Ok(()),
+			Err(e) => {
+				#[cfg(any(test, feature = "fuzztarget"))]
+				{
+					// If we errored, one of these fields still might've been set, so re-set them to None.
+					*self.next_local_commitment_tx_fee_cached.lock().unwrap() = None;
+					*self.next_remote_commitment_tx_fee_cached.lock().unwrap() = None;
+				}
+				return Err(e);
+			},
+		}
+	}
+
+	fn update_add_htlc_internal<F, L: Deref>(&mut self, msg: &msgs::UpdateAddHTLC, mut pending_forward_status: PendingHTLCStatus, create_pending_htlc_status: F, logger: &L) -> Result<(), ChannelError>
 	where F: for<'a> Fn(&'a Self, PendingHTLCStatus, u16) -> PendingHTLCStatus, L::Target: Logger {
 		// We can't accept HTLCs sent after we've sent a shutdown.
 		let local_sent_shutdown = (self.channel_state & (ChannelState::ChannelFunded as u32 | ChannelState::LocalShutdownSent as u32)) != (ChannelState::ChannelFunded as u32);
@@ -2019,15 +2111,24 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 			(commitment_tx.1, htlcs_cloned, commitment_tx.0, commitment_txid)
 		};
 
+		let total_fee = feerate_per_kw as u64 * (COMMITMENT_TX_BASE_WEIGHT + (num_htlcs as u64) * COMMITMENT_TX_WEIGHT_PER_HTLC) / 1000;
 		//If channel fee was updated by funder confirm funder can afford the new fee rate when applied to the current local commitment transaction
 		if update_fee {
-			let total_fee = feerate_per_kw as u64 * (COMMITMENT_TX_BASE_WEIGHT + (num_htlcs as u64) * COMMITMENT_TX_WEIGHT_PER_HTLC) / 1000;
-
 			let counterparty_reserve_we_require = Channel::<ChanSigner>::get_holder_selected_channel_reserve_satoshis(self.channel_value_satoshis);
 			if self.channel_value_satoshis - self.value_to_self_msat / 1000 < total_fee + counterparty_reserve_we_require {
 				return Err((None, ChannelError::Close("Funding remote cannot afford proposed new fee".to_owned())));
 			}
 		}
+		#[cfg(any(test, feature = "fuzztarget"))]
+		{
+			let projected_fee = match self.is_outbound() {
+				true => self.next_local_commitment_tx_fee_cached.lock().unwrap().take(),
+				false => self.next_remote_commitment_tx_fee_cached.lock().unwrap().take()
+			};
+			if let Some(fee) = projected_fee {
+				assert_eq!(total_fee, fee / 1000);
+			}
+		}
 
 		if msg.htlc_signatures.len() != num_htlcs {
 			return Err((None, ChannelError::Close(format!("Got wrong number of HTLC signatures ({}) from remote. It must be {}", msg.htlc_signatures.len(), num_htlcs))));
@@ -3726,6 +3827,28 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 	/// You MUST call send_commitment prior to any other calls on this Channel
 	/// If an Err is returned, it's a ChannelError::Ignore!
 	pub fn send_htlc(&mut self, amount_msat: u64, payment_hash: PaymentHash, cltv_expiry: u32, source: HTLCSource, onion_routing_packet: msgs::OnionPacket) -> Result<Option<msgs::UpdateAddHTLC>, ChannelError> {
+		match self.send_htlc_internal(amount_msat, payment_hash, cltv_expiry, source, onion_routing_packet) {
+			Ok(res) => return Ok(res),
+			Err(e) => {
+				#[cfg(any(test, feature = "fuzztarget"))]
+				{
+					// If we errored, one of these fields still might've been set, so re-set them to None.
+					*self.next_local_commitment_tx_fee_cached.lock().unwrap() = None;
+					*self.next_remote_commitment_tx_fee_cached.lock().unwrap() = None;
+				}
+				return Err(e);
+			}
+		}
+	}
+
+	/// Adds a pending outbound HTLC to this channel, note that you probably want
+	/// send_htlc_and_commit instead cause you'll want both messages at once.
+	/// This returns an option instead of a pure UpdateAddHTLC as we may be in a state where we are
+	/// waiting on the remote peer to send us a revoke_and_ack during which time we cannot add new
+	/// HTLCs on the wire or we wouldn't be able to determine what they actually ACK'ed.
+	/// You MUST call send_commitment prior to any other calls on this Channel
+	/// If an Err is returned, it's a ChannelError::Ignore!
+	fn send_htlc_internal(&mut self, amount_msat: u64, payment_hash: PaymentHash, cltv_expiry: u32, source: HTLCSource, onion_routing_packet: msgs::OnionPacket) -> Result<Option<msgs::UpdateAddHTLC>, ChannelError> {
 		if (self.channel_state & (ChannelState::ChannelFunded as u32 | BOTH_SIDES_SHUTDOWN_MASK)) != (ChannelState::ChannelFunded as u32) {
 			return Err(ChannelError::Ignore("Cannot send HTLC until channel is fully established and we haven't started shutting down".to_owned()));
 		}
@@ -3921,6 +4044,18 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 		let counterparty_commitment_txid = counterparty_commitment_tx.0.trust().txid();
 		let (signature, htlc_signatures);
 
+		#[cfg(any(test, feature = "fuzztarget"))]
+		{
+			let projected_fee = match self.is_outbound() {
+				true => self.next_local_commitment_tx_fee_cached.lock().unwrap().take(),
+				false => self.next_remote_commitment_tx_fee_cached.lock().unwrap().take()
+			};
+			if let Some(fee) = projected_fee {
+				let actual_fee = self.commit_tx_fee_msat(counterparty_commitment_tx.1);
+				assert_eq!(actual_fee, fee);
+			}
+		}
+
 		{
 			let mut htlcs = Vec::with_capacity(counterparty_commitment_tx.2.len());
 			for &(ref htlc, _) in counterparty_commitment_tx.2.iter() {
@@ -4511,6 +4646,11 @@ impl<'a, ChanSigner: ChannelKeys, K: Deref> ReadableArgs<&'a K> for Channel<Chan
 			commitment_secrets,
 
 			network_sync: UpdateStatus::Fresh,
+
+			#[cfg(any(test, feature = "fuzztarget"))]
+			next_local_commitment_tx_fee_cached: Mutex::new(None),
+			#[cfg(any(test, feature = "fuzztarget"))]
+			next_remote_commitment_tx_fee_cached: Mutex::new(None),
 		})
 	}
 }