f get_update_fulfill_htlc use enum

TheBlueMatt · TheBlueMatt · commit 363fe345c543 · 2021-07-27T00:06:00.000Z
diff --git a/lightning/src/ln/channel.rs b/lightning/src/ln/channel.rs
@@ -301,6 +301,33 @@ pub struct CounterpartyForwardingInfo {
 	pub cltv_expiry_delta: u16,
 }
 
+/// A return value enum for get_update_fulfill_htlc. See UpdateFulfillCommitFetch variants for
+/// description
+enum UpdateFulfillFetch {
+	NewClaim {
+		monitor_update: ChannelMonitorUpdate,
+		msg: Option<msgs::UpdateFulfillHTLC>,
+	},
+	DuplicateClaim {},
+}
+
+/// The return type of get_update_fulfill_htlc_and_commit.
+pub enum UpdateFulfillCommitFetch {
+	/// Indicates the HTLC fulfill is new, and either generated an update_fulfill message, placed
+	/// it in the holding cell, or re-generated the update_fulfill message after the same claim was
+	/// previously placed in the holding cell (and has since been removed).
+	NewClaim {
+		/// The ChannelMonitorUpdate which places the new payment preimage in the channel monitor
+		monitor_update: ChannelMonitorUpdate,
+		/// The update_fulfill message and commitment_signed message (if the claim was not placed
+		/// in the holding cell).
+		msgs: Option<(msgs::UpdateFulfillHTLC, msgs::CommitmentSigned)>,
+	},
+	/// Indicates the HTLC fulfill is duplicative, and already existed either in the holding cell,
+	/// or the HTLC is forgotten completely (presumably previously claimed).
+	DuplicateClaim {},
+}
+
 // TODO: We should refactor this to be an Inbound/OutboundChannel until initial setup handshaking
 // has been completed, and then turn into a Channel to get compiler-time enforcement of things like
 // calling channel_id() before we're set up or things like get_outbound_funding_signed on an
@@ -1232,7 +1259,7 @@ impl<Signer: Sign> Channel<Signer> {
 		make_funding_redeemscript(&self.get_holder_pubkeys().funding_pubkey, self.counterparty_funding_pubkey())
 	}
 
-	fn get_update_fulfill_htlc<L: Deref>(&mut self, htlc_id_arg: u64, payment_preimage_arg: PaymentPreimage, logger: &L) -> Option<(Option<msgs::UpdateFulfillHTLC>, ChannelMonitorUpdate)> where L::Target: Logger {
+	fn get_update_fulfill_htlc<L: Deref>(&mut self, htlc_id_arg: u64, payment_preimage_arg: PaymentPreimage, logger: &L) -> UpdateFulfillFetch where L::Target: Logger {
 		// Either ChannelFunded got set (which means it won't be unset) or there is no way any
 		// caller thought we could have something claimed (cause we wouldn't have accepted in an
 		// incoming HTLC anyway). If we got to ShutdownComplete, callers aren't allowed to call us,
@@ -1260,7 +1287,7 @@ impl<Signer: Sign> Channel<Signer> {
 							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()));
 							debug_assert!(false, "Tried to fulfill an HTLC that was already failed");
 						}
-						return None;
+						return UpdateFulfillFetch::DuplicateClaim {};
 					},
 					_ => {
 						debug_assert!(false, "Have an inbound HTLC we tried to claim before it was fully committed to");
@@ -1276,7 +1303,7 @@ impl<Signer: Sign> Channel<Signer> {
 			// If we failed to find an HTLC to fulfill, make sure it was previously fulfilled and
 			// this is simply a duplicate claim, not previously failed and we lost funds.
 			debug_assert!(self.historical_inbound_htlc_fulfills.contains(&htlc_id_arg));
-			return None;
+			return UpdateFulfillFetch::DuplicateClaim {};
 		}
 
 		// Now update local state:
@@ -1300,7 +1327,7 @@ impl<Signer: Sign> Channel<Signer> {
 							self.latest_monitor_update_id -= 1;
 							#[cfg(any(test, feature = "fuzztarget"))]
 							debug_assert!(self.historical_inbound_htlc_fulfills.contains(&htlc_id_arg));
-							return None;
+							return UpdateFulfillFetch::DuplicateClaim {};
 						}
 					},
 					&HTLCUpdateAwaitingACK::FailHTLC { htlc_id, .. } => {
@@ -1309,7 +1336,7 @@ impl<Signer: Sign> Channel<Signer> {
 							// TODO: We may actually be able to switch to a fulfill here, though its
 							// rare enough it may not be worth the complexity burden.
 							debug_assert!(false, "Tried to fulfill an HTLC that was already failed");
-							return Some((None, monitor_update));
+							return UpdateFulfillFetch::NewClaim { monitor_update, msg: None };
 						}
 					},
 					_ => {}
@@ -1321,7 +1348,7 @@ impl<Signer: Sign> Channel<Signer> {
 			});
 			#[cfg(any(test, feature = "fuzztarget"))]
 			self.historical_inbound_htlc_fulfills.insert(htlc_id_arg);
-			return Some((None, monitor_update));
+			return UpdateFulfillFetch::NewClaim { monitor_update, msg: None };
 		}
 		#[cfg(any(test, feature = "fuzztarget"))]
 		self.historical_inbound_htlc_fulfills.insert(htlc_id_arg);
