From 5fe598c073cbaa2b45e44ca5e671ee6fab2719fe Mon Sep 17 00:00:00 2001
From: Matt Corallo <git@bluematt.me>
Date: Fri, 7 Sep 2018 15:11:52 -0400
Subject: [PATCH] Drop channels on disconnection if we haven't funded them yet

---
 src/ln/channel.rs        |  8 +++++++-
 src/ln/channelmanager.rs | 17 ++++++++++++-----
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/ln/channel.rs b/src/ln/channel.rs
index e960250b..15e0a76a 100644
--- a/src/ln/channel.rs
+++ b/src/ln/channel.rs
@@ -1771,6 +1771,13 @@ impl Channel {
 	/// 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])> {
+		let mut outbound_drops = Vec::new();
+
+		if self.channel_state < ChannelState::FundingSent as u32 {
+			self.channel_state = ChannelState::ShutdownComplete as u32;
+			return outbound_drops;
+		}
+
 		self.pending_inbound_htlcs.retain(|htlc| {
 			match htlc.state {
 				InboundHTLCState::RemoteAnnounced => {
@@ -1806,7 +1813,6 @@ impl Channel {
 			}
 		}
 
-		let mut outbound_drops = Vec::new();
 		self.holding_cell_htlc_updates.retain(|htlc_update| {
 			match htlc_update {
 				&HTLCUpdateAwaitingACK::AddHTLC { ref payment_hash, ref source, .. } => {
diff --git a/src/ln/channelmanager.rs b/src/ln/channelmanager.rs
index f52db6bc..333d4257 100644
--- a/src/ln/channelmanager.rs
+++ b/src/ln/channelmanager.rs
@@ -2149,16 +2149,23 @@ impl ChannelMessageHandler for ChannelManager {
 					}
 				});
 			} else {
-				for chan in channel_state.by_id {
-					if chan.1.get_their_node_id() == *their_node_id {
+				channel_state.by_id.retain(|_, chan| {
+					if chan.get_their_node_id() == *their_node_id {
 						//TODO: mark channel disabled (and maybe announce such after a timeout).
-						let failed_adds = chan.1.remove_uncommitted_htlcs();
+						let failed_adds = chan.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
+							let chan_update = self.get_channel_update(&chan).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));
 						}
+						if chan.is_shutdown() {
+							if let Some(short_id) = chan.get_short_channel_id() {
+								short_to_id.remove(&short_id);
+							}
+							return false;
+						}
 					}
-				}
+					true
+				})
 			}
 		}
 		for failure in failed_channels.drain(..) {
-- 
2.30.2