Add would_broadcast_at_height functionality to Channel

In service to the larger refactor of removing the Channel's reference
to its ChannelMonitor.

Author: valentinewallace

lightning/src/ln/channel.rs

@@ -33,7 +33,7 @@ use ln::chan_utils;
 use chain::chaininterface::{FeeEstimator,ConfirmationTarget};
 use chain::transaction::OutPoint;
 use chain::keysinterface::{ChannelKeys, KeysInterface};
-use util::transaction_utils;
+use util::{byte_utils, transaction_utils};
 use util::ser::{Readable, Writeable, Writer};
 use util::logger::Logger;
 use util::errors::APIError;
@@ -43,6 +43,7 @@ use std;
 use std::default::Default;
 use std::{cmp,mem,fmt};
 use std::ops::Deref;
+use std::collections::HashMap;
 use bitcoin::hashes::hex::ToHex;
 
 #[cfg(test)]
@@ -104,16 +105,18 @@ enum InboundHTLCState {
 	AwaitingAnnouncedRemoteRevoke(PendingHTLCStatus),
 	Committed,
 	/// Removed by us and a new commitment_signed was sent (if we were AwaitingRemoteRevoke when we
-	/// created it we would have put it in the holding cell instead). When they next revoke_and_ack
-	/// we'll drop it.
+	/// created it we would have put it in the holding cell instead).
 	/// Note that we have to keep an eye on the HTLC until we've received a broadcastable
 	/// commitment transaction without it as otherwise we'll have to force-close the channel to
 	/// claim it before the timeout (obviously doesn't apply to revoked HTLCs that we can't claim
-	/// anyway). That said, ChannelMonitor does this for us (see
-	/// ChannelMonitor::would_broadcast_at_height) so we actually remove the HTLC from our own
-	/// local state before then, once we're sure that the next commitment_signed and
-	/// ChannelMonitor::provide_latest_local_commitment_tx_info will not include this HTLC.
+	/// anyway).
 	LocalRemoved(InboundHTLCRemovalReason),
+	/// Removed by us and by the remote.
+	/// Note that we have to keep an eye on the HTLC until we've received a broadcastable
+	/// commitment transaction without it as otherwise we'll have to force-close the channel to
+	/// claim it before the timeout (obviously doesn't apply to revoked HTLCs that we can't claim
+	/// anyway).
+	RemoteRemoved,
 }
 
 struct InboundHTLCOutput {
@@ -293,6 +296,7 @@ pub(super) struct Channel<ChanSigner: ChannelKeys> {
 	pending_inbound_htlcs: Vec<InboundHTLCOutput>,
 	pending_outbound_htlcs: Vec<OutboundHTLCOutput>,
 	holding_cell_htlc_updates: Vec<HTLCUpdateAwaitingACK>,
+	payment_preimages: HashMap<PaymentHash, PaymentPreimage>,
 
 	/// When resending CS/RAA messages on channel monitor restoration or on reconnect, we always
 	/// need to ensure we resend them in the order we originally generated them. Note that because
@@ -510,6 +514,7 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 			pending_inbound_htlcs: Vec::new(),
 			pending_outbound_htlcs: Vec::new(),
 			holding_cell_htlc_updates: Vec::new(),
+			payment_preimages: HashMap::new(),
 			pending_update_fee: None,
 			holding_cell_update_fee: None,
 			next_local_htlc_id: 0,
@@ -738,6 +743,7 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 			pending_inbound_htlcs: Vec::new(),
 			pending_outbound_htlcs: Vec::new(),
 			holding_cell_htlc_updates: Vec::new(),
+			payment_preimages: HashMap::new(),
 			pending_update_fee: None,
 			holding_cell_update_fee: None,
 			next_local_htlc_id: 0,
@@ -795,6 +801,57 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 		Ok(chan)
 	}
 
+	pub(super) fn monitor_would_broadcast_at_height<L: Deref>(&self, height: u32, logger: &L) -> bool where L::Target: Logger {
+		macro_rules! get_htlc_output {
+			($htlc: expr, $offered: expr) => {
+				HTLCOutputInCommitment{
+					offered: $offered,
+					amount_msat: $htlc.amount_msat,
+					cltv_expiry: $htlc.cltv_expiry,
+					payment_hash: $htlc.payment_hash,
+					transaction_output_index: None
+				}
+			}
+		}
+		let inbound_local_outputs = self.pending_inbound_htlcs.iter()
+			.map(|htlc| get_htlc_output!(htlc, false));
+
+		let outbound_local_outputs = self.pending_outbound_htlcs.iter()
+			.filter_map(|htlc| {
+				match htlc.state {
+					OutboundHTLCState::AwaitingRemovedRemoteRevoke(_) => None,
+					_ => Some(get_htlc_output!(htlc, true))
+				}
+			});
+		let awaiting_raa = (self.channel_state & ChannelState::AwaitingRemoteRevoke as u32) != 0;
+		let inbound_remote_outputs = self.pending_inbound_htlcs.iter()
+			.filter_map(|htlc| {
+				match htlc.state {
+					InboundHTLCState::LocalRemoved(_) => {
+						if awaiting_raa {
+							Some(get_htlc_output!(htlc, true))
+						} else { None }
+					},
+					InboundHTLCState::AwaitingRemoteRevokeToAnnounce(_) => None,
+					InboundHTLCState::RemoteRemoved => None,
+					_ => Some(get_htlc_output!(htlc, true))
+				}
+			});
+		let outbound_remote_outputs = self.pending_outbound_htlcs.iter()
+			.filter_map(|htlc| {
+				match htlc.state {
+					OutboundHTLCState::AwaitingRemovedRemoteRevoke(_) => {
+						if awaiting_raa {
+							Some(get_htlc_output!(htlc, false))
+						} else { None }
+					},
+					_ => Some(get_htlc_output!(htlc, false))
+				}
+			});
+
+		ChannelMonitor::<ChanSigner>::would_broadcast_at_height_given_htlcs(inbound_local_outputs.chain(outbound_local_outputs), inbound_remote_outputs.chain(outbound_remote_outputs), height, &self.payment_preimages, logger)
+	}
+
 	// Utilities to build transactions:
 
 	fn get_commitment_transaction_number_obscure_factor(&self) -> u64 {
@@ -909,6 +966,7 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 				InboundHTLCState::AwaitingAnnouncedRemoteRevoke(_) => (true, "AwaitingAnnouncedRemoteRevoke"),
 				InboundHTLCState::Committed => (true, "Committed"),
 				InboundHTLCState::LocalRemoved(_) => (!generated_by_local, "LocalRemoved"),
+				InboundHTLCState::RemoteRemoved => (false, "RemoteRemoved"),
 			};
 
 			if include {
@@ -1216,6 +1274,7 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 		// We have to put the payment_preimage in the channel_monitor right away here to ensure we
 		// can claim it even if the channel hits the chain before we see their next commitment.
 		self.latest_monitor_update_id += 1;
+		self.payment_preimages.insert(payment_hash_calc, payment_preimage_arg.clone());
 		let monitor_update = ChannelMonitorUpdate {
 			update_id: self.latest_monitor_update_id,
 			updates: vec![ChannelMonitorUpdateStep::PaymentPreimage {
@@ -2011,6 +2070,23 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 			return Err((None, ChannelError::Close(format!("Got wrong number of HTLC signatures ({}) from remote. It must be {}", msg.htlc_signatures.len(), local_commitment_tx.1))));
 		}
 
+		{
+
+			let pending_inbound_htlcs: &mut Vec<_> = &mut self.pending_inbound_htlcs;
+			let payment_preimages: &mut HashMap<_, _> = &mut self.payment_preimages;
+			// A RemoteRemoved HTLC need to be monitored for expiration until we receive a
+			// broadcastable commitment tx without said HTLC. Now that we've confirmed that
+			// this commitment signed message provides said commitment tx, we can drop the
+			// RemoteRemoved HTLCs we were previously watching for.
+			pending_inbound_htlcs.retain(|htlc| {
+				log_trace!(logger, "Removing inbound RemovedRemoved {}", log_bytes!(htlc.payment_hash.0));
+				if let &InboundHTLCState::RemoteRemoved = &htlc.state {
+					payment_preimages.remove(&htlc.payment_hash);
+					false
+				} else { true }
+			});
+		}
+
 		// TODO: Merge these two, sadly they are currently both required to be passed separately to
 		// ChannelMonitor:
 		let mut htlcs_without_source = Vec::with_capacity(local_commitment_tx.2.len());
@@ -2303,30 +2379,20 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 			// Take references explicitly so that we can hold multiple references to self.
 			let pending_inbound_htlcs: &mut Vec<_> = &mut self.pending_inbound_htlcs;
 			let pending_outbound_htlcs: &mut Vec<_> = &mut self.pending_outbound_htlcs;
-
-			// We really shouldnt have two passes here, but retain gives a non-mutable ref (Rust bug)
-			pending_inbound_htlcs.retain(|htlc| {
+			let payment_preimages: &mut HashMap<_, _> = &mut self.payment_preimages;
+			for htlc in pending_inbound_htlcs.iter_mut() {
+				let mut is_local_removed = false;
 				if let &InboundHTLCState::LocalRemoved(ref reason) = &htlc.state {
-					log_trace!(logger, " ...removing inbound LocalRemoved {}", log_bytes!(htlc.payment_hash.0));
 					if let &InboundHTLCRemovalReason::Fulfill(_) = reason {
 						value_to_self_msat_diff += htlc.amount_msat as i64;
 					}
-					false
-				} else { true }
-			});
-			pending_outbound_htlcs.retain(|htlc| {
-				if let &OutboundHTLCState::AwaitingRemovedRemoteRevoke(ref fail_reason) = &htlc.state {
-					log_trace!(logger, " ...removing outbound AwaitingRemovedRemoteRevoke {}", log_bytes!(htlc.payment_hash.0));
-					if let Some(reason) = fail_reason.clone() { // We really want take() here, but, again, non-mut ref :(
-						revoked_htlcs.push((htlc.source.clone(), htlc.payment_hash, reason));
-					} else {
-						// They fulfilled, so we sent them money
-						value_to_self_msat_diff -= htlc.amount_msat as i64;
-					}
-					false
-				} else { true }
-			});
-			for htlc in pending_inbound_htlcs.iter_mut() {
+					is_local_removed = true;
+				}
+				if is_local_removed {
+					let mut state = InboundHTLCState::RemoteRemoved;
+					mem::swap(&mut state, &mut htlc.state);
+					continue
+				}
 				let swap = if let &InboundHTLCState::AwaitingRemoteRevokeToAnnounce(_) = &htlc.state {
 					log_trace!(logger, " ...promoting inbound AwaitingRemoteRevokeToAnnounce {} to Committed", log_bytes!(htlc.payment_hash.0));
 					true
@@ -2364,6 +2430,20 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 					}
 				}
 			}
+			// We really shouldnt have two passes here, but retain gives a non-mutable ref (Rust bug)
+			pending_outbound_htlcs.retain(|htlc| {
+				if let &OutboundHTLCState::AwaitingRemovedRemoteRevoke(ref fail_reason) = &htlc.state {
+					log_trace!(logger, " ...removing outbound AwaitingRemovedRemoteRevoke {}", log_bytes!(htlc.payment_hash.0));
+					if let Some(reason) = fail_reason.clone() { // We really want take() here, but, again, non-mut ref :(
+						revoked_htlcs.push((htlc.source.clone(), htlc.payment_hash, reason));
+					} else {
+						// They fulfilled, so we sent them money
+						value_to_self_msat_diff -= htlc.amount_msat as i64;
+					}
+					payment_preimages.remove(&htlc.payment_hash);
+					false
+				} else { true }
+			});
 			for htlc in pending_outbound_htlcs.iter_mut() {
 				if let OutboundHTLCState::LocalAnnounced(_) = htlc.state {
 					log_trace!(logger, " ...promoting outbound LocalAnnounced {} to Committed", log_bytes!(htlc.payment_hash.0));
@@ -2540,6 +2620,7 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 					// the commitment_signed we can re-transmit the update then.
 					true
 				},
+				InboundHTLCState::RemoteRemoved => true,
 			}
 		});
 		self.next_remote_htlc_id -= inbound_drop_count;
@@ -3115,14 +3196,6 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 		self.user_id
 	}
 
