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)]
@@ -109,11 +110,8 @@ enum InboundHTLCState {
 	/// 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.
-	LocalRemoved(InboundHTLCRemovalReason),
+	/// anyway).
+	LocalRemoved(bool, InboundHTLCRemovalReason),
 }
 
 struct InboundHTLCOutput {
@@ -293,6 +291,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 +509,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 +738,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 +796,71 @@ 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! add_htlc_output {
+			($htlc: expr, $offered: expr, $list: expr) => {
+				$list.push(HTLCOutputInCommitment{
+					offered: $offered,
+					amount_msat: $htlc.amount_msat,
+					cltv_expiry: $htlc.cltv_expiry,
+					payment_hash: $htlc.payment_hash,
+					transaction_output_index: None
+				});
+			}
+		}
+
+		let cap = self.pending_inbound_htlcs.len() + self.pending_outbound_htlcs.len();
+		let mut remote_outputs = Vec::with_capacity(cap);
+		let mut local_outputs = Vec::with_capacity(cap);
+		let awaiting_raa = (self.channel_state & ChannelState::AwaitingRemoteRevoke as u32) != 0;
+		for ref htlc in self.pending_inbound_htlcs.iter() {
+			match htlc.state {
+				InboundHTLCState::AwaitingAnnouncedRemoteRevoke(_) => {
+					add_htlc_output!(htlc, false, local_outputs);
+					add_htlc_output!(htlc, true, remote_outputs);
+				},
+				InboundHTLCState::Committed => {
+					add_htlc_output!(htlc, false, local_outputs);
+					add_htlc_output!(htlc, true, remote_outputs);
+				},
+				InboundHTLCState::LocalRemoved(revoked, _) => {
+					add_htlc_output!(htlc, false, local_outputs);
+					if awaiting_raa && !revoked {
+						add_htlc_output!(htlc, true, remote_outputs)
+					}
+				},
+				_ => {},
+			}
+		}
+		for ref htlc in self.pending_outbound_htlcs.iter() {
+			match htlc.state {
+				OutboundHTLCState::LocalAnnounced(_) => {
+					add_htlc_output!(htlc, true, local_outputs);
+					add_htlc_output!(htlc, false, remote_outputs);
+			},
+				OutboundHTLCState::Committed => {
+					add_htlc_output!(htlc, true, local_outputs);
+					add_htlc_output!(htlc, false, remote_outputs);
+				},
+				OutboundHTLCState::RemoteRemoved(_) => {
+					add_htlc_output!(htlc, true, local_outputs);
+					add_htlc_output!(htlc, false, remote_outputs)
+				},
+				OutboundHTLCState::AwaitingRemoteRevokeToRemove(_) => {
+					add_htlc_output!(htlc, true, local_outputs);
+					add_htlc_output!(htlc, false, remote_outputs);
+				},
+				OutboundHTLCState::AwaitingRemovedRemoteRevoke(_) => {
+					if awaiting_raa {
+						add_htlc_output!(htlc, false, remote_outputs)
+					}
+				},
+			}
+		}
+
+		ChannelMonitor::<ChanSigner>::would_broadcast_at_height_given_htlcs(local_outputs.iter(), remote_outputs.iter(), height, &self.payment_preimages, logger)
+	}
+
 	// Utilities to build transactions:
 
 	fn get_commitment_transaction_number_obscure_factor(&self) -> u64 {
@@ -908,7 +974,7 @@ 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(revoked, _) => (!generated_by_local && !revoked, "LocalRemoved"),
 			};
 
 			if include {
@@ -917,7 +983,7 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 			} else {
 				log_trace!(logger, "   ...not including inbound HTLC {} (hash {}) with value {} due to state ({})", htlc.htlc_id, log_bytes!(htlc.payment_hash.0), htlc.amount_msat, state_name);
 				match &htlc.state {
-					&InboundHTLCState::LocalRemoved(ref reason) => {
+					&InboundHTLCState::LocalRemoved(false, ref reason) => {
 						if generated_by_local {
 							if let &InboundHTLCRemovalReason::Fulfill(_) = reason {
 								value_to_self_msat_offset += htlc.amount_msat as i64;
@@ -1190,7 +1256,7 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 				assert_eq!(htlc.payment_hash, payment_hash_calc);
 				match htlc.state {
 					InboundHTLCState::Committed => {},
-					InboundHTLCState::LocalRemoved(ref reason) => {
+					InboundHTLCState::LocalRemoved(_, ref reason) => {
 						if let &InboundHTLCRemovalReason::Fulfill(_) = reason {
 						} else {
 							log_warn!(logger, "Have preimage and want to fulfill HTLC with payment hash {} we already failed against channel {}", log_bytes!(htlc.payment_hash.0), log_bytes!(self.channel_id()));
@@ -1216,6 +1282,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 {
@@ -1262,7 +1329,7 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 				return Ok((None, Some(monitor_update)));
 			}
 			log_trace!(logger, "Upgrading HTLC {} to LocalRemoved with a Fulfill!", log_bytes!(htlc.payment_hash.0));
-			htlc.state = InboundHTLCState::LocalRemoved(InboundHTLCRemovalReason::Fulfill(payment_preimage_arg.clone()));
+			htlc.state = InboundHTLCState::LocalRemoved(false, InboundHTLCRemovalReason::Fulfill(payment_preimage_arg.clone()));
 		}
 
 		Ok((Some(msgs::UpdateFulfillHTLC {
@@ -1312,7 +1379,7 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 			if htlc.htlc_id == htlc_id_arg {
 				match htlc.state {
 					InboundHTLCState::Committed => {},
-					InboundHTLCState::LocalRemoved(_) => {
+					InboundHTLCState::LocalRemoved(_, _) => {
 						debug_assert!(false, "Tried to fail an HTLC that was already fail/fulfilled");
 						return Ok(None);
 					},
@@ -1356,7 +1423,7 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 
 		{
 			let htlc = &mut self.pending_inbound_htlcs[pending_idx];
-			htlc.state = InboundHTLCState::LocalRemoved(InboundHTLCRemovalReason::FailRelay(err_packet.clone()));
+			htlc.state = InboundHTLCState::LocalRemoved(false, InboundHTLCRemovalReason::FailRelay(err_packet.clone()));
 		}
 
 		Ok(Some(msgs::UpdateFailHTLC {
@@ -2011,6 +2078,17 @@ 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))));
 		}
 
+		// A LocalRemoved 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
+		// LocalRemoved HTLCs we were previously watching for.
+		self.pending_inbound_htlcs.retain(|htlc| {
+			log_trace!(logger, "Removing inbound LocalRemoved {}", log_bytes!(htlc.payment_hash.0));
+			if let &InboundHTLCState::LocalRemoved(true, _) = &htlc.state {
+				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 +2381,16 @@ 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| {
-				if let &InboundHTLCState::LocalRemoved(ref reason) = &htlc.state {
-					log_trace!(logger, " ...removing inbound LocalRemoved {}", log_bytes!(htlc.payment_hash.0));
+			for htlc in pending_inbound_htlcs.iter_mut() {
+				if let &mut InboundHTLCState::LocalRemoved(ref mut revoked, ref reason) = &mut htlc.state {
 					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;
+						if !*revoked {
+							value_to_self_msat_diff += htlc.amount_msat as i64;
+						}
 					}
-					false
-				} else { true }
-			});
-			for htlc in pending_inbound_htlcs.iter_mut() {
+					*revoked = true;
+					continue
+				}
 				let swap = if let &InboundHTLCState::AwaitingRemoteRevokeToAnnounce(_) = &htlc.state {
 					log_trace!(logger, " ...promoting inbound AwaitingRemoteRevokeToAnnounce {} to Committed", log_bytes!(htlc.payment_hash.0));
 					true
@@ -2347,11 +2411,11 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 								require_commitment = true;
 								match fail_msg {
 									HTLCFailureMsg::Relay(msg) => {
-										htlc.state = InboundHTLCState::LocalRemoved(InboundHTLCRemovalReason::FailRelay(msg.reason.clone()));
+										htlc.state = InboundHTLCState::LocalRemoved(false, InboundHTLCRemovalReason::FailRelay(msg.reason.clone()));
 										update_fail_htlcs.push(msg)
 									},
 									HTLCFailureMsg::Malformed(msg) => {
-										htlc.state = InboundHTLCState::LocalRemoved(InboundHTLCRemovalReason::FailMalformed((msg.sha256_of_onion, msg.failure_code)));
+										htlc.state = InboundHTLCState::LocalRemoved(false, InboundHTLCRemovalReason::FailMalformed((msg.sha256_of_onion, msg.failure_code)));
 										update_fail_malformed_htlcs.push(msg)
 									},
 								}
@@ -2364,6 +2428,19 @@ 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;
+					}
+					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));
@@ -2533,7 +2610,7 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 					true
 				},
 				InboundHTLCState::Committed => true,
-				InboundHTLCState::LocalRemoved(_) => {
+				InboundHTLCState::LocalRemoved(_, _) => {
 					// We (hopefully) sent a commitment_signed updating this HTLC (which we can
 					// re-transmit if needed) and they may have even sent a revoke_and_ack back
 					// (that we missed). Keep this around for now and if they tell us they missed
@@ -2688,7 +2765,7 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 		}
 
 		for htlc in self.pending_inbound_htlcs.iter() {
-			if let &InboundHTLCState::LocalRemoved(ref reason) = &htlc.state {
+			if let &InboundHTLCState::LocalRemoved(false, ref reason) = &htlc.state {
 				match reason {
 					&InboundHTLCRemovalReason::FailRelay(ref err_packet) => {
 						update_fail_htlcs.push(msgs::UpdateFailHTLC {
@@ -3115,14 +3192,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)
@@ -3819,7 +3888,7 @@ impl<ChanSigner: ChannelKeys> Channel<ChanSigner> {
 			if have_updates { break; }
 		}
 		for htlc in self.pending_inbound_htlcs.iter() {
-			if let InboundHTLCState::LocalRemoved(_) = htlc.state {
+			if let InboundHTLCState::LocalRemoved(false, _) = htlc.state {
 				have_updates = true;
 			}
 			if have_updates { break; }
@@ -4114,8 +4183,9 @@ impl<ChanSigner: ChannelKeys + Writeable> Writeable for Channel<ChanSigner> {
 				&InboundHTLCState::Committed => {
 					3u8.write(writer)?;
 				},
-				&InboundHTLCState::LocalRemoved(ref removal_reason) => {
+				&InboundHTLCState::LocalRemoved(ref revoked, ref removal_reason) => {
 					4u8.write(writer)?;
+					revoked.write(writer)?;
 					removal_reason.write(writer)?;
 				},
 			}
@@ -4151,6 +4221,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 +4323,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)?;
@@ -4286,7 +4363,7 @@ impl<ChanSigner: ChannelKeys + Readable> Readable for Channel<ChanSigner> {
 					1 => InboundHTLCState::AwaitingRemoteRevokeToAnnounce(Readable::read(reader)?),
 					2 => InboundHTLCState::AwaitingAnnouncedRemoteRevoke(Readable::read(reader)?),
 					3 => InboundHTLCState::Committed,
-					4 => InboundHTLCState::LocalRemoved(Readable::read(reader)?),
+					4 => InboundHTLCState::LocalRemoved(Readable::read(reader)?, Readable::read(reader)?),
 					_ => return Err(DecodeError::InvalidValue),
 				},
 			});
@@ -4312,6 +4389,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 +4515,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);
 					}