From 4ca5bcf8cfc5cb3868974028b23ab7829c21426b Mon Sep 17 00:00:00 2001
From: Matt Corallo <git@bluematt.me>
Date: Sun, 26 Aug 2018 16:34:47 -0400
Subject: [PATCH] Update PendingHTLCStatus to hold malformed HTLC error
 messages

---
 src/ln/channel.rs        | 22 +++++++++++++++-------
 src/ln/channelmanager.rs | 16 +++++++++++-----
 2 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/src/ln/channel.rs b/src/ln/channel.rs
index 932bebb31..44875fcd6 100644
--- a/src/ln/channel.rs
+++ b/src/ln/channel.rs
@@ -16,7 +16,7 @@ use crypto::hkdf::{hkdf_extract,hkdf_expand};
 use ln::msgs;
 use ln::msgs::{ErrorAction, HandleError, MsgEncodable};
 use ln::channelmonitor::ChannelMonitor;
-use ln::channelmanager::{PendingHTLCStatus, PendingForwardHTLCInfo, HTLCFailReason};
+use ln::channelmanager::{PendingHTLCStatus, PendingForwardHTLCInfo, HTLCFailReason, HTLCFailureMsg};
 use ln::chan_utils::{TxCreationKeys,HTLCOutputInCommitment,HTLC_SUCCESS_TX_WEIGHT,HTLC_TIMEOUT_TX_WEIGHT};
 use ln::chan_utils;
 use chain::chaininterface::{FeeEstimator,ConfirmationTarget};
@@ -1681,7 +1681,8 @@ impl Channel {
 
 		let mut to_forward_infos = Vec::new();
 		let mut revoked_htlcs = Vec::new();
-		let mut failed_htlcs = Vec::new();
+		let mut update_fail_htlcs = Vec::new();
+		let mut update_fail_malformed_htlcs = Vec::new();
 		let mut require_commitment = false;
 		let mut value_to_self_msat_diff: i64 = 0;
 		// We really shouldnt have two passes here, but retain gives a non-mutable ref (Rust bug)
@@ -1709,7 +1710,10 @@ impl Channel {
 					PendingHTLCStatus::Fail(fail_msg) => {
 						htlc.state = HTLCState::LocalRemoved;
 						require_commitment = true;
-						failed_htlcs.push(fail_msg);
+						match fail_msg {
+							HTLCFailureMsg::Relay(msg) => update_fail_htlcs.push(msg),
+							HTLCFailureMsg::Malformed(msg) => update_fail_malformed_htlcs.push(msg),
+						}
 					},
 					PendingHTLCStatus::Forward(forward_info) => {
 						to_forward_infos.push(forward_info);
@@ -1728,10 +1732,14 @@ impl Channel {
 
 		match self.free_holding_cell_htlcs()? {
 			Some(mut commitment_update) => {
-				commitment_update.0.update_fail_htlcs.reserve(failed_htlcs.len());
-				for fail_msg in failed_htlcs.drain(..) {
+				commitment_update.0.update_fail_htlcs.reserve(update_fail_htlcs.len());
+				for fail_msg in update_fail_htlcs.drain(..) {
 					commitment_update.0.update_fail_htlcs.push(fail_msg);
 				}
+				commitment_update.0.update_fail_malformed_htlcs.reserve(update_fail_malformed_htlcs.len());
+				for fail_msg in update_fail_malformed_htlcs.drain(..) {
+					commitment_update.0.update_fail_malformed_htlcs.push(fail_msg);
+				}
 				Ok((Some(commitment_update.0), to_forward_infos, revoked_htlcs, commitment_update.1))
 			},
 			None => {
@@ -1740,8 +1748,8 @@ impl Channel {
 					Ok((Some(msgs::CommitmentUpdate {
 						update_add_htlcs: Vec::new(),
 						update_fulfill_htlcs: Vec::new(),
-						update_fail_htlcs: failed_htlcs,
-						update_fail_malformed_htlcs: Vec::new(),
+						update_fail_htlcs,
+						update_fail_malformed_htlcs,
 						commitment_signed
 					}), to_forward_infos, revoked_htlcs, monitor_update))
 				} else {
diff --git a/src/ln/channelmanager.rs b/src/ln/channelmanager.rs
index 9daa7ba6a..88eb4b126 100644
--- a/src/ln/channelmanager.rs
+++ b/src/ln/channelmanager.rs
@@ -50,11 +50,17 @@ mod channel_held_info {
 		pub(super) outgoing_cltv_value: u32,
 	}
 
+	#[derive(Clone)] // See Channel::revoke_and_ack for why, tl;dr: Rust bug
+	pub enum HTLCFailureMsg {
+		Relay(msgs::UpdateFailHTLC),
+		Malformed(msgs::UpdateFailMalformedHTLC),
+	}
+
 	/// Stores whether we can't forward an HTLC or relevant forwarding info
 	#[derive(Clone)] // See Channel::revoke_and_ack for why, tl;dr: Rust bug
 	pub enum PendingHTLCStatus {
 		Forward(PendingForwardHTLCInfo),
-		Fail(msgs::UpdateFailHTLC),
+		Fail(HTLCFailureMsg),
 	}
 
 	#[cfg(feature = "fuzztarget")]
@@ -699,11 +705,11 @@ impl ChannelManager {
 					if channel_state.is_none() {
 						channel_state = Some(self.channel_state.lock().unwrap());
 					}
-					return (PendingHTLCStatus::Fail(msgs::UpdateFailHTLC {
+					return (PendingHTLCStatus::Fail(HTLCFailureMsg::Relay(msgs::UpdateFailHTLC {
 						channel_id: msg.channel_id,
 						htlc_id: msg.htlc_id,
 						reason: ChannelManager::build_first_hop_failure_packet(&shared_secret, $err_code, $data),
-					}), shared_secret, channel_state.unwrap());
+					})), shared_secret, channel_state.unwrap());
 				}
 			}
 		}
@@ -1726,11 +1732,11 @@ impl ChannelMessageHandler for ChannelManager {
 				}
 				if !acceptable_cycle {
 					log_info!(self, "Failed to accept incoming HTLC: Payment looped through us twice");
-					pending_forward_info = PendingHTLCStatus::Fail(msgs::UpdateFailHTLC {
+					pending_forward_info = PendingHTLCStatus::Fail(HTLCFailureMsg::Relay(msgs::UpdateFailHTLC {
 						channel_id: msg.channel_id,
 						htlc_id: msg.htlc_id,
 						reason: ChannelManager::build_first_hop_failure_packet(&shared_secret, 0x4000 | 0x2000 | 2, &[0;0]),
-					});
+					}));
 				} else {
 					will_forward = true;
 				}
-- 
2.39.5