-	/// May only be called after funding has been initiated (ie is_funding_initiated() is true)
-	pub fn channel_monitor(&mut self) -> &mut ChannelMonitor<ChanSigner> {
-		if self.channel_state < ChannelState::FundingSent as u32 {
-			panic!("Can't get a channel monitor until funding has been created");
-		}
-		self.channel_monitor.as_mut().unwrap()
-	}
-
 	/// Guaranteed to be Some after both FundingLocked messages have been exchanged (and, thus,
 	/// is_usable() returns true).
 	/// Allowed in any state (including after shutdown)
@@ -4118,6 +4191,9 @@ impl<ChanSigner: ChannelKeys + Writeable> Writeable for Channel<ChanSigner> {
 					4u8.write(writer)?;
 					removal_reason.write(writer)?;
 				},
+				&InboundHTLCState::RemoteRemoved => {
+					5u8.write(writer)?;
+				}
 			}
 		}
 
@@ -4151,6 +4227,11 @@ impl<ChanSigner: ChannelKeys + Writeable> Writeable for Channel<ChanSigner> {
 			}
 		}
 
+		writer.write_all(&byte_utils::be64_to_array(self.payment_preimages.len() as u64))?;
+		for payment_preimage in self.payment_preimages.values() {
+			writer.write_all(&payment_preimage.0[..])?;
+		}
+
 		(self.holding_cell_htlc_updates.len() as u64).write(writer)?;
 		for update in self.holding_cell_htlc_updates.iter() {
 			match update {
@@ -4248,6 +4329,8 @@ impl<ChanSigner: ChannelKeys + Writeable> Writeable for Channel<ChanSigner> {
 	}
 }
 
+const MAX_ALLOC_SIZE: usize = 64*1024;
+
 impl<ChanSigner: ChannelKeys + Readable> Readable for Channel<ChanSigner> {
 	fn read<R : ::std::io::Read>(reader: &mut R) -> Result<Self, DecodeError> {
 		let _ver: u8 = Readable::read(reader)?;
@@ -4287,6 +4370,7 @@ impl<ChanSigner: ChannelKeys + Readable> Readable for Channel<ChanSigner> {
 					2 => InboundHTLCState::AwaitingAnnouncedRemoteRevoke(Readable::read(reader)?),
 					3 => InboundHTLCState::Committed,
 					4 => InboundHTLCState::LocalRemoved(Readable::read(reader)?),
+					5 => InboundHTLCState::RemoteRemoved,
 					_ => return Err(DecodeError::InvalidValue),
 				},
 			});
@@ -4312,6 +4396,16 @@ impl<ChanSigner: ChannelKeys + Readable> Readable for Channel<ChanSigner> {
 			});
 		}
 
