From f55f055b4edcc7b0d2d38ffc41872d0c6571db64 Mon Sep 17 00:00:00 2001
From: Matt Corallo <git@bluematt.me>
Date: Wed, 5 Sep 2018 18:32:55 -0400
Subject: [PATCH] Remove/fail uncommitted HTLCs upon peer disconnection

---
 src/ln/channel.rs        | 53 ++++++++++++++++++++++++++++++++++++++++
 src/ln/channelmanager.rs | 15 +++++++++---
 2 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/src/ln/channel.rs b/src/ln/channel.rs
index 4338f9bfa..e960250b8 100644
--- a/src/ln/channel.rs
+++ b/src/ln/channel.rs
@@ -1766,6 +1766,59 @@ impl Channel {
 		}
 	}
 
+	/// Removes any uncommitted HTLCs, to be used on peer disconnection, including any pending
+	/// HTLCs that we intended to add but haven't as we were waiting on a remote revoke.
+	/// Returns the set of PendingHTLCStatuses from remote uncommitted HTLCs (which we're
+	/// implicitly dropping) and the payment_hashes of HTLCs we tried to add but are dropping.
+	pub fn remove_uncommitted_htlcs(&mut self) -> Vec<(HTLCSource, [u8; 32])> {
+		self.pending_inbound_htlcs.retain(|htlc| {
+			match htlc.state {
+				InboundHTLCState::RemoteAnnounced => {
+					// They sent us an update_add_htlc but we never got the commitment_signed.
+					// We'll tell them what commitment_signed we're expecting next and they'll drop
+					// this HTLC accordingly
+					false
+				},
+				InboundHTLCState::AwaitingRemoteRevokeToAnnounce|InboundHTLCState::AwaitingAnnouncedRemoteRevoke => {
+					// Same goes for AwaitingRemoteRevokeToRemove and AwaitingRemovedRemoteRevoke
+					// We received a commitment_signed updating this HTLC and (at least hopefully)
+					// sent a revoke_and_ack (which we can re-transmit) and have heard nothing
+					// in response to it yet, so don't touch it.
+					true
+				},
+				InboundHTLCState::Committed => true,
+				InboundHTLCState::LocalRemoved => { // Same goes for LocalAnnounced
+					// 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
+					// the commitment_signed we can re-transmit the update then.
+					true
+				},
+			}
+		});
+
+		for htlc in self.pending_outbound_htlcs.iter_mut() {
+			if htlc.state == OutboundHTLCState::RemoteRemoved {
+				// They sent us an update to remove this but haven't yet sent the corresponding
+				// commitment_signed, we need to move it back to Committed and they can re-send
+				// the update upon reconnection.
+				htlc.state = OutboundHTLCState::Committed;
+			}
+		}
+
+		let mut outbound_drops = Vec::new();
+		self.holding_cell_htlc_updates.retain(|htlc_update| {
+			match htlc_update {
+				&HTLCUpdateAwaitingACK::AddHTLC { ref payment_hash, ref source, .. } => {
+					outbound_drops.push((source.clone(), payment_hash.clone()));
+					false
+				},
+				&HTLCUpdateAwaitingACK::ClaimHTLC {..} | &HTLCUpdateAwaitingACK::FailHTLC {..} => true,
+			}
+		});
+		outbound_drops
+	}
+
 	pub fn update_fee(&mut self, fee_estimator: &FeeEstimator, msg: &msgs::UpdateFee) -> Result<(), HandleError> {
 		if self.channel_outbound {
 			return Err(HandleError{err: "Non-funding remote tried to update channel fee", action: None});
diff --git a/src/ln/channelmanager.rs b/src/ln/channelmanager.rs
index af73cc422..f52db6bc7 100644
--- a/src/ln/channelmanager.rs
+++ b/src/ln/channelmanager.rs
@@ -2126,6 +2126,7 @@ impl ChannelMessageHandler for ChannelManager {
 	fn peer_disconnected(&self, their_node_id: &PublicKey, no_connection_possible: bool) {
 		let mut new_events = Vec::new();
 		let mut failed_channels = Vec::new();
+		let mut failed_payments = Vec::new();
 		{
 			let mut channel_state_lock = self.channel_state.lock().unwrap();
 			let channel_state = channel_state_lock.borrow_parts();
@@ -2150,9 +2151,12 @@ impl ChannelMessageHandler for ChannelManager {
 			} else {
 				for chan in channel_state.by_id {
 					if chan.1.get_their_node_id() == *their_node_id {
-						//TODO: mark channel disabled (and maybe announce such after a timeout). Also
-						//fail and wipe any uncommitted outbound HTLCs as those are considered after
-						//reconnect.
+						//TODO: mark channel disabled (and maybe announce such after a timeout).
+						let failed_adds = chan.1.remove_uncommitted_htlcs();
+						if !failed_adds.is_empty() {
+							let chan_update = self.get_channel_update(&chan.1).map(|u| u.encode_with_len()).unwrap(); // Cannot add/recv HTLCs before we have a short_id so unwrap is safe
+							failed_payments.push((chan_update, failed_adds));
+						}
 					}
 				}
 			}
@@ -2166,6 +2170,11 @@ impl ChannelMessageHandler for ChannelManager {
 				pending_events.push(event);
 			}
 		}
+		for (chan_update, mut htlc_sources) in failed_payments {
+			for (htlc_source, payment_hash) in htlc_sources.drain(..) {
+				self.fail_htlc_backwards_internal(self.channel_state.lock().unwrap(), htlc_source, &payment_hash, HTLCFailReason::Reason { failure_code: 0x1000 | 7, data: chan_update.clone() });
+			}
+		}
 	}
 
 	fn handle_error(&self, their_node_id: &PublicKey, msg: &msgs::ErrorMessage) {
-- 
2.39.5