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(test)]
+use std::sync::Mutex;
 use bitcoin::hashes::hex::ToHex;
 
 #[cfg(test)]
@@ -391,6 +393,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(test)]
+	next_local_commitment_tx_fee_cached: Mutex<Option<u64>>,
+	#[cfg(test)]
+	next_remote_commitment_tx_fee_cached: Mutex<Option<u64>>,
 }
 
 pub const OUR_MAX_HTLCS: u16 = 50; //TODO
@@ -559,6 +570,11 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 			commitment_secrets: CounterpartyCommitmentSecrets::new(),
 
 			network_sync: UpdateStatus::Fresh,
+
+			#[cfg(test)]
+			next_local_commitment_tx_fee_cached: Mutex::new(None),
+			#[cfg(test)]
+			next_remote_commitment_tx_fee_cached: Mutex::new(None),
 		})
 	}
 
@@ -787,6 +803,11 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 			commitment_secrets: CounterpartyCommitmentSecrets::new(),
 
 			network_sync: UpdateStatus::Fresh,
+
+			#[cfg(test)]
+			next_local_commitment_tx_fee_cached: Mutex::new(None),
+			#[cfg(test)]
+			next_remote_commitment_tx_fee_cached: Mutex::new(None),
 		};
 
 		Ok(chan)
@@ -905,7 +926,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(test)]
+					{
+						// 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 {
@@ -928,9 +963,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(test)]
+					{
+						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(test)]
+						{
+							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"),
 			};
@@ -1780,7 +1833,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(test)]
+		{
+			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
@@ -1825,10 +1889,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(test)]
+		{
+			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, 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(test)]
+				{
+					// 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);
+			},
+		}
 	}
 
-	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>
+	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);
@@ -2054,16 +2146,25 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 			return Err((None, ChannelError::Close("Invalid commitment tx signature from peer".to_owned())));
 		}
 
+		let num_htlcs = commitment_tx.1;
+		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 num_htlcs = commitment_tx.1;
-			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(test)]
+		{
+			let projected_fee = match self.channel_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() != commitment_tx.1 {
 			return Err((None, ChannelError::Close(format!("Got wrong number of HTLC signatures ({}) from remote. It must be {}", msg.htlc_signatures.len(), commitment_tx.1))));
@@ -3744,6 +3845,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(test)]
+				{
+					// 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()));
 		}
@@ -3938,6 +4061,18 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 		let counterparty_commitment_tx = self.build_commitment_transaction(self.cur_counterparty_commitment_transaction_number, &counterparty_keys, false, true, feerate_per_kw, logger);
 		let (signature, htlc_signatures);
 
+		#[cfg(test)]
+		{
+			let projected_fee = match self.channel_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() {
@@ -4524,6 +4659,11 @@ impl<ChanSigner: ChannelKeys + Readable> Readable for Channel<ChanSigner> {
 			commitment_secrets,
 
 			network_sync: UpdateStatus::Fresh,
+
+			#[cfg(test)]
+			next_local_commitment_tx_fee_cached: Mutex::new(None),
+			#[cfg(test)]
+			next_remote_commitment_tx_fee_cached: Mutex::new(None),
 		})
 	}
 }