+		let payment_preimages_len: u64 = Readable::read(reader)?;
+		let mut payment_preimages = HashMap::with_capacity(cmp::min(payment_preimages_len as usize, MAX_ALLOC_SIZE / 32));
+		for _ in 0..payment_preimages_len {
+			let preimage: PaymentPreimage = Readable::read(reader)?;
+			let hash = PaymentHash(Sha256::hash(&preimage.0[..]).into_inner());
+			if let Some(_) = payment_preimages.insert(hash, preimage) {
+				return Err(DecodeError::InvalidValue);
+			}
+		}
+
 		let holding_cell_htlc_update_count: u64 = Readable::read(reader)?;
 		let mut holding_cell_htlc_updates = Vec::with_capacity(cmp::min(holding_cell_htlc_update_count as usize, OUR_MAX_HTLCS as usize*2));
 		for _ in 0..holding_cell_htlc_update_count {
@@ -4428,6 +4522,7 @@ impl<ChanSigner: ChannelKeys + Readable> Readable for Channel<ChanSigner> {
 			pending_inbound_htlcs,
 			pending_outbound_htlcs,
 			holding_cell_htlc_updates,
+			payment_preimages,
 
 			resend_order,
 

lightning/src/ln/channelmanager.rs

@@ -3104,7 +3104,7 @@ impl<ChanSigner: ChannelKeys, M: Deref + Sync + Send, T: Deref + Sync + Send, K:
 						}
 					}
 				}