@@ -1331,31 +1358,34 @@ impl<Signer: Sign> Channel<Signer> {
 			if let InboundHTLCState::Committed = htlc.state {
 			} else {
 				debug_assert!(false, "Have an inbound HTLC we tried to claim before it was fully committed to");
-				return Some((None, monitor_update));
+				return UpdateFulfillFetch::NewClaim { monitor_update, msg: None };
 			}
 			log_trace!(logger, "Upgrading HTLC {} to LocalRemoved with a Fulfill in channel {}!", log_bytes!(htlc.payment_hash.0), log_bytes!(self.channel_id));
 			htlc.state = InboundHTLCState::LocalRemoved(InboundHTLCRemovalReason::Fulfill(payment_preimage_arg.clone()));
 		}
 
-		Some((Some(msgs::UpdateFulfillHTLC {
-			channel_id: self.channel_id(),
-			htlc_id: htlc_id_arg,
-			payment_preimage: payment_preimage_arg,
-		}), monitor_update))
+		UpdateFulfillFetch::NewClaim {
+			monitor_update,
+			msg: Some(msgs::UpdateFulfillHTLC {
+				channel_id: self.channel_id(),
+				htlc_id: htlc_id_arg,
+				payment_preimage: payment_preimage_arg,
+			}),
+		}
 	}
 
-	pub fn get_update_fulfill_htlc_and_commit<L: Deref>(&mut self, htlc_id: u64, payment_preimage: PaymentPreimage, logger: &L) -> Result<Option<(Option<(msgs::UpdateFulfillHTLC, msgs::CommitmentSigned)>, ChannelMonitorUpdate)>, ChannelError> where L::Target: Logger {
+	pub fn get_update_fulfill_htlc_and_commit<L: Deref>(&mut self, htlc_id: u64, payment_preimage: PaymentPreimage, logger: &L) -> Result<UpdateFulfillCommitFetch, ChannelError> where L::Target: Logger {
 		match self.get_update_fulfill_htlc(htlc_id, payment_preimage, logger) {
-			Some((Some(update_fulfill_htlc), mut monitor_update)) => {
+			UpdateFulfillFetch::NewClaim { mut monitor_update, msg: Some(update_fulfill_htlc) } => {
 				let (commitment, mut additional_update) = self.send_commitment_no_status_check(logger)?;
 				// send_commitment_no_status_check may bump latest_monitor_id but we want them to be
 				// strictly increasing by one, so decrement it here.
 				self.latest_monitor_update_id = monitor_update.update_id;
 				monitor_update.updates.append(&mut additional_update.updates);
-				Ok(Some((Some((update_fulfill_htlc, commitment)), monitor_update)))
+				Ok(UpdateFulfillCommitFetch::NewClaim { monitor_update, msgs: Some((update_fulfill_htlc, commitment)) })
 			},
-			Some((None, monitor_update)) => Ok(Some((None, monitor_update))),
-			None => Ok(None)
+			UpdateFulfillFetch::NewClaim { monitor_update, msg: None } => Ok(UpdateFulfillCommitFetch::NewClaim { monitor_update, msgs: None }),
+			UpdateFulfillFetch::DuplicateClaim {} => Ok(UpdateFulfillCommitFetch::DuplicateClaim {}),
 		}
 	}
 
@@ -2463,7 +2493,10 @@ impl<Signer: Sign> Channel<Signer> {
 						// not fail - any in between attempts to claim the HTLC will have resulted
 						// in it hitting the holding cell again and we cannot change the state of a
 						// holding cell HTLC from fulfill to anything else.
-						let (update_fulfill_msg_option, mut additional_monitor_update) = self.get_update_fulfill_htlc(htlc_id, *payment_preimage, logger).unwrap();
+						let (update_fulfill_msg_option, mut additional_monitor_update) =
+							if let UpdateFulfillFetch::NewClaim { msg, monitor_update } = self.get_update_fulfill_htlc(htlc_id, *payment_preimage, logger) {
+								(msg, monitor_update)
+							} else { unreachable!() };
 						update_fulfill_htlcs.push(update_fulfill_msg_option.unwrap());
 						monitor_update.updates.append(&mut additional_monitor_update.updates);
 					},
diff --git a/lightning/src/ln/channelmanager.rs b/lightning/src/ln/channelmanager.rs
@@ -44,7 +44,7 @@ use chain::transaction::{OutPoint, TransactionData};
 // construct one themselves.
 use ln::{PaymentHash, PaymentPreimage, PaymentSecret};
 pub use ln::channel::CounterpartyForwardingInfo;
-use ln::channel::{Channel, ChannelError, ChannelUpdateStatus};
+use ln::channel::{Channel, ChannelError, ChannelUpdateStatus, UpdateFulfillCommitFetch};
 use ln::features::{InitFeatures, NodeFeatures};
 use routing::router::{Route, RouteHop};
 use ln::msgs;
@@ -2682,7 +2682,7 @@ impl<Signer: Sign, M: Deref, T: Deref, K: Deref, F: Deref, L: Deref> ChannelMana
 			let was_frozen_for_monitor = chan.get().is_awaiting_monitor_update();
 			match chan.get_mut().get_update_fulfill_htlc_and_commit(prev_hop.htlc_id, payment_preimage, &self.logger) {
 				Ok(msgs_monitor_option) => {
-					if let Some((msgs, monitor_update)) = msgs_monitor_option {
+					if let UpdateFulfillCommitFetch::NewClaim { msgs, monitor_update } = msgs_monitor_option {
 						if let Err(e) = self.chain_monitor.update_channel(chan.get().get_funding_txo().unwrap(), monitor_update) {
 							if was_frozen_for_monitor {
 								assert!(msgs.is_none());