-				if channel.is_funding_initiated() && channel.channel_monitor().would_broadcast_at_height(height, &self.logger) {
+				if channel.is_funding_initiated() && channel.monitor_would_broadcast_at_height(height, &self.logger) {
 					if let Some(short_id) = channel.get_short_channel_id() {
 						short_to_id.remove(&short_id);
 					}

lightning/src/ln/channelmonitor.rs

@@ -49,6 +49,7 @@ use util::ser::{Readable, MaybeReadable, Writer, Writeable, U48};
 use util::{byte_utils, events};
 
 use std::collections::{HashMap, hash_map};
+use std::borrow::Borrow;
 use std::sync::Mutex;
 use std::{hash,cmp, mem};
 use std::ops::Deref;
@@ -1993,7 +1994,30 @@ impl<ChanSigner: ChannelKeys> ChannelMonitor<ChanSigner> {
 		self.last_block_hash = block_hash.clone();
 	}
 
-	pub(super) fn would_broadcast_at_height<L: Deref>(&self, height: u32, logger: &L) -> bool where L::Target: Logger {
+	fn would_broadcast_at_height<L: Deref>(&self, height: u32, logger: &L) -> bool where L::Target: Logger {
+		let local_outputs = self.current_local_commitment_tx.htlc_outputs.iter().map(|&(ref a, _, _)| a);
+
+		let map_fn: fn(&(HTLCOutputInCommitment, Option<Box<HTLCSource>>)) -> &HTLCOutputInCommitment = |a| &a.0;
+		let dummy = Vec::new();
+		let curr_remote_outputs = if let Some(ref txid) = self.current_remote_commitment_txid {
+			if let Some(ref htlc_outputs) = self.remote_claimable_outpoints.get(txid) {
+				htlc_outputs.iter().map(map_fn)
+			} else { dummy.iter().map(map_fn) }
+		} else { dummy.iter().map(map_fn) };
+
+		let map_fn: fn(&(HTLCOutputInCommitment, Option<Box<HTLCSource>>)) -> &HTLCOutputInCommitment = |a| &a.0;
+		let prev_remote_outputs = if let Some(ref txid) = self.prev_remote_commitment_txid {
+			if let Some(ref htlc_outputs) = self.remote_claimable_outpoints.get(txid) {
+				htlc_outputs.iter().map(map_fn)
+			} else { dummy.iter().map(map_fn) }
+		} else { dummy.iter().map(map_fn) };
+
+		ChannelMonitor::<ChanSigner>::would_broadcast_at_height_given_htlcs(local_outputs, curr_remote_outputs.chain(prev_remote_outputs), height, &self.payment_preimages, logger)
+	}
+
+	pub(super) fn would_broadcast_at_height_given_htlcs<L: Deref, IL: Iterator<Item = impl Borrow<HTLCOutputInCommitment>>, IR: Iterator<Item = impl Borrow<HTLCOutputInCommitment>>>
+		(local_htlc_outputs: IL, remote_htlc_outputs: IR, height: u32, preimages: &HashMap<PaymentHash, PaymentPreimage>, logger: &L) -> bool
+		where L::Target: Logger {
 		// We need to consider all HTLCs which are:
 		//  * in any unrevoked remote commitment transaction, as they could broadcast said
 		//    transactions and we'd end up in a race, or
@@ -2030,28 +2054,18 @@ impl<ChanSigner: ChannelKeys> ChannelMonitor<ChanSigner> {
 					//      LATENCY_GRACE_PERIOD_BLOCKS + 2*CLTV_CLAIM_BUFFER <= CLTV_EXPIRY_DELTA
 					//  The final, above, condition is checked for statically in channelmanager
 					//  with CHECK_CLTV_EXPIRY_SANITY_2.
-					let htlc_outbound = $local_tx == htlc.offered;
-					if ( htlc_outbound && htlc.cltv_expiry + LATENCY_GRACE_PERIOD_BLOCKS <= height) ||
-					   (!htlc_outbound && htlc.cltv_expiry <= height + CLTV_CLAIM_BUFFER && self.payment_preimages.contains_key(&htlc.payment_hash)) {
-						log_info!(logger, "Force-closing channel due to {} HTLC timeout, HTLC expiry is {}", if htlc_outbound { "outbound" } else { "inbound "}, htlc.cltv_expiry);
+					let htlc_outbound = $local_tx == htlc.borrow().offered;
+					if ( htlc_outbound && htlc.borrow().cltv_expiry + LATENCY_GRACE_PERIOD_BLOCKS <= height) ||
+					   (!htlc_outbound && htlc.borrow().cltv_expiry <= height + CLTV_CLAIM_BUFFER && preimages.contains_key(&htlc.borrow().payment_hash)) {
+						log_info!(logger, "Force-closing channel due to {} HTLC timeout, HTLC expiry is {}", if htlc_outbound { "outbound" } else { "inbound "}, htlc.borrow().cltv_expiry);
 						return true;
 					}
 				}
 			}
 		}
 
-		scan_commitment!(self.current_local_commitment_tx.htlc_outputs.iter().map(|&(ref a, _, _)| a), true);
-
-		if let Some(ref txid) = self.current_remote_commitment_txid {
-			if let Some(ref htlc_outputs) = self.remote_claimable_outpoints.get(txid) {
-				scan_commitment!(htlc_outputs.iter().map(|&(ref a, _)| a), false);
-			}
-		}
-		if let Some(ref txid) = self.prev_remote_commitment_txid {
-			if let Some(ref htlc_outputs) = self.remote_claimable_outpoints.get(txid) {
-				scan_commitment!(htlc_outputs.iter().map(|&(ref a, _)| a), false);
-			}
-		}
+		scan_commitment!(local_htlc_outputs, true);
+		scan_commitment!(remote_htlc_outputs, false);
 
 		false
 	}